YOLO Plugin¶

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

    Loads a YOLO dataset from the specified directory. The directory should
    contain a YAML configuration file and subdirectories for images and
    labels.

    Args:
        path: Path to YOLO dataset YAML configuration file.
        name: Optional custom name for the dataset. If None, uses directory
            name.
        splits: Which split(s) to load. Can be:
            - None: Load all splits (train, val, test)
            - str: Load single split (e.g., "train")
            - list[str]: Load specific splits (e.g., ["train", "val"])
        **_kwargs: Additional parameters (currently unused, reserved for
            future extensions).

    Returns:
        Loaded Dataset instance containing all images, annotations, and
        categories.

    Raises:
        FileNotFoundError: If the YAML configuration file is not found.
        ValueError: If YAML configuration is missing required 'names' field.
    """
    yaml_path = pathlib.Path(path)

    if not yaml_path.exists():
        raise FileNotFoundError(f"YAML file not found: {yaml_path}")

    # Load YAML configuration
    with yaml_path.open(mode="r") as f:
        yaml_config = yaml.safe_load(f)

    dataset_name = name or yaml_path.name
    dataset = Dataset(name=dataset_name)
    dataset_dir = (
        pathlib.Path(yaml_config["path"]) if "path" in yaml_config else yaml_path.parent
    )

    logger.info(f"Loading YOLOv5 dataset from {dataset_dir}")

    # Load categories
    self._load_categories(yaml_config, dataset)

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

    # Determine splits to load
    splits_to_load = self._determine_splits(splits)

    # Load each split
    total_images = 0
    total_annotations = 0

    for split in splits_to_load:
        images_dir = dataset_dir / "images" / split
        labels_dir = dataset_dir / "labels" / split

        if not images_dir.exists():
            logger.warning(f"Images directory not found for {split}: {images_dir}")
            continue

        split_images, split_annotations = self._load_split(
            dataset,
            images_dir,
            labels_dir,
            total_images,
            total_annotations,
        )

        total_images += split_images
        total_annotations += split_annotations

        logger.info(
            f"Loaded {split} split: {split_images} images, {split_annotations} annotations"
        )

    logger.info(f"Total loaded: {total_images} images, {total_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 dataset to YOLO format.

    Creates a YOLO-compatible dataset with proper directory structure,
    label files, and YAML configuration.

    Args:
        dataset: Dataset instance to export.
        output_dir: Output directory path. Will be created if it doesn't
            exist.
        split_ratio: Optional SplitRatio for train/val/test division. If
            None, exports entire dataset as 'train' split.
        seed: Random seed for reproducible splits. Only used if split_ratio
            is provided.
        naming_strategy: Strategy for generating output file names. If None,
            uses OriginalNaming (preserves original filenames).
        copy_images: If True, copies image files to output directory. If
            False, only creates label files.
        unified_structure: If True, uses 'annotations' directory instead of
            labels'. Useful for compatibility with some training frameworks.
        **_kwargs: Additional parameters (currently unused, reserved for
            future extensions).

    Note:
        Category IDs in label files are 0-indexed (YOLO convention), even
        though the Dataset uses 1-indexed IDs internally.
    """
    output_dir = pathlib.Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    naming_strategy = naming_strategy or OriginalNaming()

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

    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,
        )
        splits_to_write = ["train"]
    else:
        splits = dataset.split(split_ratio, seed)
        splits_to_write = []
        for split_name, image_ids in splits.items():
            if image_ids:
                self._export_split(
                    dataset,
                    output_dir,
                    split_name,
                    image_ids,
                    naming_strategy,
                    copy_images,
                    unified_structure,
                )
                splits_to_write.append(split_name)

    # Create data.yaml
    self._create_yaml(dataset, output_dir, splits_to_write)

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