From 23832b93733db22c64437c925cdcc7515987d7fb Mon Sep 17 00:00:00 2001 From: gtfrans2re Date: Mon, 28 Sep 2026 14:03:49 -0400 Subject: [PATCH] Fix #529: support YOLO datasets in the image detection evaluator Signed-off-by: gtfrans2re --- gui/tasks/image_detection/dataset_viewer.py | 160 +++++++++++--------- gui/tasks/image_detection/evaluator.py | 57 +++---- tests/gui/test_image_detection.py | 82 +++++++++- 3 files changed, 191 insertions(+), 108 deletions(-) diff --git a/gui/tasks/image_detection/dataset_viewer.py b/gui/tasks/image_detection/dataset_viewer.py index ad6c58fb..96690996 100644 --- a/gui/tasks/image_detection/dataset_viewer.py +++ b/gui/tasks/image_detection/dataset_viewer.py @@ -1,5 +1,6 @@ import streamlit as st import os +import tempfile from perceptionmetrics.datasets.coco import find_img_dir_and_ann_file from gui.tasks.utils import render_image_grid @@ -7,9 +8,6 @@ def render_image_detection_viewer(): """Render the image detection dataset viewer tab in Streamlit.""" - import tempfile - from perceptionmetrics.datasets.coco import CocoDataset - from perceptionmetrics.datasets.yolo import YOLODataset import numpy as np from PIL import Image from supervision.draw.color import ColorPalette @@ -28,94 +26,40 @@ def render_image_detection_viewer(): st.warning("⚠️ Please select a valid dataset folder.") return - # Setup paths and pagination + # Setup image directory used for pagination if dataset_type == "coco": try: - img_dir, ann_file = find_img_dir_and_ann_file( + img_dir, _ = find_img_dir_and_ann_file( dataset_path=dataset_path, split=split ) except FileNotFoundError: st.warning("Dataset files not found. Check path and split.") return - elif dataset_type == "yolo": - dataset_config_file = st.session_state.get("dataset_config_file", None) - img_dir = os.path.join(dataset_path, f"images/{split}") + img_dir = os.path.join(dataset_path, "images", split) if not os.path.isdir(img_dir): st.warning("Image directory not found.") return - if dataset_config_file is None: - st.warning("Dataset configuration file not found. Please upload it.") - return else: st.error("Unsupported dataset type.") return - # Load dataset - dataset_key = f"{dataset_path}_{split}" + # Load dataset (shared with the evaluator tab through session state) + dataset_key = get_image_detection_dataset_key(dataset_type, dataset_path, split) if dataset_key not in st.session_state: try: - if dataset_type == "coco": - st.session_state[dataset_key] = CocoDataset( - annotation_file=ann_file, - image_dir=img_dir, - split=split, - ) - elif dataset_type == "yolo": - if dataset_config_file is not None: - # Save uploaded config file to a temporary location - with tempfile.NamedTemporaryFile( - delete=False, suffix=".yaml" - ) as tmp: - tmp.write(dataset_config_file.read()) - tmp_path = tmp.name - - # Load YOLO dataset - yolo_dataset = YOLODataset(tmp_path, dataset_path) - st.session_state["full_dataset_df"] = yolo_dataset.dataset - - # Filter dataset for the selected split - yolo_dataset.dataset = yolo_dataset.dataset[ - yolo_dataset.dataset["split"] == split - ].reset_index(drop=True) - st.session_state[dataset_key] = yolo_dataset - - os.unlink(tmp_path) # Clean up temp file - else: - st.warning( - "Dataset configuration file not found. Please upload it." - ) - return - else: - st.error("Unsupported dataset type.") - return - + st.session_state[dataset_key] = load_image_detection_dataset( + dataset_type, + dataset_path, + split, + st.session_state.get("dataset_config_file"), + ) + except FileNotFoundError as e: + st.warning(f"⚠️ {e}") + return except Exception as e: st.error(f"Failed to load dataset: {e}") return - else: - # Ensure cached dataset has the correct split; if not, rebuild it - try: - cached_ds = st.session_state[dataset_key] - cached_split = getattr(cached_ds, "split", None) - if cached_split != split: - if dataset_type == "coco": - st.session_state[dataset_key] = CocoDataset( - annotation_file=ann_file, - image_dir=img_dir, - split=split, - ) - elif dataset_type == "yolo": - yolo_dataset = st.session_state[dataset_key] - yolo_dataset.dataset = st.session_state["full_dataset_df"][ - st.session_state["full_dataset_df"]["split"] == split - ].reset_index(drop=True) - st.session_state[dataset_key] = yolo_dataset - else: - st.error("Unsupported dataset type.") - return - except Exception: - pass dataset = st.session_state[dataset_key] # Get image files @@ -196,3 +140,77 @@ def render_image_detection_viewer(): st.error(f"Error displaying image: {e}") else: st.info("Select an image to view with annotations.") + + +def get_image_detection_dataset_key(dataset_type, dataset_path, split): + """Build the session state key used to cache an image detection dataset. + + The dataset type is part of the key so that switching between COCO and YOLO + for the same folder and split never returns a stale dataset. + + :param dataset_type: Type of the dataset (e.g., "COCO", "YOLO") + :type dataset_type: str + :param dataset_path: Path to the dataset directory + :type dataset_path: str + :param split: Dataset split (e.g., "train", "val", "test") + :type split: str + :return: Session state key for the dataset + :rtype: str + """ + return f"{dataset_type.lower()}_{dataset_path}_{split}" + + +def load_image_detection_dataset( + dataset_type, dataset_path, split, dataset_config_file=None +): + """Load an image detection dataset based on the provided type, path, and split. + + Filenames are kept relative to the dataset directory. Models make them global + on their own copy of the dataset during evaluation. + + :param dataset_type: Type of the dataset ("COCO" or "YOLO", case-insensitive) + :type dataset_type: str + :param dataset_path: Path to the dataset directory + :type dataset_path: str + :param split: Dataset split to load (e.g., "train", "val", "test") + :type split: str + :param dataset_config_file: YOLO dataset configuration (YAML) as an uploaded + file-like object exposing ``getvalue()``. Required for YOLO datasets. + :type dataset_config_file: Optional[streamlit.runtime.uploaded_file_manager.UploadedFile] + :raises FileNotFoundError: If required dataset files or configuration are missing + :raises ValueError: If the dataset type is not supported + :return: Loaded dataset containing only the requested split + :rtype: CocoDataset or YOLODataset + """ + from perceptionmetrics.datasets.coco import CocoDataset + from perceptionmetrics.datasets.yolo import YOLODataset + + dataset_type = dataset_type.lower() + + if dataset_type == "coco": + img_dir, ann_file = find_img_dir_and_ann_file( + dataset_path=dataset_path, split=split + ) + return CocoDataset(annotation_file=ann_file, image_dir=img_dir, split=split) + + if dataset_type == "yolo": + if dataset_config_file is None: + raise FileNotFoundError( + "Please upload the YOLO dataset configuration (.yaml) in the sidebar." + ) + + # getvalue() returns the full content even if the file was read before + with tempfile.NamedTemporaryFile(delete=False, suffix=".yaml") as tmp: + tmp.write(dataset_config_file.getvalue()) + tmp_path = tmp.name + try: + yolo_dataset = YOLODataset(tmp_path, dataset_path) + finally: + os.unlink(tmp_path) + + yolo_dataset.dataset = yolo_dataset.dataset[ + yolo_dataset.dataset["split"] == split + ].reset_index(drop=True) + return yolo_dataset + + raise ValueError(f"{dataset_type} image detection dataset is not supported.") diff --git a/gui/tasks/image_detection/evaluator.py b/gui/tasks/image_detection/evaluator.py index 0d65ff37..061eff32 100644 --- a/gui/tasks/image_detection/evaluator.py +++ b/gui/tasks/image_detection/evaluator.py @@ -2,11 +2,12 @@ import os import tempfile import json -from perceptionmetrics.datasets.coco import CocoDataset - from gui.tasks.utils import browse_folder -from perceptionmetrics.datasets.coco import find_img_dir_and_ann_file +from gui.tasks.image_detection.dataset_viewer import ( + get_image_detection_dataset_key, + load_image_detection_dataset, +) def browse_predictions_outdir(): @@ -28,43 +29,27 @@ def render_image_detection_evaluator(): # Check for dataset from sidebar inputs dataset_path = st.session_state.get("dataset_path", "") - dataset_type = st.session_state.get("dataset_type", "Coco") + dataset_type = st.session_state.get("dataset_type", "COCO") split = st.session_state.get("split", "val") - # Try to get existing dataset from session state first - dataset_key = f"{dataset_path}_{split}" - if dataset_key in st.session_state: - dataset = st.session_state[dataset_key] - dataset_available = True - st.success( - f"✅ Dataset loaded: {dataset_path} ({split} split) - {len(dataset.dataset)} samples" - ) - elif dataset_path and os.path.isdir(dataset_path): + # Reuse the dataset loaded by the viewer if available, otherwise load it here + dataset_key = get_image_detection_dataset_key(dataset_type, dataset_path, split) + if dataset_path and os.path.isdir(dataset_path): try: - if dataset_type.lower() == "coco": - img_dir, ann_file = find_img_dir_and_ann_file( - dataset_path=dataset_path, split=split - ) - - if os.path.isdir(img_dir) and os.path.isfile(ann_file): - st.session_state[dataset_key] = CocoDataset( - annotation_file=ann_file, image_dir=img_dir, split=split - ) - # Make filenames global - this is crucial for evaluation - st.session_state[dataset_key].make_fname_global() - dataset = st.session_state[dataset_key] - dataset_available = True - st.success( - f"✅ Dataset loaded: {dataset_path} ({split} split) - {len(dataset.dataset)} samples" - ) - else: - st.warning( - "⚠️ Dataset files not found. Please check the dataset path and split in the sidebar." - ) - else: - st.warning( - "⚠️ Only COCO datasets are currently supported for evaluation." + if dataset_key not in st.session_state: + st.session_state[dataset_key] = load_image_detection_dataset( + dataset_type, + dataset_path, + split, + st.session_state.get("dataset_config_file"), ) + dataset = st.session_state[dataset_key] + dataset_available = True + st.success( + f"✅ Dataset loaded: {dataset_path} ({split} split) - {len(dataset.dataset)} samples" + ) + except FileNotFoundError as e: + st.warning(f"⚠️ {e}") except Exception as e: st.error(f"❌ Error loading dataset: {e}") else: diff --git a/tests/gui/test_image_detection.py b/tests/gui/test_image_detection.py index 32953ff5..e79989a5 100644 --- a/tests/gui/test_image_detection.py +++ b/tests/gui/test_image_detection.py @@ -1,8 +1,15 @@ +import io import json +import pytest from PIL import Image from streamlit.testing.v1 import AppTest +from gui.tasks.image_detection.dataset_viewer import ( + get_image_detection_dataset_key, + load_image_detection_dataset, +) + def create_tiny_coco_dataset(tmp_path): """Create a minimal COCO-style validation split for GUI tests.""" @@ -47,6 +54,24 @@ def create_tiny_coco_dataset(tmp_path): return tmp_path +def create_tiny_yolo_dataset(tmp_path): + """Create a minimal YOLO dataset with one train and one val sample. + + :return: Dataset directory and the YAML configuration as bytes + :rtype: tuple + """ + for split, name in [("train", "train_sample"), ("val", "val_sample")]: + image_dir = tmp_path / "images" / split + label_dir = tmp_path / "labels" / split + image_dir.mkdir(parents=True) + label_dir.mkdir(parents=True) + Image.new("RGB", (32, 32), color=(40, 80, 120)).save(image_dir / f"{name}.jpg") + (label_dir / f"{name}.txt").write_text("0 0.5 0.5 0.25 0.25\n") + + config = "train: images/train\nval: images/val\nnames:\n 0: car\n" + return tmp_path, config.encode("utf-8") + + def get_by_label(elements, label): """Return the first Streamlit testing element with the given label.""" return next(element for element in elements if element.label == label) @@ -173,7 +198,7 @@ def test_image_detection_gui_loads_tiny_coco_dataset(tmp_path): assert app.session_state["split"] == "val" assert app.session_state["dataset_path"] == str(dataset_path) - dataset_key = f"{dataset_path}_val" + dataset_key = get_image_detection_dataset_key("COCO", str(dataset_path), "val") assert dataset_key in app.session_state assert len(app.session_state[dataset_key].dataset) == 1 @@ -207,3 +232,58 @@ def test_image_detection_evaluator_uses_loaded_coco_dataset(tmp_path): button for button in app.button if button.label == "🚀 Run Evaluation" ) assert run_evaluation.disabled + + +def test_image_detection_evaluator_does_not_reject_yolo(tmp_path): + """Verify that selecting YOLO no longer shows a COCO-only warning (#529).""" + dataset_path, _ = create_tiny_yolo_dataset(tmp_path) + app = AppTest.from_file("app.py").run(timeout=10) + + get_by_label(app.selectbox, "Type").select("YOLO").run(timeout=10) + get_by_label(app.selectbox, "Split").select("val").run(timeout=10) + get_by_label(app.text_input, "Dataset Folder").input(str(dataset_path)).run( + timeout=10 + ) + + warning_text = "\n".join(warning.value for warning in app.warning) + + assert not app.exception + assert "Only COCO datasets are currently supported" not in warning_text + assert "Please upload the YOLO dataset configuration" in warning_text + + +def test_load_image_detection_dataset_yolo(tmp_path): + """Verify that the shared loader builds a YOLO dataset filtered by split.""" + dataset_path, config = create_tiny_yolo_dataset(tmp_path) + + dataset = load_image_detection_dataset( + "YOLO", str(dataset_path), "val", io.BytesIO(config) + ) + + assert len(dataset.dataset) == 1 + assert set(dataset.dataset["split"]) == {"val"} + assert "car" in dataset.ontology + + +def test_load_image_detection_dataset_yolo_requires_config(tmp_path): + """Verify that YOLO loading fails with a clear message without a config.""" + dataset_path, _ = create_tiny_yolo_dataset(tmp_path) + + with pytest.raises(FileNotFoundError, match="YOLO dataset configuration"): + load_image_detection_dataset("YOLO", str(dataset_path), "val") + + +def test_load_image_detection_dataset_coco(tmp_path): + """Verify that the shared loader builds a COCO dataset.""" + dataset_path = create_tiny_coco_dataset(tmp_path) + + dataset = load_image_detection_dataset("COCO", str(dataset_path), "val") + + assert len(dataset.dataset) == 1 + + +def test_image_detection_dataset_key_depends_on_type(): + """Verify that COCO and YOLO datasets on the same path are cached separately.""" + assert get_image_detection_dataset_key( + "COCO", "/data", "val" + ) != get_image_detection_dataset_key("YOLO", "/data", "val")