From 73b1a2cce6ff840f79015b7d80afa4b10dae2d9e Mon Sep 17 00:00:00 2001 From: luandalmazo Date: Thu, 20 Aug 2026 09:09:55 -0300 Subject: [PATCH] add datamint import command --- datamint/__main__.py | 2 + datamint/client_cmd_tools/datamint_import.py | 249 +++++++++++++++++++ datamint/importers/__init__.py | 2 + datamint/importers/_detect.py | 108 ++++++++ docs/source/command_line_tools.rst | 32 ++- 5 files changed, 392 insertions(+), 1 deletion(-) create mode 100644 datamint/client_cmd_tools/datamint_import.py create mode 100644 datamint/importers/_detect.py diff --git a/datamint/__main__.py b/datamint/__main__.py index b913f3ba..2c4b7f41 100644 --- a/datamint/__main__.py +++ b/datamint/__main__.py @@ -13,6 +13,7 @@ "train": "datamint.client_cmd_tools.datamint_train", "inference": "datamint.client_cmd_tools.datamint_inference", "example": "datamint.client_cmd_tools.datamint_example", + "import": "datamint.client_cmd_tools.datamint_import", } _COMMAND_HELP: dict[str, str] = { @@ -22,6 +23,7 @@ "train": "Train a model on a Datamint project", "inference": "Run local inference with a registered model", "example": "Populate a project with an example dataset", + "import": "Import a labeled dataset (COCO, Pascal VOC, or YOLO)", } diff --git a/datamint/client_cmd_tools/datamint_import.py b/datamint/client_cmd_tools/datamint_import.py new file mode 100644 index 00000000..87baad1e --- /dev/null +++ b/datamint/client_cmd_tools/datamint_import.py @@ -0,0 +1,249 @@ +"""datamint import command-line tool. + +Uploads an already-labeled dataset (COCO, Pascal VOC, or YOLO format) to a +Datamint project in one call -- images plus their bounding-box annotations. +See ``datamint import --help`` for the available options. +""" +import argparse +import logging +import sys +from pathlib import Path + +from datamint.client_cmd_tools.datamint_upload import handle_api_key +from datamint.exceptions import DatamintException +from datamint.importers import ( + COCOImporter, + DatasetFormat, + PascalVOCImporter, + YOLOImporter, + detect_format, + sniff_single_format, +) +from datamint.utils.logging_utils import load_cmdline_logging_config + +_LOGGER = logging.getLogger(__name__) +_USER_LOGGER = logging.getLogger('user_logger') + +_IMPORTER_CLASSES = { + 'coco': COCOImporter, + 'pascal_voc': PascalVOCImporter, + 'yolo': YOLOImporter, +} + +# Which CLI override flags apply to each format's importer constructor. +_FORMAT_OVERRIDE_KEYS = { + 'coco': ('annotations_file', 'images_dir'), + 'pascal_voc': ('annotations_dir', 'images_dir'), + 'yolo': ('images_dir', 'labels_dir', 'class_names', 'data_yaml'), +} + +_REQUIRED_KEYS = { + 'coco': ('annotations_file', 'images_dir'), + 'pascal_voc': ('annotations_dir', 'images_dir'), + 'yolo': ('images_dir', 'labels_dir'), +} + + +def _build_parser(subparsers: argparse._SubParsersAction | None = None) -> argparse.ArgumentParser: + """Build the argument parser. + + When ``subparsers`` is given, the parser is registered as an ``import`` subparser + (used by ``datamint``'s combined completion tree) instead of a standalone parser. + """ + kwargs = { + 'description': 'Upload an already-labeled dataset (COCO, Pascal VOC, or YOLO) to a Datamint project.', + 'epilog': """ +Examples: + datamint import ./my_coco_dataset --project MyProject # auto-detect the format + datamint import ./my_yolo_dataset --project MyProject --format yolo + datamint import ./voc_dataset --project MyProject --dry-run # preview without uploading + +More Documentation: https://sonanceai.github.io/datamint-python-api/command_line_tools.html + """, + 'formatter_class': argparse.RawDescriptionHelpFormatter, + } + if subparsers is not None: + parser = subparsers.add_parser('import', **kwargs) + else: + parser = argparse.ArgumentParser(**kwargs) + + parser.add_argument('path', type=str, metavar='PATH', help='Root directory of the labeled dataset.') + parser.add_argument('--project', type=str, required=True, + help='Name of the Datamint project to import into. Created if it does not exist yet.') + parser.add_argument('--format', type=str, choices=sorted(_IMPORTER_CLASSES), + help='Annotation format. If omitted, the format is auto-detected from PATH.') + + detect_group = parser.add_argument_group( + 'Path overrides', + 'Override the paths that auto-detection would otherwise infer. Required when --format ' + 'is given and the layout could not be located automatically.', + ) + detect_group.add_argument('--annotations-file', type=str, help='(COCO) Path to the annotations JSON file.') + detect_group.add_argument('--annotations-dir', type=str, help='(Pascal VOC) Directory of annotation XML files.') + detect_group.add_argument('--labels-dir', type=str, help='(YOLO) Directory of label .txt files.') + detect_group.add_argument('--images-dir', type=str, help='Directory of image files.') + detect_group.add_argument('--class-names', type=str, nargs='+', + help='(YOLO) Ordered class names, index 0 first. Alternative to --data-yaml.') + detect_group.add_argument('--data-yaml', type=str, help='(YOLO) data.yaml file providing class names.') + + parser.add_argument('--tag', type=str, action='append', help='A tag to apply to every uploaded resource.') + parser.add_argument('--imported-from', type=str, default=None, + help="Provenance label stored on each annotation. Defaults to '-import'.") + parser.add_argument('--on-error', type=str, choices=('raise', 'skip'), default='raise', + help='Whether to stop or skip past a single resource/annotation upload failure.') + parser.add_argument('--dry-run', action='store_true', default=False, + help='Parse and print a summary, but do not create the project or upload anything.') + parser.add_argument('--yes', action='store_true', default=False, + help='Skip the auto-detected-format confirmation prompt.') + parser.add_argument('--verbose', action='store_true', default=False, help='Print debug messages.') + + return parser + + +def _parse_args() -> argparse.Namespace: + parser = _build_parser() + import argcomplete + argcomplete.autocomplete(parser) + return parser.parse_args() + + +def _resolve_dataset(args: argparse.Namespace) -> tuple[DatasetFormat, dict]: + """Figure out the format + importer constructor kwargs, applying CLI overrides. + + Prompts for confirmation when the format was auto-detected (unless ``--yes``). + Raises ``DatamintException`` when the format/paths can't be resolved. + """ + path = Path(args.path) + + if args.format is not None: + fmt: DatasetFormat = args.format + detected = sniff_single_format(path, fmt) + base_kwargs = dict(detected.importer_kwargs) if detected is not None else {} + else: + detected = detect_format(path) + if detected is None: + raise DatamintException( + f"Could not auto-detect a dataset format under '{path}'. " + f"Pass --format explicitly ({', '.join(sorted(_IMPORTER_CLASSES))}) " + 'together with the matching path override flags.' + ) + fmt = detected.format + base_kwargs = dict(detected.importer_kwargs) + + _USER_LOGGER.info(f"Detected a {fmt.upper()} dataset under '{path}':") + for key, value in base_kwargs.items(): + _USER_LOGGER.info(f' {key}: {value}') + if not args.yes: + answer = input(f'Import as {fmt.upper()}? (y/n): ') + if answer.strip().lower() != 'y': + raise DatamintException( + 'Aborted. Re-run with --format to choose a different format explicitly.' + ) + + overrides = { + 'annotations_file': args.annotations_file, + 'annotations_dir': args.annotations_dir, + 'labels_dir': args.labels_dir, + 'images_dir': args.images_dir, + 'class_names': args.class_names, + 'data_yaml': args.data_yaml, + } + for key in _FORMAT_OVERRIDE_KEYS[fmt]: + value = overrides.get(key) + if value is not None: + base_kwargs[key] = Path(value) if key != 'class_names' else value + + missing = [key for key in _REQUIRED_KEYS[fmt] if key not in base_kwargs] + if missing: + flags = ', '.join(f'--{key.replace("_", "-")}' for key in missing) + raise DatamintException( + f"Could not locate the {fmt.upper()} dataset layout under '{path}'. " + f'Pass {flags} explicitly.' + ) + + return fmt, base_kwargs + + +def main() -> None: + load_cmdline_logging_config() + + args = _parse_args() + + if args.verbose: + logging.getLogger().handlers[0].setLevel(logging.DEBUG) + logging.getLogger('datamint').setLevel(logging.DEBUG) + _LOGGER.setLevel(logging.DEBUG) + _USER_LOGGER.setLevel(logging.DEBUG) + + try: + path = Path(args.path) + if not path.is_dir(): + raise DatamintException(f"'{path}' is not a directory.") + + fmt, importer_kwargs = _resolve_dataset(args) + importer = _IMPORTER_CLASSES[fmt](**importer_kwargs) + + result = importer.parse() + _USER_LOGGER.info( + f'Parsed {result.num_images} image(s), {result.num_boxes} box(es), ' + f'{len(result.class_names)} class(es): {result.class_names}' + ) + if result.missing_images: + _USER_LOGGER.warning(f'{len(result.missing_images)} referenced image(s) not found on disk.') + if result.unsupported_annotations: + _USER_LOGGER.warning( + f'{result.unsupported_annotations} annotation(s) use an unsupported type and will be skipped.' + ) + + if args.dry_run: + _USER_LOGGER.info('Dry run: nothing was uploaded.') + return + + api_key = handle_api_key() + if api_key is None: + _USER_LOGGER.error('API key not provided. Aborting.') + sys.exit(1) + import os + + from datamint import Api, configs + os.environ[configs.ENV_VARS[configs.APIKEY_KEY]] = api_key + + api = Api(check_connection=True) + project = api.projects.create( + name=args.project, + description=f"Imported via 'datamint import' ({fmt} format).", + exists_ok=True, + ) + + import_kwargs = {} + if args.imported_from is not None: + import_kwargs['imported_from'] = args.imported_from + + import_result = importer.import_to_project( + project=project, + api=api, + tags=args.tag, + on_error=args.on_error, + progress_bar=True, + **import_kwargs, + ) + + _USER_LOGGER.info( + f'Uploaded {import_result.n_images_uploaded} image(s) and ' + f'{import_result.n_boxes_uploaded} box annotation(s) to project {args.project!r}.' + ) + if import_result.errors: + _USER_LOGGER.warning(f'{len(import_result.errors)} error(s) during upload:') + for file_name, error in import_result.errors: + _USER_LOGGER.warning(f' {file_name}: {error}') + _USER_LOGGER.info(project.url) + except DatamintException as e: + _USER_LOGGER.error(f'❌ {e}') + sys.exit(1) + except KeyboardInterrupt: + _USER_LOGGER.warning('\nCancelled by user.') + sys.exit(1) + + +if __name__ == '__main__': + main() diff --git a/datamint/importers/__init__.py b/datamint/importers/__init__.py index 11adee3d..61ddbb7d 100644 --- a/datamint/importers/__init__.py +++ b/datamint/importers/__init__.py @@ -1,3 +1,4 @@ +from ._detect import DatasetFormat, DetectedDataset, detect_format, sniff_single_format from .coco import COCOBox, COCOImporter, COCOImportResult, COCOParseResult, COCOSample from .pascal_voc import PascalVOCBox, PascalVOCImporter, PascalVOCImportResult, PascalVOCParseResult, PascalVOCSample from .yolo import YOLOBox, YOLOImporter, YOLOImportResult, YOLOParseResult, YOLOSample @@ -6,4 +7,5 @@ 'COCOImporter', 'COCOParseResult', 'COCOImportResult', 'COCOSample', 'COCOBox', 'PascalVOCImporter', 'PascalVOCParseResult', 'PascalVOCImportResult', 'PascalVOCSample', 'PascalVOCBox', 'YOLOImporter', 'YOLOParseResult', 'YOLOImportResult', 'YOLOSample', 'YOLOBox', + 'DatasetFormat', 'DetectedDataset', 'detect_format', 'sniff_single_format', ] diff --git a/datamint/importers/_detect.py b/datamint/importers/_detect.py new file mode 100644 index 00000000..807fa047 --- /dev/null +++ b/datamint/importers/_detect.py @@ -0,0 +1,108 @@ +"""Best-effort detection of a labeled dataset's annotation format. + +Looks at what's on disk under a root directory and guesses which +``*Importer`` in :mod:`datamint.importers` applies, along with the concrete +paths that importer's constructor needs. Ambiguous or non-standard layouts should +fall back to an explicit format choice + path overrides rather than a +silent wrong guess. +""" +import json +import xml.etree.ElementTree as ET +from dataclasses import dataclass, field +from pathlib import Path +from typing import Literal + +DatasetFormat = Literal['coco', 'yolo', 'pascal_voc'] + +_IMAGE_DIR_NAMES = {'images', 'image', 'imgs', 'jpegimages', 'img'} + + +@dataclass +class DetectedDataset: + """Result of a format sniff: which format, and the constructor kwargs + to feed the matching ``*Importer``.""" + format: DatasetFormat + importer_kwargs: dict[str, Path] = field(default_factory=dict) + + +def _find_dir(root: Path, names: set[str]) -> Path | None: + for p in sorted(root.rglob('*')): + if p.is_dir() and p.name.lower() in names: + return p + return None + + +def sniff_coco(root: Path) -> DetectedDataset | None: + """Look for a JSON file with the COCO ``images``/``annotations``/``categories`` keys.""" + for json_path in sorted(root.rglob('*.json')): + try: + with open(json_path) as f: + data = json.load(f) + except (json.JSONDecodeError, OSError): + continue + if isinstance(data, dict) and {'images', 'annotations', 'categories'} <= data.keys(): + images_dir = _find_dir(root, _IMAGE_DIR_NAMES) or json_path.parent + return DetectedDataset('coco', {'annotations_file': json_path, 'images_dir': images_dir}) + return None + + +def sniff_pascal_voc(root: Path) -> DetectedDataset | None: + """Look for ``.xml`` files whose root element is a Pascal VOC ````.""" + xml_paths = sorted(root.rglob('*.xml')) + for xml_path in xml_paths[:5]: + try: + root_el = ET.parse(xml_path).getroot() + except ET.ParseError: + continue + if root_el.tag == 'annotation' and root_el.find('object') is not None: + annotations_dir = xml_path.parent + images_dir = _find_dir(root, _IMAGE_DIR_NAMES) or annotations_dir + return DetectedDataset('pascal_voc', {'annotations_dir': annotations_dir, 'images_dir': images_dir}) + return None + + +def sniff_yolo(root: Path) -> DetectedDataset | None: + """Look for a ``labels/`` dir of YOLO ``.txt`` files (class + 4 normalized floats).""" + labels_dir = _find_dir(root, {'labels', 'label'}) + if labels_dir is None: + return None + txt_paths = [p for p in sorted(labels_dir.glob('*.txt')) if p.name != 'classes.txt'] + if not txt_paths: + return None + if not any(len(p.read_text().split()) % 5 == 0 and p.read_text().strip() for p in txt_paths[:5]): + return None + + images_dir = _find_dir(root, _IMAGE_DIR_NAMES) or labels_dir.parent / 'images' + kwargs: dict[str, Path] = {'images_dir': images_dir, 'labels_dir': labels_dir} + + data_yaml = next(root.glob('*.yaml'), None) or next(root.glob('*.yml'), None) + if data_yaml is not None: + kwargs['data_yaml'] = data_yaml + return DetectedDataset('yolo', kwargs) + + +_SNIFFERS = { + 'coco': sniff_coco, + 'pascal_voc': sniff_pascal_voc, + 'yolo': sniff_yolo, +} + + +def detect_format(root: str | Path) -> DetectedDataset | None: + """Guess the annotation format of a dataset directory. + + Tries COCO, then Pascal VOC, then YOLO (see the corresponding + ``sniff_*`` function for what triggers each). Returns ``None`` if + nothing matched. + """ + root = Path(root) + for sniff in _SNIFFERS.values(): + result = sniff(root) + if result is not None: + return result + return None + + +def sniff_single_format(root: str | Path, format: DatasetFormat) -> DetectedDataset | None: + """Locate paths for a single, already-chosen format.""" + return _SNIFFERS[format](Path(root)) diff --git a/docs/source/command_line_tools.rst b/docs/source/command_line_tools.rst index 4616afef..4b1cd77c 100644 --- a/docs/source/command_line_tools.rst +++ b/docs/source/command_line_tools.rst @@ -26,10 +26,11 @@ You should see this in the first line: usage: datamint config [-h] [--api-key API_KEY] [--default-url DEFAULT_URL] [-i] [...] -There are six command-line tools available: +There are seven command-line tools available: - ``datamint config``: To configure the Datamint API key and URL. - ``datamint upload``: To upload DICOM, NIfTI, video, image, and segmentation files to the Datamint server. +- ``datamint import``: To import an already-labeled dataset (COCO, Pascal VOC, or YOLO) into a Datamint project. - ``datamint init``: To scaffold a ready-to-run project (upload, train, and deploy scripts). - ``datamint example``: To populate a project with a ready-made example dataset, no data of your own required. - ``datamint train``: To train a model on a Datamint project using a built-in one-line trainer. @@ -267,6 +268,35 @@ See all available options by running ``datamint upload --help``: --version show program's version number and exit --verbose Print debug messages +Importing an already-labeled dataset +-------------------------------------- + +If you already have a labeled dataset in COCO, Pascal VOC, or YOLO format, ``datamint import`` +uploads the images and their bounding-box annotations to a Datamint project in one call, instead +of uploading resources and annotations separately. + +By default, the format and the file/folder layout (annotations file, images directory, etc.) are +auto-detected from the given path. You'll be shown what was detected and asked to confirm before +anything is uploaded: + +.. code-block:: bash + + datamint import ./my_dataset --project MyProject + +If detection guesses wrong, or the layout can't be found automatically, tell it the format explicitly +with ``--format`` and supply the paths auto-detection couldn't infer (``--annotations-file``, +``--annotations-dir``, ``--images-dir``, ``--labels-dir``, etc.): + +.. code-block:: bash + + datamint import ./my_dataset --project MyProject --format yolo \ + --images-dir ./my_dataset/images --labels-dir ./my_dataset/labels + +Use ``--dry-run`` to parse and print a summary (image/box/class counts) without creating the +project or uploading anything, and ``--yes`` to skip the auto-detected-format confirmation prompt. + +See all available options by running ``datamint import --help``. + Populating a project with example data ---------------------------------------