COCO Plugin¶

def load(
    self,
    path: str | os.PathLike[str],
    name: str | None = None,
    images_dir: str | os.PathLike[str] | None = None,
    **_kwargs: t.Any,
) -> Dataset:
    """Load COCO dataset.

    Args:
        path: Path to COCO annotation JSON file
        name: Optional dataset name. If not provided, the filename (without
            extension) will be used.
        images_dir: Optional custom path to images directory
        **_kwargs: Additional parameters (ignored)

    Returns:
        Loaded Dataset
    """
    annotation_path = pathlib.Path(path)

    if not annotation_path.exists():
        raise FileNotFoundError(f"Annotation file not found: {annotation_path}")

    # Auto-detect images directory
    if images_dir is not None:
        images_dir = pathlib.Path(images_dir)
    else:
        images_dir = self._find_images_dir(annotation_path)

    with annotation_path.open(mode="r") as f:
        coco_data = json.load(f)

    dataset_name = name or annotation_path.stem
    dataset = Dataset(name=dataset_name)

    logger.info(f"Loading COCO dataset from {annotation_path}")

    # Load categories
    for cat in coco_data.get("categories", []):
        dataset.add_category(cat["id"], cat["name"])

    logger.info(f"Loaded {len(dataset.categories)} categories")

    # Load images
    images_found = 0
    images_missing = 0

    for img in coco_data.get("images", []):
        img_path = self._find_image_path(img["file_name"], images_dir)

        if img_path:
            images_found += 1
        else:
            images_missing += 1

        img_info = ImageInfo(
            image_id=str(img["id"]),
            file_name=img["file_name"],
            width=img["width"],
            height=img["height"],
            path=img_path,
        )
        dataset.add_image(img_info, source_name=dataset_name)

    logger.info(
        f"Loaded {len(dataset.images)} images (found: {images_found}, missing: {images_missing})"
    )

    # Load annotations
    for ann in coco_data.get("annotations", []):
        bbox_xywh = ann["bbox"]
        bbox = BBox.from_xywh(*bbox_xywh)

        cat_name = dataset.get_category_name(ann["category_id"])
        if cat_name is None:
            continue

        annotation = Annotation(
            bbox=bbox,
            category_id=ann["category_id"],
            category_name=cat_name,
            image_id=str(ann["image_id"]),
            annotation_id=str(ann["id"]),
            area=ann.get("area", bbox.area),
            iscrowd=ann.get("iscrowd", 0),
        )
        dataset.add_annotation(annotation)

    logger.info(f"Loaded {dataset.num_annotations()} annotations")

    return dataset

def export(
    self,
    dataset: Dataset,
    output_dir: str | os.PathLike[str],
    split_ratio: SplitRatio | None = None,
    seed: int | None = None,
    naming_strategy: NamingStrategy | None = None,
    copy_images: bool = True,
    unified_structure: bool = False,
    **_kwargs: t.Any,
) -> None:
    """Export to COCO format.

    Args:
        dataset: Dataset to export
        output_dir: Output directory
        split_ratio: Optional split ratios
        seed: Random seed
        naming_strategy: Naming strategy
        copy_images: Whether to copy images
        unified_structure: Use unified directory structure
        **_kwargs: Additional parameters (ignored)
    """
    output_dir = pathlib.Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    naming_strategy = naming_strategy or OriginalNaming()

    logger.info(f"Exporting COCO dataset to {output_dir}")

    image_counter = 0
    ann_counter = 0
    if split_ratio is None:
        all_image_ids = list(dataset.images.keys())
        self._export_split(
            dataset,
            output_dir,
            "train",
            all_image_ids,
            naming_strategy,
            copy_images,
            unified_structure,
        )
    else:
        splits = dataset.split(split_ratio, seed)
        for split_name, image_ids in splits.items():
            if image_ids:
                image_counter, ann_counter = self._export_split(
                    dataset,
                    output_dir,
                    split_name,
                    image_ids,
                    naming_strategy,
                    copy_images,
                    unified_structure,
                    image_counter,
                    ann_counter,
                )

    logger.info(f"COCO dataset exported to: {output_dir}")