Fix #529: support YOLO datasets in the image detection evaluator

gui/tasks/image_detection/dataset_viewer.py

@@ -1,15 +1,13 @@
 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
 
 
 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.")

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:

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")