diff --git a/README.md b/README.md index 509e8a8..b7a52b6 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,14 @@ imgeda plot all -m manifest.jsonl imgeda report -m manifest.jsonl ``` -Or just run `imgeda` with no arguments for an interactive wizard that walks you through everything. +Or just run `imgeda` with no arguments for an interactive wizard that walks you through everything: + +```bash +# Interactive mode — auto-detects dataset format (YOLO, COCO, VOC, classification, flat) +imgeda +``` + +The wizard detects your dataset structure, shows a summary panel with image counts, splits, and class info, then lets you pick which splits and analyses to run. ## Features @@ -50,7 +57,8 @@ Or just run `imgeda` with no arguments for an interactive wizard that walks you - **Quality checks**: corrupt files, dark/overexposed images, border artifacts, exact and near-duplicate detection - **7 plot types** with automatic large-dataset adaptations - **Single-page HTML report** with embedded plots and summary tables -- **Interactive configurator** for guided setup +- **Dataset format detection** — auto-detects YOLO, COCO, Pascal VOC, classification, and flat image directories with split-aware scanning +- **Interactive configurator** with Rich panels, split selection, and smart defaults - **Lambda-compatible core** — the analysis functions have zero CLI dependencies, ready for serverless deployment ## Example Output diff --git a/src/imgeda/cli/interactive.py b/src/imgeda/cli/interactive.py index 303e337..343bae9 100644 --- a/src/imgeda/cli/interactive.py +++ b/src/imgeda/cli/interactive.py @@ -7,7 +7,10 @@ import questionary from rich.console import Console +from rich.panel import Panel +from rich.table import Table +from imgeda.core.format_detector import DatasetInfo, detect_format from imgeda.io.image_reader import discover_images from imgeda.models.config import ScanConfig from imgeda.utils import fmt_bytes @@ -17,9 +20,67 @@ DEFAULT_EXTENSIONS = ScanConfig().extensions +def _format_dataset_panel(info: DatasetInfo) -> Panel: + """Build a Rich panel showing detected dataset info.""" + table = Table(show_header=False, box=None, padding=(0, 2)) + table.add_column("Key", style="bold cyan") + table.add_column("Value") + + table.add_row("Format", info.format.upper()) + table.add_row("Images", f"{info.num_images:,} (~{fmt_bytes(info.estimated_size_bytes)})") + + if info.splits: + split_parts = [f"{k} ({v:,})" for k, v in info.splits.items()] + table.add_row("Splits", " \u00b7 ".join(split_parts)) + + if info.num_classes is not None: + names_str = "" + if info.class_names: + preview = ", ".join(info.class_names[:5]) + if info.num_classes > 5: + preview += ", \u2026" + names_str = f" ({preview})" + table.add_row("Classes", f"{info.num_classes}{names_str}") + + if info.annotations_path: + table.add_row("Annotations", str(Path(info.annotations_path).name) + "/") + + return Panel(table, title="Dataset Info", border_style="blue") + + +def _build_split_choices(info: DatasetInfo) -> list[questionary.Choice]: + """Build questionary choices for split selection.""" + choices = [] + for split_name, count in info.splits.items(): + choices.append( + questionary.Choice( + f"{split_name} ({count:,} images)", + value=split_name, + checked=True, + ) + ) + return choices + + +def _resolve_image_dirs(info: DatasetInfo, selected_splits: list[str]) -> list[str]: + """Resolve which image directories to scan based on selected splits.""" + if not info.splits or not selected_splits: + return info.image_dirs + + # For YOLO/COCO/VOC with splits, filter image_dirs to selected splits + selected: list[str] = [] + for img_dir in info.image_dirs: + dir_name = Path(img_dir).name + if dir_name in selected_splits: + selected.append(img_dir) + + # If we couldn't match by dir name, return all image_dirs + return selected if selected else info.image_dirs + + def run_interactive() -> None: """Launch the interactive configuration wizard.""" - console.print("\n[bold blue]Welcome to imgeda — Image Dataset EDA Tool[/bold blue]\n") + console.print("\n[bold blue]Welcome to imgeda \u2014 Image Dataset EDA Tool[/bold blue]\n") # 1. Directory directory = questionary.path( @@ -35,21 +96,43 @@ def run_interactive() -> None: console.print(f"[red]Error: {directory} is not a valid directory[/red]") return - # Quick count - images = discover_images(str(dir_path), DEFAULT_EXTENSIONS) - try: - total_size = sum(os.path.getsize(p) for p in images[:1000]) - except OSError: - total_size = 0 - est_size = total_size * len(images) / min(len(images), 1000) if images else 0 - - console.print(f" Found [bold]{len(images):,}[/bold] images (~{fmt_bytes(est_size)})\n") - - if not images: - console.print("[yellow]No images found. Check the directory path.[/yellow]") - return - - # 2. Analyses + # 2. Format detection + console.print(" Detecting dataset format\u2026\n") + info = detect_format(str(dir_path)) + console.print(_format_dataset_panel(info)) + console.print() + + if info.num_images == 0: + # Fallback: try discover_images on root in case format detector missed something + images = discover_images(str(dir_path), DEFAULT_EXTENSIONS) + if not images: + console.print("[yellow]No images found. Check the directory path.[/yellow]") + return + # Update info with discovered count + try: + total_size = sum(os.path.getsize(p) for p in images[:1000]) + except OSError: + total_size = 0 + est_size = total_size * len(images) / min(len(images), 1000) if images else 0 + console.print(f" Found [bold]{len(images):,}[/bold] images (~{fmt_bytes(est_size)})\n") + + # 3. Split selection (if splits detected) + selected_splits: list[str] = [] + if info.splits: + split_choices = _build_split_choices(info) + selected_splits = questionary.checkbox( + "Which splits to analyze?", + choices=split_choices, + ).ask() + + if selected_splits is None: + return + + if not selected_splits: + console.print("[yellow]No splits selected. Exiting.[/yellow]") + return + + # 4. Analyses analyses = questionary.checkbox( "What analyses would you like to run?", choices=[ @@ -66,10 +149,10 @@ def run_interactive() -> None: skip_pixel_stats = "Pixel statistics (brightness, color channels)" not in analyses include_hashes = "Perceptual hashing (duplicate detection)" in analyses - # 3. Workers + # 5. Workers cpu = os.cpu_count() or 4 workers_str = questionary.text( - f"How many workers? (default: {cpu}, your CPU has {cpu} cores)", + f"How many workers? ({cpu} cores available)", default=str(cpu), ).ask() @@ -78,7 +161,7 @@ def run_interactive() -> None: workers = int(workers_str) if workers_str.isdigit() else cpu - # 4. Output + # 6. Output output = questionary.text( "Output manifest path?", default="./imgeda_manifest.jsonl", @@ -87,13 +170,13 @@ def run_interactive() -> None: if output is None: return - # 5. Plots - generate_plots = questionary.confirm( - "Generate plots after scanning?", + # 7. Combined report prompt + generate_report = questionary.confirm( + "Generate plots and HTML report after scanning?", default=True, ).ask() - if generate_plots is None: + if generate_report is None: return # Build config and run @@ -103,14 +186,19 @@ def run_interactive() -> None: skip_pixel_stats=skip_pixel_stats, ) - console.print("\n[bold]Starting scan...[/bold]\n") + # Resolve which directories to scan + scan_dirs = _resolve_image_dirs(info, selected_splits) + + console.print("\n[bold]Starting scan\u2026[/bold]\n") from imgeda.pipeline.runner import run_scan - run_scan(str(dir_path), output, config) + # Scan each directory (or the root if flat) + for scan_dir in scan_dirs: + run_scan(scan_dir, output, config) - # Generate plots if requested - if generate_plots: + # Generate plots and report if requested + if generate_report: from imgeda.io.manifest_io import read_manifest from imgeda.models.config import PlotConfig from imgeda.plotting.aspect_ratio import plot_aspect_ratio @@ -120,7 +208,7 @@ def run_interactive() -> None: from imgeda.plotting.file_size import plot_file_size from imgeda.plotting.pixel_stats import plot_brightness, plot_channels - console.print("\n[bold]Generating plots...[/bold]\n") + console.print("\n[bold]Generating plots\u2026[/bold]\n") _meta, records = read_manifest(output) plot_config = PlotConfig(output_dir="./plots") @@ -139,4 +227,13 @@ def run_interactive() -> None: except Exception as e: console.print(f" [red]{name}: {e}[/red]") + # Generate HTML report + console.print("\n[bold]Generating HTML report\u2026[/bold]\n") + from imgeda.cli.report import report as report_cmd + + try: + report_cmd(manifest=output, output="./imgeda_report.html") + except SystemExit: + pass # typer.Exit raised when manifest has no records + console.print("\n[bold green]Done![/bold green]") diff --git a/src/imgeda/core/format_detector.py b/src/imgeda/core/format_detector.py new file mode 100644 index 0000000..9deb21a --- /dev/null +++ b/src/imgeda/core/format_detector.py @@ -0,0 +1,395 @@ +"""Dataset format detection — probes directory structure to identify ML dataset formats.""" + +from __future__ import annotations + +import json +import os +from dataclasses import dataclass, field +from pathlib import Path + +IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif", ".webp", ".gif"} + + +@dataclass(slots=True) +class DatasetInfo: + """Detected dataset format and metadata.""" + + format: str # "yolo" | "coco" | "voc" | "classification" | "flat" + image_dirs: list[str] # paths to image directories (may be per-split) + num_images: int + estimated_size_bytes: int + splits: dict[str, int] # {"train": 2940, "val": 740} or {} for flat + num_classes: int | None = None + class_names: list[str] | None = None + annotations_path: str | None = None + extra: dict[str, str] = field(default_factory=dict) + + +def detect_format(root: str) -> DatasetInfo: + """Detect ML dataset format by inspecting directory structure. + + Checks in order: YOLO, COCO, Pascal VOC, Classification, Flat (fallback). + """ + root_path = Path(root) + + # 1. YOLO — data.yaml at root + info = _try_yolo(root_path) + if info is not None: + return info + + # 2. COCO — annotations/*.json with COCO keys + info = _try_coco(root_path) + if info is not None: + return info + + # 3. Pascal VOC — Annotations/ + JPEGImages/ + info = _try_voc(root_path) + if info is not None: + return info + + # 4. Classification — >3 subdirs each containing images + info = _try_classification(root_path) + if info is not None: + return info + + # 5. Flat fallback + return _build_flat(root_path) + + +def _count_images_in(directory: Path) -> int: + """Count image files recursively under a directory.""" + count = 0 + if not directory.is_dir(): + return 0 + for dirpath, _dirnames, filenames in os.walk(directory): + for fn in filenames: + if Path(fn).suffix.lower() in IMAGE_EXTENSIONS: + count += 1 + return count + + +def _estimate_size(directory: Path, sample_limit: int = 100) -> int: + """Estimate total image size by sampling up to sample_limit files.""" + sizes: list[int] = [] + total_images = 0 + for dirpath, _dirnames, filenames in os.walk(directory): + for fn in filenames: + if Path(fn).suffix.lower() in IMAGE_EXTENSIONS: + total_images += 1 + if len(sizes) < sample_limit: + try: + sizes.append(os.path.getsize(os.path.join(dirpath, fn))) + except OSError: + pass + if not sizes: + return 0 + avg = sum(sizes) / len(sizes) + return int(avg * total_images) + + +def _parse_simple_yaml(path: Path) -> dict[str, str | list[str]]: + """Parse a simple YAML file without pyyaml dependency. + + Handles basic key: value pairs and simple lists (names: [...] or + names:\n - item lines). Enough for YOLO data.yaml files. + """ + result: dict[str, str | list[str]] = {} + try: + text = path.read_text(encoding="utf-8") + except (OSError, UnicodeDecodeError): + return result + + lines = text.splitlines() + current_key: str | None = None + current_list: list[str] | None = None + + for line in lines: + stripped = line.strip() + if not stripped or stripped.startswith("#"): + continue + + # Check for list continuation (indented "- item") + if current_key is not None and current_list is not None: + if stripped.startswith("- "): + item = stripped[2:].strip().strip("'\"") + current_list.append(item) + continue + else: + # End of list + result[current_key] = current_list + current_key = None + current_list = None + + if ":" not in stripped: + continue + + key, _, value = stripped.partition(":") + key = key.strip() + value = value.strip() + + if not value: + # Could be start of a list + current_key = key + current_list = [] + elif value.startswith("[") and value.endswith("]"): + # Inline list: names: [cat, dog, bird] + items = value[1:-1].split(",") + result[key] = [item.strip().strip("'\"") for item in items if item.strip()] + else: + result[key] = value.strip("'\"") + + # Flush any pending list + if current_key is not None and current_list is not None: + result[current_key] = current_list + + return result + + +def _try_yolo(root: Path) -> DatasetInfo | None: + """Detect YOLO format via data.yaml.""" + yaml_path = root / "data.yaml" + if not yaml_path.is_file(): + return None + + parsed = _parse_simple_yaml(yaml_path) + if not parsed: + return None + + # Extract class names + names_val = parsed.get("names") + class_names: list[str] | None = None + num_classes: int | None = None + if isinstance(names_val, list): + class_names = names_val[:10] + num_classes = len(names_val) + elif "nc" in parsed: + nc_val = parsed["nc"] + if isinstance(nc_val, str) and nc_val.isdigit(): + num_classes = int(nc_val) + + # Detect splits from data.yaml or directory structure + splits: dict[str, int] = {} + image_dirs: list[str] = [] + + for split_name in ("train", "val", "test"): + split_val = parsed.get(split_name) + if isinstance(split_val, str): + split_path = root / split_val + # YOLO convention: images dir mirrors the path + # data.yaml may point to images/train or just train + if split_path.is_dir(): + count = _count_images_in(split_path) + if count > 0: + splits[split_name] = count + image_dirs.append(str(split_path)) + continue + # Try under images/ + img_split = root / "images" / split_name + if img_split.is_dir(): + count = _count_images_in(img_split) + if count > 0: + splits[split_name] = count + image_dirs.append(str(img_split)) + + # Fallback: check images/ dir directly + if not image_dirs: + images_dir = root / "images" + if images_dir.is_dir(): + image_dirs.append(str(images_dir)) + + # Detect annotations path + labels_dir = root / "labels" + annotations_path = str(labels_dir) if labels_dir.is_dir() else None + + num_images = ( + sum(splits.values()) if splits else sum(_count_images_in(Path(d)) for d in image_dirs) + ) + + return DatasetInfo( + format="yolo", + image_dirs=image_dirs or [str(root)], + num_images=num_images, + estimated_size_bytes=_estimate_size(root), + splits=splits, + num_classes=num_classes, + class_names=class_names, + annotations_path=annotations_path, + extra={}, + ) + + +def _try_coco(root: Path) -> DatasetInfo | None: + """Detect COCO format via annotations/*.json with COCO keys.""" + ann_dir = root / "annotations" + if not ann_dir.is_dir(): + return None + + json_files = list(ann_dir.glob("*.json")) + if not json_files: + return None + + # Check first JSON for COCO structure + coco_file: Path | None = None + categories: list[dict[str, str]] = [] + for jf in json_files: + try: + # Read first portion to check keys + with open(jf, encoding="utf-8") as f: + data = json.load(f) + if isinstance(data, dict) and "images" in data and "annotations" in data: + coco_file = jf + categories = data.get("categories", []) + break + except (json.JSONDecodeError, OSError): + continue + + if coco_file is None: + return None + + # Extract class info + class_names: list[str] | None = None + num_classes: int | None = None + if categories: + all_names = [c.get("name", "") for c in categories if isinstance(c, dict)] + num_classes = len(all_names) + class_names = all_names[:10] if all_names else None + + # Detect splits from annotation filenames + splits: dict[str, int] = {} + image_dirs: list[str] = [] + + images_dir = root / "images" + if images_dir.is_dir(): + image_dirs.append(str(images_dir)) + + # Try to detect splits from annotation file names (e.g. instances_train2017.json) + for jf in json_files: + name = jf.stem.lower() + for split_name in ("train", "val", "test"): + if split_name in name: + try: + with open(jf, encoding="utf-8") as f: + data = json.load(f) + if isinstance(data, dict) and "images" in data: + splits[split_name] = len(data["images"]) + except (json.JSONDecodeError, OSError): + pass + + num_images = ( + sum(splits.values()) + if splits + else (_count_images_in(images_dir) if images_dir.is_dir() else 0) + ) + + return DatasetInfo( + format="coco", + image_dirs=image_dirs or [str(root)], + num_images=num_images, + estimated_size_bytes=_estimate_size(root), + splits=splits, + num_classes=num_classes, + class_names=class_names, + annotations_path=str(ann_dir), + extra={}, + ) + + +def _try_voc(root: Path) -> DatasetInfo | None: + """Detect Pascal VOC format via Annotations/ + JPEGImages/.""" + ann_dir = root / "Annotations" + img_dir = root / "JPEGImages" + + if not ann_dir.is_dir() or not img_dir.is_dir(): + return None + + # Check for XML files in Annotations + xml_files = list(ann_dir.glob("*.xml")) + if not xml_files: + return None + + num_images = _count_images_in(img_dir) + + # Check for ImageSets/Main/ split files + splits: dict[str, int] = {} + imagesets_dir = root / "ImageSets" / "Main" + if imagesets_dir.is_dir(): + for txt_file in imagesets_dir.glob("*.txt"): + split_name = txt_file.stem.lower() + if split_name in ("train", "val", "test", "trainval"): + try: + lines = txt_file.read_text().strip().splitlines() + count = len([ln for ln in lines if ln.strip()]) + if count > 0: + splits[split_name] = count + except OSError: + pass + + return DatasetInfo( + format="voc", + image_dirs=[str(img_dir)], + num_images=num_images, + estimated_size_bytes=_estimate_size(img_dir), + splits=splits, + num_classes=None, + class_names=None, + annotations_path=str(ann_dir), + extra={}, + ) + + +def _try_classification(root: Path) -> DatasetInfo | None: + """Detect classification format: >3 subdirs each containing images.""" + # Skip if annotation-style dirs exist + for ann_dir_name in ("labels", "annotations", "Annotations"): + if (root / ann_dir_name).is_dir(): + return None + + subdirs = [d for d in root.iterdir() if d.is_dir() and not d.name.startswith(".")] + + if len(subdirs) < 3: + return None + + # Check that most subdirs contain images + counts: dict[str, int] = {} + total_images = 0 + for sd in subdirs: + count = _count_images_in(sd) + counts[sd.name] = count + if count > 0: + total_images += count + + dirs_with_images = sum(1 for c in counts.values() if c > 0) + + # Require majority of subdirs to contain images + if dirs_with_images < len(subdirs) * 0.5: + return None + + class_names_all = sorted(name for name, c in counts.items() if c > 0) + + return DatasetInfo( + format="classification", + image_dirs=[str(root)], + num_images=total_images, + estimated_size_bytes=_estimate_size(root), + splits={}, + num_classes=len(class_names_all), + class_names=class_names_all[:10], + annotations_path=None, + extra={}, + ) + + +def _build_flat(root: Path) -> DatasetInfo: + """Fallback: flat directory with images.""" + num_images = _count_images_in(root) + return DatasetInfo( + format="flat", + image_dirs=[str(root)], + num_images=num_images, + estimated_size_bytes=_estimate_size(root) if num_images > 0 else 0, + splits={}, + num_classes=None, + class_names=None, + annotations_path=None, + extra={}, + ) diff --git a/tests/test_format_detector.py b/tests/test_format_detector.py new file mode 100644 index 0000000..90f6f62 --- /dev/null +++ b/tests/test_format_detector.py @@ -0,0 +1,231 @@ +"""Tests for dataset format detection.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +from PIL import Image + +from imgeda.core.format_detector import DatasetInfo, detect_format + + +def _create_image(path: Path, w: int = 100, h: int = 100) -> None: + """Create a small test image at the given path.""" + path.parent.mkdir(parents=True, exist_ok=True) + arr = np.random.randint(60, 200, (h, w, 3), dtype=np.uint8) + Image.fromarray(arr).save(path) + + +class TestYoloDetection: + def test_yolo_with_data_yaml(self, tmp_path: Path) -> None: + # Create data.yaml + (tmp_path / "data.yaml").write_text( + "train: images/train\nval: images/val\nnc: 3\nnames: [cat, dog, bird]\n" + ) + # Create images dirs + for split in ("train", "val"): + for i in range(3): + _create_image(tmp_path / "images" / split / f"img_{i}.jpg") + # Create labels dir + (tmp_path / "labels" / "train").mkdir(parents=True) + (tmp_path / "labels" / "train" / "img_0.txt").write_text("0 0.5 0.5 0.1 0.1\n") + + info = detect_format(str(tmp_path)) + assert info.format == "yolo" + assert "train" in info.splits + assert "val" in info.splits + assert info.splits["train"] == 3 + assert info.splits["val"] == 3 + assert info.num_classes == 3 + assert info.class_names == ["cat", "dog", "bird"] + assert info.annotations_path is not None + assert info.num_images == 6 + + def test_yolo_with_list_names(self, tmp_path: Path) -> None: + (tmp_path / "data.yaml").write_text( + "train: images/train\nval: images/val\nnames:\n - cat\n - dog\n - bird\n" + ) + for split in ("train", "val"): + _create_image(tmp_path / "images" / split / "img_0.jpg") + (tmp_path / "labels").mkdir() + + info = detect_format(str(tmp_path)) + assert info.format == "yolo" + assert info.num_classes == 3 + assert info.class_names == ["cat", "dog", "bird"] + + def test_yolo_no_images(self, tmp_path: Path) -> None: + """data.yaml exists but no images — still detects as YOLO.""" + (tmp_path / "data.yaml").write_text("nc: 2\nnames: [a, b]\n") + info = detect_format(str(tmp_path)) + assert info.format == "yolo" + assert info.num_images == 0 + + +class TestCocoDetection: + def test_coco_format(self, tmp_path: Path) -> None: + ann_dir = tmp_path / "annotations" + ann_dir.mkdir() + img_dir = tmp_path / "images" + img_dir.mkdir() + + coco_data = { + "images": [ + {"id": 1, "file_name": "img_0.jpg"}, + {"id": 2, "file_name": "img_1.jpg"}, + ], + "annotations": [ + {"id": 1, "image_id": 1, "category_id": 1}, + ], + "categories": [ + {"id": 1, "name": "cat"}, + {"id": 2, "name": "dog"}, + ], + } + (ann_dir / "instances_train.json").write_text(json.dumps(coco_data)) + + for i in range(3): + _create_image(img_dir / f"img_{i}.jpg") + + info = detect_format(str(tmp_path)) + assert info.format == "coco" + assert info.num_classes == 2 + assert info.class_names == ["cat", "dog"] + assert info.annotations_path is not None + assert "train" in info.splits + assert info.splits["train"] == 2 + + def test_coco_no_categories(self, tmp_path: Path) -> None: + ann_dir = tmp_path / "annotations" + ann_dir.mkdir() + coco_data = {"images": [{"id": 1}], "annotations": [{"id": 1}]} + (ann_dir / "data.json").write_text(json.dumps(coco_data)) + + info = detect_format(str(tmp_path)) + assert info.format == "coco" + assert info.num_classes is None + + def test_non_coco_json_skipped(self, tmp_path: Path) -> None: + """JSON without COCO keys should not match.""" + ann_dir = tmp_path / "annotations" + ann_dir.mkdir() + (ann_dir / "config.json").write_text(json.dumps({"key": "value"})) + + info = detect_format(str(tmp_path)) + assert info.format != "coco" + + +class TestVocDetection: + def test_voc_format(self, tmp_path: Path) -> None: + ann_dir = tmp_path / "Annotations" + ann_dir.mkdir() + img_dir = tmp_path / "JPEGImages" + img_dir.mkdir() + + (ann_dir / "img_0.xml").write_text("") + (ann_dir / "img_1.xml").write_text("") + + for i in range(5): + _create_image(img_dir / f"img_{i}.jpg") + + info = detect_format(str(tmp_path)) + assert info.format == "voc" + assert info.num_images == 5 + assert info.annotations_path is not None + + def test_voc_with_imagesets(self, tmp_path: Path) -> None: + (tmp_path / "Annotations").mkdir() + (tmp_path / "JPEGImages").mkdir() + (tmp_path / "Annotations" / "a.xml").write_text("") + _create_image(tmp_path / "JPEGImages" / "a.jpg") + + imagesets = tmp_path / "ImageSets" / "Main" + imagesets.mkdir(parents=True) + (imagesets / "train.txt").write_text("img_0\nimg_1\nimg_2\n") + (imagesets / "val.txt").write_text("img_3\nimg_4\n") + + info = detect_format(str(tmp_path)) + assert info.format == "voc" + assert info.splits.get("train") == 3 + assert info.splits.get("val") == 2 + + def test_voc_needs_both_dirs(self, tmp_path: Path) -> None: + """Only Annotations/ without JPEGImages/ should not match VOC.""" + (tmp_path / "Annotations").mkdir() + (tmp_path / "Annotations" / "a.xml").write_text("") + + info = detect_format(str(tmp_path)) + assert info.format != "voc" + + +class TestClassificationDetection: + def test_classification_format(self, tmp_path: Path) -> None: + for cls in ("cat", "dog", "bird", "fish"): + for i in range(3): + _create_image(tmp_path / cls / f"img_{i}.jpg") + + info = detect_format(str(tmp_path)) + assert info.format == "classification" + assert info.num_classes == 4 + assert info.class_names is not None + assert "cat" in info.class_names + assert info.num_images == 12 + + def test_classification_needs_3_plus_subdirs(self, tmp_path: Path) -> None: + """Only 2 subdirs should not match classification.""" + for cls in ("cat", "dog"): + _create_image(tmp_path / cls / "img_0.jpg") + + info = detect_format(str(tmp_path)) + assert info.format != "classification" + + def test_classification_skipped_with_labels_dir(self, tmp_path: Path) -> None: + """If labels/ exists, should not match classification (probably YOLO).""" + for cls in ("cat", "dog", "bird", "fish"): + _create_image(tmp_path / cls / "img_0.jpg") + (tmp_path / "labels").mkdir() + + info = detect_format(str(tmp_path)) + assert info.format != "classification" + + +class TestFlatDetection: + def test_flat_format(self, tmp_path: Path) -> None: + for i in range(5): + _create_image(tmp_path / f"img_{i}.jpg") + + info = detect_format(str(tmp_path)) + assert info.format == "flat" + assert info.num_images == 5 + + def test_empty_directory(self, tmp_path: Path) -> None: + info = detect_format(str(tmp_path)) + assert info.format == "flat" + assert info.num_images == 0 + assert info.estimated_size_bytes == 0 + + def test_flat_with_subdirs(self, tmp_path: Path) -> None: + """Only 1-2 subdirs with images -> flat, not classification.""" + _create_image(tmp_path / "img_0.jpg") + _create_image(tmp_path / "subdir" / "img_1.jpg") + + info = detect_format(str(tmp_path)) + assert info.format == "flat" + assert info.num_images == 2 + + +class TestDatasetInfo: + def test_dataclass_defaults(self) -> None: + info = DatasetInfo( + format="flat", + image_dirs=["/tmp"], + num_images=0, + estimated_size_bytes=0, + splits={}, + ) + assert info.num_classes is None + assert info.class_names is None + assert info.annotations_path is None + assert info.extra == {}