Add dataset format detection and enhanced interactive wizard - #1
Merged
Conversation
Detect YOLO, COCO, Pascal VOC, classification, and flat dataset formats automatically when the interactive wizard starts. Show a Rich panel with format, image count, splits, classes, and annotations. Add split selection for structured datasets and combined plots+report generation prompt. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
![amazon-q-developer[bot]](https://avatars.githubusercontent.com/in/1220912?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary
This PR adds dataset format detection and enhances the interactive wizard with split-aware scanning. The implementation is comprehensive with good test coverage (16 new tests), but contains 5 critical issues that block merge:
Critical Issues Requiring Fixes
- Security Vulnerability (format_detector.py:177-192): Path traversal vulnerability in YOLO split path processing - user-controlled paths from data.yaml are used without validation
- Crash Risk (interactive.py:112-116): Division by zero when no images are discovered in fallback path
- Logic Error (format_detector.py:265-276): Duplicate COCO split names overwrite earlier entries, causing incorrect image counts
- Crash Risk (format_detector.py:338-345): Unhandled OSError/PermissionError when iterating directories
- Performance Regression (format_detector.py:350-365): Double iteration in classification detection degrades performance on large datasets
All issues have code suggestions attached for quick resolution. Please address these defects before merge.
You can now have the agent implement changes and create commits directly on your pull request's source branch. Simply comment with /q followed by your request in natural language to ask the agent to make changes.
Comment on lines
+177
to
+192
| 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)) |
There was a problem hiding this comment.
🛑 Security Vulnerability: Path traversal risk when processing YOLO data.yaml split paths. User-controlled YAML values are used directly in Path operations without validation, allowing directory traversal attacks.1
Replace lines 177-192 with validation to ensure split paths remain within the root directory.
Suggested change
| 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)) | |
| split_path = root / split_val | |
| # Validate path stays within root directory | |
| try: | |
| split_path_resolved = split_path.resolve() | |
| root_resolved = root.resolve() | |
| split_path_resolved.relative_to(root_resolved) | |
| except (ValueError, OSError): | |
| # Path escapes root or cannot be resolved - skip it | |
| continue | |
| # 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)) |
Footnotes
-
CWE-22: Path Traversal - https://cwe.mitre.org/data/definitions/22.html ↩
Comment on lines
+112
to
+116
| 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 |
There was a problem hiding this comment.
🛑 Crash Risk: Division by zero when len(images) is 0. If discover_images returns an empty list, the expression min(len(images), 1000) returns 0, causing a division by zero error on line 116.
Suggested change
| 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 | |
| try: | |
| total_size = sum(os.path.getsize(p) for p in images[:1000]) | |
| except OSError: | |
| total_size = 0 | |
| sample_count = min(len(images), 1000) | |
| est_size = total_size * len(images) / sample_count if sample_count > 0 else 0 |
Comment on lines
+265
to
+276
| # 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 |
There was a problem hiding this comment.
🛑 Logic Error: Duplicate split detection overwrites earlier splits. When multiple JSON files contain the same split name (e.g., two files with "train" in the name), the second one overwrites the first without warning or merging counts, causing incorrect image counts.
Suggested change
| # 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 | |
| # 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 and split_name not in splits: | |
| 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 |
Comment on lines
+338
to
+345
| 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(".")] |
There was a problem hiding this comment.
🛑 Crash Risk: Unhandled exception when directory operations fail. If root.iterdir() fails due to permission issues or I/O errors, the function crashes instead of gracefully falling back to the next detection method.
Suggested change
| 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(".")] | |
| 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 | |
| try: | |
| subdirs = [d for d in root.iterdir() if d.is_dir() and not d.name.startswith(".")] | |
| except (OSError, PermissionError): | |
| return None |
Comment on lines
+350
to
+365
| # Check that most subdirs contain images | ||
| dirs_with_images = 0 | ||
| total_images = 0 | ||
| for sd in subdirs: | ||
| count = _count_images_in(sd) | ||
| if count > 0: | ||
| dirs_with_images += 1 | ||
| total_images += count | ||
|
|
||
| # Require majority of subdirs to contain images | ||
| if dirs_with_images < len(subdirs) * 0.5: | ||
| return None | ||
|
|
||
| class_names_all = sorted( | ||
| sd.name for sd in subdirs if _count_images_in(sd) > 0 | ||
| ) |
There was a problem hiding this comment.
🛑 Performance Regression: Double iteration over subdirs in classification detection. _count_images_in(sd) is called twice for each subdir - once in the counting loop (line 354) and again in the list comprehension (line 364), causing significant slowdown for large datasets.
Suggested change
| # Check that most subdirs contain images | |
| dirs_with_images = 0 | |
| total_images = 0 | |
| for sd in subdirs: | |
| count = _count_images_in(sd) | |
| if count > 0: | |
| dirs_with_images += 1 | |
| total_images += count | |
| # Require majority of subdirs to contain images | |
| if dirs_with_images < len(subdirs) * 0.5: | |
| return None | |
| class_names_all = sorted( | |
| sd.name for sd in subdirs if _count_images_in(sd) > 0 | |
| ) | |
| # Check that most subdirs contain images | |
| dirs_with_images = 0 | |
| total_images = 0 | |
| subdir_counts = {} | |
| for sd in subdirs: | |
| count = _count_images_in(sd) | |
| if count > 0: | |
| dirs_with_images += 1 | |
| total_images += count | |
| subdir_counts[sd.name] = count | |
| # Require majority of subdirs to contain images | |
| if dirs_with_images < len(subdirs) * 0.5: | |
| return None | |
| class_names_all = sorted(subdir_counts.keys()) |
Apply ruff format to pass CI, add slots=True to DatasetInfo dataclass per project conventions, and cache _count_images_in results in _try_classification to avoid walking each subdir twice. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
format_detector.pymodule that auto-detects YOLO, COCO, Pascal VOC, classification, and flat dataset formats by inspecting directory structureTest plan
uv run pytest tests/test_format_detector.py -v— all 16 detection tests passuv run pytest tests/ -q— full suite passes (65 tests)uv run ruff check src/ tests/— cleanuv run mypy src/imgeda/— clean (37 source files)detect_formaton oxford-petdata/— correctly detects "flat" with 3,680 images--help,--version,scan --help)🤖 Generated with Claude Code