diff --git a/app.py b/app.py index 564396ef..b05d0ef3 100644 --- a/app.py +++ b/app.py @@ -1,15 +1,12 @@ import streamlit as st + +from perceptionmetrics.utils.torch import get_device_info from tabs.dataset_viewer import dataset_viewer_tab -from tabs.inference import inference_tab from tabs.evaluator import evaluator_tab -from perceptionmetrics.utils.gui import browse_folder -from perceptionmetrics.utils.torch import get_device_info - - -def browse_dataset_path(): - st.session_state.dataset_path = browse_folder() - - +from tabs.inference import inference_tab +from tabs.tasks.image_detection.sidebar import render_image_detection_sidebar +from tabs.tasks.image_segmentation.sidebar import render_image_segmentation_sidebar +from tabs.tasks.lidar_segmentation.sidebar import render_lidar_segmentation_sidebar st.set_page_config(page_title="PerceptionMetrics", layout="wide") PAGES = { @@ -20,408 +17,51 @@ def browse_dataset_path(): best_device, available_devices = get_device_info() -# Initialize commonly used session state keys +# Shared state +st.session_state.setdefault("task", "Image Detection") st.session_state.setdefault("dataset_path", "") -st.session_state.setdefault("dataset_type", "YOLO") st.session_state.setdefault("split", "test") +st.session_state.setdefault("device", best_device) + +# Image detection state +st.session_state.setdefault("dataset_type", "YOLO") st.session_state.setdefault("config_option", "Manual Configuration") st.session_state.setdefault("confidence_threshold", 0.5) st.session_state.setdefault("nms_threshold", 0.5) -st.session_state.setdefault("max_detections", -1) -st.session_state.setdefault("device", best_device) +st.session_state.setdefault("max_detections", 100) st.session_state.setdefault("batch_size", 1) st.session_state.setdefault("evaluation_step", 5) st.session_state.setdefault("detection_model", None) st.session_state.setdefault("detection_model_loaded", False) -# Sidebar: Dataset Inputs -with st.sidebar: - with st.expander("Dataset Inputs", expanded=True): - # First row: Type and Split - col1, col2 = st.columns(2) - with col1: - st.selectbox( - "Type", - ["COCO", "YOLO"], - key="dataset_type", - ) - with col2: - st.selectbox( - "Split", - ["train", "val", "test"], - key="split", - ) - - # Second row: Path and Browse button - col1, col2 = st.columns([3.5, 2.5]) - with col1: - st.text_input("Dataset Folder", key="dataset_path") - with col2: - st.markdown( - "
", unsafe_allow_html=True - ) - st.button("Browse", on_click=browse_dataset_path, use_container_width=True) - - # Additional input for YOLO config file - if st.session_state.get("dataset_type", "COCO") == "YOLO": - st.file_uploader( - "Dataset Configuration (.yaml)", - type=["yaml"], - key="dataset_config_file", - help="Upload a YAML dataset configuration file.", - ) - - with st.expander("Model Inputs", expanded=False): - st.file_uploader( - "Model File (.pt, .onnx, .h5, .pb, .pth, .torchscript)", - type=["pt", "onnx", "h5", "pb", "pth", "torchscript"], - key="model_file", - help="Upload your trained model file.", - max_upload_size=1024, # MB - ) - st.file_uploader( - "Ontology File (.json)", - type=["json"], - key="ontology_file", - help="Upload a JSON file with class labels.", - ) - st.radio( - "Configuration Method:", - ["Manual Configuration", "Upload Config File"], - key="config_option", - horizontal=True, - ) - if ( - st.session_state.get("config_option", "Manual Configuration") - == "Upload Config File" - ): - st.file_uploader( - "Configuration File (.json)", - type=["json"], - key="config_file", - help="Upload a JSON configuration file.", - ) - else: - col1, col2 = st.columns(2) - with col1: - st.slider( - "Confidence Threshold", - min_value=0.0, - max_value=1.0, - step=0.01, - key="confidence_threshold", - help="Minimum confidence score for detections", - ) - st.slider( - "NMS Threshold", - min_value=0.0, - max_value=1.0, - step=0.01, - key="nms_threshold", - help="Non-maximum suppression threshold", - ) - st.number_input( - "Max Detections/Image", - min_value=-1, - max_value=1000, - step=1, - key="max_detections", - ) - with col2: - st.selectbox( - "Device", - available_devices, - key="device", - ) - st.selectbox( - "Model Format", - ["torchvision", "YOLO"], - index=( - 0 - if st.session_state.get("model_format", "torchvision") - == "torchvision" - else 1 - ), - key="model_format", - ) - st.number_input( - "Batch Size", - min_value=1, - max_value=256, - step=1, - key="batch_size", - ) - st.number_input( - "Evaluation Step", - min_value=0, - max_value=1000, - step=1, - key="evaluation_step", - help="Update UI with intermediate metrics every N images (0 = disable intermediate updates)", - ) - - st.write("---") - st.write("**Image Size Configuration**") +# Image segmentation state +st.session_state.setdefault("segmentation_model", None) +st.session_state.setdefault("segmentation_model_loaded", False) +st.session_state.setdefault("segmentation_model_type", "Torch Model File") +st.session_state.setdefault("segmentation_model_path", "") +st.session_state.setdefault("segmentation_config_path", "") +st.session_state.setdefault("segmentation_ontology_path", "") - # Resize Logic - enable_resize = st.checkbox( - "Enable Resize", value=True, key="enable_resize" - ) - - if enable_resize: - resize_strategy = st.radio( - "Resize Strategy", - ["Min Side", "Max Side", "Fixed Dimensions"], - key="resize_strategy", - horizontal=True, - label_visibility="collapsed", - ) - - if resize_strategy == "Fixed Dimensions": - c1, c2 = st.columns(2) - with c1: - st.number_input( - "Image Resize Height", - min_value=1, - max_value=4096, - value=640, - step=1, - key="resize_height", - help="Height to resize images for inference", - ) - with c2: - st.number_input( - "Image Resize Width", - min_value=1, - max_value=4096, - value=640, - step=1, - key="resize_width", - help="Width to resize images for inference", - ) - elif resize_strategy == "Min Side": - st.number_input( - "Min Side", - min_value=1, - max_value=4096, - value=640, - step=1, - key="min_side", - help="Minimum size of the shorter side of the image", - ) - elif resize_strategy == "Max Side": - st.number_input( - "Max Side", - min_value=1, - max_value=4096, - value=640, - step=1, - key="max_side", - help="Maximum size of the longer side of the image", - ) - else: - st.error("Invalid resize strategy selected") - - # Pad to closest multiple - enable_pad = st.checkbox( - "Enable Padding to Closest Multiple", value=True, key="enable_pad" - ) - - if enable_pad: - st.number_input( - "Divisor", - min_value=1, - max_value=128, - value=32, - step=1, - key="pad_divisor", - help="Pad image dimensions to the closest multiple of this value", - ) - - # Crop Logic - enable_crop = st.checkbox("Enable Center Crop", key="enable_crop") - - if enable_crop: - c1, c2 = st.columns(2) - with c1: - st.number_input( - "Crop Height", - min_value=1, - max_value=4096, - value=640, - step=1, - key="crop_height", - help="Center crop height", - ) - with c2: - st.number_input( - "Crop Width", - min_value=1, - max_value=4096, - value=640, - step=1, - key="crop_width", - help="Center crop width", - ) - - # Load model action in sidebar - from perceptionmetrics.models.torch_detection import TorchImageDetectionModel - import json, tempfile - - load_model_btn = st.button( - "Load Model", - type="primary", - width="stretch", - help="Load and save the model for use in the Inference tab", - key="sidebar_load_model_btn", - ) - - if load_model_btn: - model_file = st.session_state.get("model_file") - ontology_file = st.session_state.get("ontology_file") - config_option = st.session_state.get( - "config_option", "Manual Configuration" - ) - config_file = ( - st.session_state.get("config_file") - if config_option == "Upload Config File" - else None - ) - - # Prepare configuration - config_data = None - config_path = None - try: - if config_option == "Upload Config File": - if config_file is not None: - config_data = json.load(config_file) - with tempfile.NamedTemporaryFile( - delete=False, suffix=".json", mode="w" - ) as tmp_cfg: - json.dump(config_data, tmp_cfg) - config_path = tmp_cfg.name - else: - st.error("Please upload a configuration file") - else: - confidence_threshold = float( - st.session_state.get("confidence_threshold", 0.5) - ) - nms_threshold = float(st.session_state.get("nms_threshold", 0.5)) - max_detections = int(st.session_state.get("max_detections", -1)) - device = st.session_state.get("device", "cpu") - batch_size = int(st.session_state.get("batch_size", 1)) - evaluation_step = int(st.session_state.get("evaluation_step", 5)) - model_format = st.session_state.get("model_format", "torchvision") - - # Resize Logic extraction - enable_resize = st.session_state.get("enable_resize", True) - resize_cfg = None - if enable_resize: - resize_strategy = st.session_state.get( - "resize_strategy", "Fixed Dimensions" - ) - if resize_strategy == "Fixed Dimensions": - resize_height = int( - st.session_state.get("resize_height", 640) - ) - resize_width = int( - st.session_state.get("resize_width", 640) - ) - resize_cfg = { - "height": resize_height, - "width": resize_width, - } - elif resize_strategy == "Min Side": - min_side = int(st.session_state.get("min_side", 640)) - resize_cfg = {"min_side": min_side} - elif resize_strategy == "Max Side": - max_side = int(st.session_state.get("max_side", 640)) - resize_cfg = {"max_side": max_side} - else: - st.error("Invalid resize strategy selected") - - if enable_pad: - pad_divisor = int(st.session_state.get("pad_divisor", 32)) - if resize_cfg is not None: - resize_cfg["closest_divisor"] = pad_divisor - else: - resize_cfg = {"closest_divisor": pad_divisor} - - config_data = { - "confidence_threshold": confidence_threshold, - "nms_threshold": nms_threshold, - "max_detections_per_image": max_detections, - "device": device, - "batch_size": batch_size, - "evaluation_step": evaluation_step, - "model_format": model_format.lower(), - } - if resize_cfg is not None: - config_data["resize"] = resize_cfg - - if enable_crop: - crop_height = int(st.session_state.get("crop_height", 640)) - crop_width = int(st.session_state.get("crop_width", 640)) - crop_cfg = {"height": crop_height, "width": crop_width} - config_data["crop"] = crop_cfg - - with tempfile.NamedTemporaryFile( - delete=False, suffix=".json", mode="w" - ) as tmp_cfg: - json.dump(config_data, tmp_cfg) - config_path = tmp_cfg.name - except Exception as e: - st.error(f"Failed to prepare configuration: {e}") - config_path = None - - if model_file is None: - st.error("Please upload a model file") - elif config_path is None: - st.error("Please provide a valid model configuration") - elif ontology_file is None: - st.error("Please upload an ontology file") - else: - with st.spinner("Loading model..."): - # Persist ontology to temp file - try: - ontology_data = json.load(ontology_file) - with tempfile.NamedTemporaryFile( - delete=False, suffix=".json", mode="w" - ) as tmp_ont: - json.dump(ontology_data, tmp_ont) - ontology_path = tmp_ont.name - except Exception as e: - st.error(f"Failed to load ontology: {e}") - ontology_path = None - - # Persist model to temp file - try: - with tempfile.NamedTemporaryFile( - delete=False, suffix=".pt", mode="wb" - ) as tmp_model: - tmp_model.write(model_file.read()) - model_temp_path = tmp_model.name - except Exception as e: - st.error(f"Failed to save model file: {e}") - model_temp_path = None +with st.sidebar: + task = st.selectbox( + "Task", + ["Image Detection", "Image Segmentation", "Lidar Segmentation"], + key="task", + help="Image segmentation is currently a placeholder.", + ) + + if task == "Image Detection": + render_image_detection_sidebar(available_devices) + elif task == "Image Segmentation": + render_image_segmentation_sidebar(available_devices) + elif task == "Lidar Segmentation": + render_lidar_segmentation_sidebar(available_devices) + else: + st.error(f"Unsupported task: {task}") + + - if ontology_path and model_temp_path: - try: - model = TorchImageDetectionModel( - model=model_temp_path, - model_cfg=config_path, - ontology_fname=ontology_path, - device=st.session_state.get("device", "cpu"), - ) - st.session_state.detection_model = model - st.session_state.detection_model_loaded = True - st.success("Model loaded and saved for inference") - except Exception as e: - st.session_state.detection_model = None - st.session_state.detection_model_loaded = False - st.error(f"Failed to load model: {e}") -# Main content area with horizontal tabs tab1, tab2, tab3 = st.tabs(["Dataset Viewer", "Inference", "Evaluator"]) with tab1: diff --git a/perceptionmetrics/models/torch_detection.py b/perceptionmetrics/models/torch_detection.py index dbdfe12c..9b45fcc8 100644 --- a/perceptionmetrics/models/torch_detection.py +++ b/perceptionmetrics/models/torch_detection.py @@ -9,7 +9,10 @@ import torch from torch.utils.data import DataLoader, Dataset from torchvision.transforms import v2 as transforms -from torchvision import tv_tensors +try: + from torchvision import tv_tensors +except ImportError: + from torchvision import datapoints as tv_tensors from tqdm.auto import tqdm from perceptionmetrics.datasets import detection as detection_dataset @@ -191,9 +194,14 @@ def __getitem__( # Convert boxes/labels to tensors if len(boxes) == 0: boxes = torch.zeros((0, 4), dtype=torch.float32) - boxes = tv_tensors.BoundingBoxes( - boxes, format="XYXY", canvas_size=(image.height, image.width) - ) + if hasattr(tv_tensors, "BoundingBoxes"): + boxes = tv_tensors.BoundingBoxes( + boxes, format="XYXY", canvas_size=(image.height, image.width) + ) + else: + boxes = tv_tensors.BoundingBox( + boxes, format="XYXY", spatial_size=(image.height, image.width) + ) category_indices = torch.as_tensor(category_indices, dtype=torch.int64) target = { diff --git a/perceptionmetrics/models/torch_segmentation.py b/perceptionmetrics/models/torch_segmentation.py index ddf03d25..66851090 100644 --- a/perceptionmetrics/models/torch_segmentation.py +++ b/perceptionmetrics/models/torch_segmentation.py @@ -2,7 +2,7 @@ import os import time import tempfile -from typing import Any, List, Optional, Tuple, Union +from typing import Any, Callable, List, Optional, Tuple, Union import numpy as np import pandas as pd @@ -374,6 +374,8 @@ def eval( translation_direction: str = "dataset_to_model", predictions_outdir: Optional[str] = None, results_per_sample: bool = False, + progress_callback: Optional[Callable] = None, + metrics_callback: Optional[Callable] = None, ) -> pd.DataFrame: """Perform evaluation for an image segmentation dataset @@ -389,6 +391,10 @@ def eval( :type predictions_outdir: Optional[str], optional :param results_per_sample: Whether to store results per sample or not, defaults to False. If True, predictions_outdir must be provided. :type results_per_sample: bool, optional + :param progress_callback: Optional callback called as progress_callback(processed, total), defaults to None + :type progress_callback: Optional[Callable], optional + :param metrics_callback: Optional callback called as metrics_callback(metrics_df, processed, total), defaults to None + :type metrics_callback: Optional[Callable], optional :return: DataFrame containing evaluation results :rtype: pd.DataFrame """ @@ -444,6 +450,9 @@ def eval( # Init metrics metrics_factory = um.SegmentationMetricsFactory(n_classes) + total_samples = len(dataset) + processed_samples = 0 + evaluation_step = self.model_cfg.get("evaluation_step", 1) # Evaluation loop with torch.no_grad(): @@ -504,6 +513,26 @@ def eval( os.path.join(predictions_outdir, f"{sample_idx}.png") ) + processed_samples += len(idx) + + if progress_callback is not None: + progress_callback(processed_samples, total_samples) + + if ( + metrics_callback is not None + and evaluation_step is not None + and evaluation_step > 0 + and ( + processed_samples % evaluation_step == 0 + or processed_samples == total_samples + ) + ): + metrics_callback( + um.get_metrics_dataframe(metrics_factory, eval_ontology), + processed_samples, + total_samples, + ) + return um.get_metrics_dataframe(metrics_factory, eval_ontology) def get_computational_cost( diff --git a/perceptionmetrics/utils/gui.py b/perceptionmetrics/utils/gui.py deleted file mode 100644 index b1295338..00000000 --- a/perceptionmetrics/utils/gui.py +++ /dev/null @@ -1,83 +0,0 @@ -import sys -import subprocess -import platform - - -def is_wsl(): - """ - Detect if running in Windows Subsystem for Linux (WSL). - Returns True if WSL is detected, False otherwise. - """ - return ( - "wsl" in platform.release().lower() or "microsoft" in platform.release().lower() - ) - - -def browse_folder(): - """ - Opens a native folder selection dialog and returns the selected folder path. - Works on Windows, macOS, and Linux (with zenity or kdialog). - Returns None if cancelled or error. - """ - try: - is_windows = sys.platform.startswith("win") - is_wsl_env = is_wsl() - if is_windows or is_wsl_env: - script = ( - "Add-Type -AssemblyName System.windows.forms;" - "$f=New-Object System.Windows.Forms.FolderBrowserDialog;" - 'if($f.ShowDialog() -eq "OK"){Write-Output $f.SelectedPath}' - ) - result = subprocess.run( - ["powershell.exe", "-NoProfile", "-Command", script], - capture_output=True, - text=True, - timeout=30, - ) - folder = result.stdout.strip() - if folder and is_wsl_env: # Convert Windows path to WSL path - result = subprocess.run( - ["wslpath", "-u", folder], - capture_output=True, - text=True, - timeout=30, - ) - folder = result.stdout.strip() - return folder if folder else None - elif sys.platform == "darwin": - script = 'POSIX path of (choose folder with prompt "Select folder:")' - result = subprocess.run( - ["osascript", "-e", script], capture_output=True, text=True, timeout=30 - ) - folder = result.stdout.strip() - return folder if folder else None - else: - # Linux: try zenity, then kdialog - for cmd in [ - [ - "zenity", - "--file-selection", - "--directory", - "--title=Select folder", - ], - [ - "kdialog", - "--getexistingdirectory", - "--title", - "Select folder", - ], - ]: - try: - result = subprocess.run( - cmd, capture_output=True, text=True, timeout=30 - ) - if result.returncode == 0 or result.returncode == 1: # zenity and kdialog return 1 on cancel - folder = result.stdout.strip() - return folder if folder else None - except subprocess.TimeoutExpired: - return None - except (FileNotFoundError, Exception): - continue - return None - except Exception: - return None \ No newline at end of file diff --git a/tabs/dataset_viewer.py b/tabs/dataset_viewer.py index 85077129..8ded02a9 100644 --- a/tabs/dataset_viewer.py +++ b/tabs/dataset_viewer.py @@ -1,316 +1,25 @@ import streamlit as st -import os -from streamlit_image_select import image_select - -from perceptionmetrics.datasets.coco import find_img_dir_and_ann_file +from tabs.tasks.image_detection.dataset_viewer import render_image_detection_viewer +from tabs.tasks.image_segmentation.dataset_viewer import ( + render_image_segmentation_viewer, +) +from tabs.tasks.lidar_segmentation.dataset_viewer import ( + render_lidar_segmentation_viewer, +) def dataset_viewer_tab(): - 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 - from supervision.detection.annotate import BoxAnnotator - from supervision.detection.core import Detections - - # Get inputs from session state - dataset_path = st.session_state.get("dataset_path", "") - dataset_type = st.session_state.get("dataset_type", "COCO").lower() - split = st.session_state.get("split", "val") + task = st.session_state.get("task", "Image Detection") - # Header row only - st.header("Dataset Viewer") - - if not dataset_path or not os.path.isdir(dataset_path): - st.warning("⚠️ Please select a valid dataset folder.") + if task == "Image Detection": + render_image_detection_viewer() return - # Setup paths and pagination - if dataset_type == "coco": - try: - img_dir, ann_file = 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}") - 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.") + if task == "Image Segmentation": + render_image_segmentation_viewer() return - - # Pagination and search row - nav_col1, nav_col2, nav_col3, nav_col4 = st.columns([1, 1, 2, 1.5]) - with nav_col1: - pass # Placeholder for "< page" button, to be added later in the code - with nav_col2: - pass # Placeholder for ">" button, to be added later in the code - with nav_col3: - pass # Placeholder for page info, to be added later in the code - with nav_col4: - # Move the button up by reducing the margin and decrease button size with custom CSS - st.markdown( - """ - - """, - unsafe_allow_html=True, - ) - st.markdown("
", unsafe_allow_html=True) - - # Load dataset - dataset_key = f"{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 - - 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 - image_files = [ - f for f in os.listdir(img_dir) if f.lower().endswith((".jpg", ".jpeg", ".png")) - ] - if not image_files: - st.warning("No images found.") + if task == "Lidar Segmentation": + render_lidar_segmentation_viewer() return - # Pagination - IMAGES_PER_PAGE = 12 - _, total_pages = ( - len(image_files), - (len(image_files) + IMAGES_PER_PAGE - 1) // IMAGES_PER_PAGE, - ) - page_key = f"image_page_{dataset_path}_{split}" - - if page_key not in st.session_state: - st.session_state[page_key] = 0 - current_page = max(0, min(st.session_state[page_key], total_pages - 1)) - st.session_state[page_key] = current_page - - start_idx = current_page * IMAGES_PER_PAGE - sample_images = image_files[start_idx : start_idx + IMAGES_PER_PAGE] - image_paths = [os.path.join(img_dir, img_name) for img_name in sample_images] - - # Navigation - col1, col2, col3, col4 = st.columns([0.5, 9.5, 0.5, 0.5]) - with col1: - if st.button("⟨", key="prev_page_btn", disabled=(current_page == 0)): - st.session_state[page_key] = current_page - 1 - st.rerun() - with col2: - st.markdown( - f"
Page {current_page + 1} of {total_pages}
", - unsafe_allow_html=True, - ) - with col3: - if st.button( - "⟩", key="next_page_btn", disabled=(current_page >= total_pages - 1) - ): - st.session_state[page_key] = current_page + 1 - st.rerun() - with col4: - if st.button( - "🔍", - key="search_icon_btn", - help="Search for an image by name", - disabled=not (dataset_path and os.path.isdir(dataset_path)), - ): - st.session_state["show_search_dropdown"] = True - - # Search dropdown - if st.session_state.get("show_search_dropdown", False): - col1, col2, col3 = st.columns([4, 1, 1]) - with col1: - selected_img = st.selectbox( - "Search image:", options=image_files, key="search_image" - ) - with col2: - st.markdown( - "
", unsafe_allow_html=True - ) - if st.button("Go to image", key="go_to_image_btn"): - new_page = image_files.index(selected_img) // IMAGES_PER_PAGE - st.session_state[page_key] = new_page - st.session_state[ - f"img_select_all_{dataset_path}_{split}_{new_page}" - ] = (image_files.index(selected_img) % IMAGES_PER_PAGE) - st.session_state["show_search_dropdown"] = False - st.rerun() - with col3: - st.markdown( - "
", unsafe_allow_html=True - ) - if st.button("Cancel", key="cancel_search_btn"): - st.session_state["show_search_dropdown"] = False - st.rerun() - - caption_len_limit = 17 - captions = [ - ( - (name[:caption_len_limit] + "..." + name[-3:]) - if len(name) > caption_len_limit - else name - ) - for name in sample_images - ] - - # Image grid - img_select_key = f"img_select_all_{dataset_path}_{split}_{current_page}" - img_select_index = st.session_state.get(img_select_key) - if img_select_index is None or not isinstance(img_select_index, int): - img_select_index = 0 - selected_img_path = ( - image_select( - label="", - images=image_paths, - captions=captions, - use_container_width=False, - key=img_select_key, - index=img_select_index, - ) - if image_paths - else None - ) - - # Display selected image with annotations - if selected_img_path: - selected_img_name = os.path.basename(selected_img_path) - try: - img = Image.open(selected_img_path).convert("RGB") - img_np = np.array(img) - - if dataset_type == "yolo": - ann_row = dataset.dataset[ - dataset.dataset["image"].str.endswith(selected_img_name) - ] - else: - ann_row = dataset.dataset[dataset.dataset["image"] == selected_img_name] - - if not ann_row.empty: - annotation_id = ann_row.iloc[0]["annotation"] - if dataset_type == "yolo": - annotation_id = os.path.join(dataset.dataset_dir, annotation_id) - - boxes, category_indices = dataset.read_annotation(annotation_id) - - # Get class names from ontology - ontology = getattr(dataset, "ontology", None) - if ontology is None and hasattr(dataset.dataset, "attrs"): - ontology = dataset.dataset.attrs.get("ontology", None) - - if ontology: - catid_to_name = {v["idx"]: k for k, v in ontology.items()} - class_names = [ - catid_to_name.get(cat_id, str(cat_id)) - for cat_id in category_indices - ] - else: - class_names = [str(cat_id) for cat_id in category_indices] - - # Annotate image - palette = ColorPalette.default() - detections = Detections( - xyxy=np.array(boxes), class_id=np.array(category_indices) - ) - annotator = BoxAnnotator( - color=palette, text_scale=0.7, text_thickness=1, text_padding=2 - ) - annotated_img = annotator.annotate( - scene=img_np, detections=detections, labels=class_names - ) - - # Resize for display - annotated_pil = Image.fromarray(annotated_img) - try: - resample = getattr(Image, "Resampling", Image).LANCZOS - except AttributeError: - resample = Image.LANCZOS - annotated_pil.thumbnail((500, 500), resample) - st.image(annotated_pil, width="content") - else: - st.warning("No annotation found for this image.") - except Exception as e: - st.error(f"Error displaying image: {e}") - else: - st.info("Select an image to view with annotations.") + st.error(f"Unsupported task for dataset viewer: {task}") diff --git a/tabs/evaluator.py b/tabs/evaluator.py index 2e5fcac8..6d49e31f 100644 --- a/tabs/evaluator.py +++ b/tabs/evaluator.py @@ -1,505 +1,22 @@ import streamlit as st -import os -import tempfile -import json -from perceptionmetrics.datasets.coco import CocoDataset - - -from perceptionmetrics.utils.gui import browse_folder -from perceptionmetrics.datasets.coco import find_img_dir_and_ann_file - - -def browse_predictions_outdir(): - folder = browse_folder() - if folder: - st.session_state.predictions_outdir = folder +from tabs.tasks.image_detection.evaluator import render_image_detection_evaluator +from tabs.tasks.image_segmentation.evaluator import render_image_segmentation_evaluator +from tabs.tasks.lidar_segmentation.evaluator import render_lidar_segmentation_evaluator def evaluator_tab(): - st.header("Evaluator") - st.markdown("Evaluate your model on the loaded dataset using PerceptionMetrics.") - - # Check if we have the required objects from sidebar inputs - dataset_available = False - model_available = False - dataset = None - model = None - - # Check for dataset from sidebar inputs - dataset_path = st.session_state.get("dataset_path", "") - 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): - 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." - ) - except Exception as e: - st.error(f"❌ Error loading dataset: {e}") - else: - st.warning( - "⚠️ No dataset path provided. Please set the dataset path in the sidebar." - ) - - # Check for model from sidebar (loaded via Load Model button) - if ( - "detection_model" in st.session_state - and st.session_state.detection_model is not None - ): - model = st.session_state.detection_model - model_available = True - st.success("✅ Model loaded and ready for evaluation") - else: - st.warning( - "⚠️ No model loaded. Please load a model using the 'Load Model' button in the sidebar." - ) - - # Evaluation configuration - st.markdown("### Evaluation Configuration") - - save_predictions = st.checkbox( - "Save Predictions", - value=False, - help="Save individual predictions and metrics per sample", - ) - - save_visualizations = st.checkbox( - "Save Visualizations", - value=False, - help="Save visualized qualitative results (Image + GT + Preds)", - ) - - predictions_outdir_input = None - if save_predictions or save_visualizations: - col1, col2 = st.columns([3, 1]) - with col1: - st.text_input("Predictions Output Directory", key="predictions_outdir") - with col2: - st.markdown( - "
", unsafe_allow_html=True - ) - st.button( - "Browse", on_click=browse_predictions_outdir, key="browse_preds_outdir" - ) - - predictions_outdir_input = st.session_state.get("predictions_outdir") - - ontology_translation = st.file_uploader( - "Ontology Translation (Optional)", - type=["json"], - help="JSON file for translating between dataset and model ontologies", - ) - - # Disable Run Evaluation button if output dir is missing - output_dir_required = save_predictions or save_visualizations - output_dir_missing = output_dir_required and not ( - predictions_outdir_input and predictions_outdir_input.strip() - ) - - if output_dir_missing: - st.warning( - "⚠️ Please provide a Predictions Output Directory to enable evaluation " - "when 'Save Predictions' or 'Save Visualizations' is turned on." - ) - - # Run evaluation button - if st.button( - "🚀 Run Evaluation", - type="primary", - disabled=not (dataset_available and model_available) or output_dir_missing, - ): - if not dataset_available or not model_available: - st.error( - "Please ensure both dataset and model are loaded before running evaluation." - ) - return - - # Prepare evaluation - with st.spinner("Running evaluation..."): - try: - # Validate dataset and model - if len(dataset.dataset) == 0: - st.error( - "Dataset has no samples. Please check the dataset configuration." - ) - return - - if not hasattr(model, "model_cfg") or model.model_cfg is None: - st.error( - "Model configuration is missing. Please reload the model in the Inference tab." - ) - return - - # Handle ontology translation if provided - ontology_translation_path = None - if ontology_translation is not None: - with tempfile.NamedTemporaryFile( - delete=False, suffix=".json", mode="w" - ) as tmp_trans: - json.dump(json.load(ontology_translation), tmp_trans) - ontology_translation_path = tmp_trans.name - - # Prepare predictions output directory if needed - predictions_outdir = None - if save_predictions or save_visualizations: - if predictions_outdir_input and predictions_outdir_input.strip(): - predictions_outdir = predictions_outdir_input.strip() - os.makedirs(predictions_outdir, exist_ok=True) - else: - predictions_outdir = tempfile.mkdtemp( - prefix="eval_predictions_" - ) - - # Create progress bar for evaluation - progress_bar = st.progress(0) - status_text = st.empty() - - # Create placeholders for intermediate metrics that will be updated in place - intermediate_metrics_placeholder = st.empty() - intermediate_table_placeholder = st.empty() - - def progress_callback(processed, total): - """Progress callback for Streamlit UI""" - try: - progress = processed / total if total > 0 else 0 - progress_bar.progress(progress) - status_text.text( - f"Processing: {processed}/{total} images ({progress:.1%})" - ) - except Exception as e: - st.error(f"Progress callback error: {e}") - - def metrics_callback(metrics_df, processed, total): - """Metrics callback for intermediate results display""" - try: - # Update the metrics placeholder with current summary metrics - if "mean" in metrics_df.columns: - mean_metrics = metrics_df["mean"] - - with intermediate_metrics_placeholder.container(): - st.markdown( - f"#### 📊 Intermediate Results (after {processed} images)" - ) - - col1, col2, col3 = st.columns(3) - with col1: - st.metric("mAP", f"{mean_metrics.get('AP', 0):.3f}") - with col2: - st.metric( - "Mean Precision", - f"{mean_metrics.get('Precision', 0):.3f}", - ) - with col3: - st.metric( - "Mean Recall", - f"{mean_metrics.get('Recall', 0):.3f}", - ) - - # Update the table placeholder with current per-class results - per_class_results = ( - metrics_df.drop(columns=["mean"]) - if "mean" in metrics_df.columns - else metrics_df - ) - per_class_results = per_class_results.drop( - ["AUC-PR", "mAP@[0.5:0.95]"], errors="ignore" - ) - - # Round for display - display_df = per_class_results.copy() - numeric_columns = display_df.select_dtypes( - include=["float64", "int64"] - ).columns - for col in numeric_columns: - if col in display_df.columns: - display_df[col] = display_df[col].round(3) - - with intermediate_table_placeholder.container(): - st.markdown("#### Per-Class Metrics (Intermediate)") - st.dataframe(display_df, width="stretch") - - except Exception as e: - st.error(f"Metrics callback error: {e}") - - # Run evaluation with progress tracking - # Use full dataset for evaluation - - try: - # Use the full dataset for evaluation - results = model.eval( - dataset=dataset, - split=split, - ontology_translation=ontology_translation_path, - predictions_outdir=predictions_outdir, - results_per_sample=save_predictions, - save_visualizations=save_visualizations, - progress_callback=progress_callback, - metrics_callback=metrics_callback, - ) - except Exception as e: - st.error(f"Error in model.eval(): {e}") - return - - # Results ready - - # Clear progress elements and intermediate results - progress_bar.empty() - status_text.empty() - intermediate_metrics_placeholder.empty() - intermediate_table_placeholder.empty() + task = st.session_state.get("task", "Image Detection") - # Store results in session state - st.session_state["evaluation_results"] = results - st.session_state["evaluation_config"] = { - "split": split, - "predictions_saved": save_predictions, - "visualizations_saved": save_visualizations, - } - - st.success("✅ Evaluation completed successfully!") - - except Exception as e: - st.error(f"❌ Evaluation failed: {e}") - import traceback - - st.code(traceback.format_exc()) - - # Display results (either from current evaluation or previous) - if "evaluation_results" in st.session_state: - display_evaluation_results(st.session_state["evaluation_results"]) - - -def display_evaluation_results(results): - """Display evaluation results in a comprehensive format""" - - if results is None: - st.warning("No evaluation results to display.") + if task == "Image Detection": + render_image_detection_evaluator() return - # Handle new results format (dictionary with metrics_df and metrics_factory) - if isinstance(results, dict): - metrics_df = results.get("metrics_df") - metrics_factory = results.get("metrics_factory") - else: - # Fallback for old format - metrics_df = results - metrics_factory = None - - if metrics_df is None or metrics_df.empty: - st.warning("No evaluation results to display.") + if task == "Image Segmentation": + render_image_segmentation_evaluator() return - # Display summary metrics - st.markdown("#### Summary Metrics") - - # Get mean metrics - mean is a column - if "mean" in metrics_df.columns: - mean_metrics = metrics_df["mean"] - - col1, col2, col3, col4, col5 = st.columns(5) - with col1: - st.metric("mAP", f"{mean_metrics.get('AP', 0):.3f}") - with col2: - st.metric("Mean Precision", f"{mean_metrics.get('Precision', 0):.3f}") - with col3: - st.metric("Mean Recall", f"{mean_metrics.get('Recall', 0):.3f}") - with col4: - coco_map = mean_metrics.get("mAP@[0.5:0.95]", 0) - st.metric("mAP@[0.5:0.95]", f"{coco_map:.3f}") - with col5: - auc_pr = mean_metrics.get("AUC-PR", 0) - st.metric("AUC-PR", f"{auc_pr:.3f}") - - # Display per-class metrics first - st.markdown("#### Per-Class Metrics") - - # Filter out the 'mean' column for per-class display - per_class_results = ( - metrics_df.drop(columns=["mean"]) - if "mean" in metrics_df.columns - else metrics_df - ) - - # Remove overall metrics rows (AUC-PR and mAP@[0.5:0.95]) from per-class display - per_class_results = per_class_results.drop( - ["AUC-PR", "mAP@[0.5:0.95]"], errors="ignore" - ) - - # Create a more readable display - display_df = per_class_results.copy() - - # Round numeric columns for better display - numeric_columns = display_df.select_dtypes(include=["float64", "int64"]).columns - for col in numeric_columns: - if col in display_df.columns: - display_df[col] = display_df[col].round(3) - - st.dataframe(display_df, width="stretch") - - # Now display Precision-Recall Curve - if metrics_factory is not None: - st.markdown("#### Precision-Recall Curve") - - try: - # Get the precision-recall curve data - curve_data = metrics_factory.get_overall_precision_recall_curve() - auc_pr = metrics_factory.compute_auc_pr() - - # Create the plot using streamlit's plotly integration - import plotly.graph_objects as go - - # Create the precision-recall curve - fig = go.Figure() - - # Add the curve - fig.add_trace( - go.Scatter( - x=curve_data["recall"], - y=curve_data["precision"], - mode="lines", - name="Precision-Recall Curve", - line=dict(color="blue", width=2), - fill="tonexty", - fillcolor="rgba(0, 0, 255, 0.1)", - ) - ) - - # Add AUC-PR annotation - fig.add_annotation( - x=0.6, - y=0.2, - text=f"AUC-PR: {auc_pr:.3f}", - showarrow=False, - font=dict(size=12), - bgcolor="white", - bordercolor="black", - borderwidth=1, - ) - - # Update layout - fig.update_layout( - # title='Overall Precision-Recall Curve', - xaxis_title="Recall", - yaxis_title="Precision", - xaxis=dict(range=[0, 1]), - yaxis=dict(range=[0, 1]), - showlegend=True, - height=500, - ) - - # Add grid - fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor="lightgray") - fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor="lightgray") - - st.plotly_chart(fig, width="stretch") - - except Exception as e: - st.error(f"Error plotting precision-recall curve: {e}") - st.info("Precision-recall curve data not available.") - - # Download results - st.markdown("#### Download Results") - - # Convert to CSV for download - csv = metrics_df.to_csv(index=True) - st.download_button( - label="📥 Download per class metrics", - data=csv, - file_name="evaluation_results.csv", - mime="text/csv", - ) - try: - curve_data = ( - metrics_factory.get_overall_precision_recall_curve() - if metrics_factory is not None - else None - ) - if curve_data is not None: - import pandas as pd - - pr_points_df = pd.DataFrame( - {"recall": curve_data["recall"], "precision": curve_data["precision"]} - ) - pr_csv = pr_points_df.to_csv(index=False) - st.download_button( - label="📈 Download precision-recall points", - data=pr_csv, - file_name="precision_recall_points.csv", - mime="text/csv", - ) - else: - st.write("No precision-recall data available.") - except Exception as e: - st.write(f"Error preparing precision-recall points: {e}") - - # Show detailed statistics - with st.expander("📊 Detailed Statistics"): - st.markdown("**Results Shape:**") - st.write(f"Rows: {metrics_df.shape[0]}, Columns: {metrics_df.shape[1]}") - - st.markdown("**Available Metrics:**") - st.write(list(metrics_df.columns)) - - st.markdown("**Class Names:**") - st.write( - list(metrics_df.index) if len(metrics_df.index) > 0 else "No classes found" - ) - - st.markdown("**DataFrame Info:**") - st.write("Index:", metrics_df.index.tolist()) - st.write("Columns:", metrics_df.columns.tolist()) - - st.markdown("**Sample Data:**") - st.dataframe(metrics_df.head(), width="stretch") - - if "evaluation_config" in st.session_state: - st.markdown("**Evaluation Configuration:**") - config = st.session_state["evaluation_config"] - for key, value in config.items(): - st.write(f"- {key}: {value}") + if task == "Lidar Segmentation": + render_lidar_segmentation_evaluator() + return - # Show precision-recall curve data if available - if metrics_factory is not None: - st.markdown("**Precision-Recall Curve Data:**") - try: - curve_data = metrics_factory.get_overall_precision_recall_curve() - st.write(f"Number of points: {len(curve_data['precision'])}") - st.write( - f"Precision range: {min(curve_data['precision']):.3f} - {max(curve_data['precision']):.3f}" - ) - st.write( - f"Recall range: {min(curve_data['recall']):.3f} - {max(curve_data['recall']):.3f}" - ) - st.write(f"AUC-PR: {metrics_factory.compute_auc_pr():.3f}") - except Exception as e: - st.write(f"Error accessing curve data: {e}") + st.error(f"Unsupported task for evaluator: {task}") diff --git a/tabs/inference.py b/tabs/inference.py index a2481828..8213aa30 100644 --- a/tabs/inference.py +++ b/tabs/inference.py @@ -1,148 +1,22 @@ -from typing import Optional - import streamlit as st -import json -from PIL import Image -try: - import torch -except ImportError: - raise ImportError( - "PyTorch is required for GUI-based inference and evaluation. " - ) - - -def draw_detections(image: Image, predictions: dict, label_map: Optional[dict] = None): - """Draw color-coded bounding boxes and labels on the image using supervision. - - :param image: PIL Image - :type image: Image.Image - :param predictions: dict with 'boxes', 'labels', 'scores' (torch tensors) - :type predictions: dict - :param label_map: dict mapping label indices to class names (optional) - :type label_map: dict - :return: np.ndarray with detections drawn (for st.image) - :rtype: np.ndarray - """ - from perceptionmetrics.utils import image as ui - - boxes = predictions.get("boxes", torch.empty(0)).cpu().numpy() - class_ids = predictions.get("labels", torch.empty(0)).cpu().numpy().astype(int) - - scores_tensor = predictions.get("scores") - if scores_tensor is not None and len(scores_tensor) > 0: - scores = scores_tensor.cpu().numpy() - else: - scores = None - - if label_map: - class_names = [label_map.get(int(label), str(label)) for label in class_ids] - else: - class_names = [str(label) for label in class_ids] - - return ui.draw_detections( - image=image, - boxes=boxes, - class_ids=class_ids, - class_names=class_names, - scores=scores, - ) +from tabs.tasks.image_detection.inference import render_image_detection_inference +from tabs.tasks.image_segmentation.inference import render_image_segmentation_inference +from tabs.tasks.lidar_segmentation.inference import render_lidar_segmentation_inference def inference_tab(): - st.header("Model Inference") - st.markdown("Select an image and run inference using the loaded model.") + task = st.session_state.get("task", "Image Detection") - # Check if a model has been loaded and saved in session - if ( - "detection_model" not in st.session_state - or st.session_state.detection_model is None - ): - st.warning("⚠️ Load a model from the sidebar to start inference") + if task == "Image Detection": + render_image_detection_inference() return - st.success("Model loaded and saved. You can now select an image.") - - # Image picker in the tab - image_file = st.file_uploader( - "Choose an image", - type=["jpg", "jpeg", "png"], - key="inference_image_file", - help="Upload an image to run inference", - ) - - if image_file is not None: - with st.spinner("Running inference..."): - try: - image = Image.open(image_file).convert("RGB") - predictions = st.session_state.detection_model.predict(image) - - label_map = getattr( - st.session_state.detection_model, "idx_to_class_name", None - ) - result_img = draw_detections(image, predictions, label_map) - - st.markdown("#### Detection Results") - st.image(result_img, caption="Detection Results", width="content") - - # Display detection statistics - if ( - predictions.get("scores") is not None - and len(predictions["scores"]) > 0 - ): - st.markdown("#### Detection Statistics") - col1, col2, col3 = st.columns(3) - with col1: - st.metric("Total Detections", len(predictions["scores"])) - with col2: - avg_confidence = float(predictions["scores"].mean()) - st.metric("Avg Confidence", f"{avg_confidence:.3f}") - with col3: - max_confidence = float(predictions["scores"].max()) - st.metric("Max Confidence", f"{max_confidence:.3f}") - - # Display and download detection results - st.markdown("#### Detection Details") - - # Convert predictions to JSON format - detection_results = [] - boxes = predictions.get("boxes", torch.empty(0)).cpu().numpy() - labels = predictions.get("labels", torch.empty(0)).cpu().numpy() - scores = predictions.get("scores", torch.empty(0)).cpu().numpy() - - for i in range(len(scores)): - class_name = ( - label_map.get(int(labels[i]), f"class_{labels[i]}") - if label_map - else f"class_{labels[i]}" - ) - detection_results.append( - { - "detection_id": i, - "class_id": int(labels[i]), - "class_name": class_name, - "confidence": float(scores[i]), - "bbox": { - "x1": float(boxes[i][0]), - "y1": float(boxes[i][1]), - "x2": float(boxes[i][2]), - "y2": float(boxes[i][3]), - }, - "bbox_xyxy": boxes[i].tolist(), - } - ) + if task == "Image Segmentation": + render_image_segmentation_inference() + return - with st.expander(" View Detection Results (JSON)", expanded=False): - st.json(detection_results) + if task == "Lidar Segmentation": + render_lidar_segmentation_inference() + return - json_str = json.dumps(detection_results, indent=2) - st.download_button( - label="Download Detection Results as JSON", - data=json_str, - file_name="detection_results.json", - mime="application/json", - help="Download the detection results as a JSON file", - ) - else: - st.info("No detections found in the image.") - except Exception as e: - st.error(f"Failed to run inference: {e}") + st.error(f"Unsupported task for inference: {task}") diff --git a/tabs/tasks/image_detection/dataset_viewer.py b/tabs/tasks/image_detection/dataset_viewer.py new file mode 100644 index 00000000..2eb3f72a --- /dev/null +++ b/tabs/tasks/image_detection/dataset_viewer.py @@ -0,0 +1,198 @@ +import streamlit as st +import os + +from perceptionmetrics.datasets.coco import find_img_dir_and_ann_file +from tabs.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 + from supervision.detection.annotate import BoxAnnotator + from supervision.detection.core import Detections + + # Get inputs from session state + dataset_path = st.session_state.get("dataset_path", "") + dataset_type = st.session_state.get("dataset_type", "COCO").lower() + split = st.session_state.get("split", "val") + + # Header row only + st.header("Dataset Viewer") + + if not dataset_path or not os.path.isdir(dataset_path): + st.warning("⚠️ Please select a valid dataset folder.") + return + + # Setup paths and pagination + if dataset_type == "coco": + try: + img_dir, ann_file = 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}") + 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}" + 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 + + 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 + image_files = [ + f for f in os.listdir(img_dir) if f.lower().endswith((".jpg", ".jpeg", ".png")) + ] + if not image_files: + st.warning("No images found.") + return + + image_paths = [os.path.join(img_dir, img_name) for img_name in image_files] + selected_img_path, _ = render_image_grid( + item_names=image_files, + image_paths=image_paths, + state_prefix="image_detection", + context=f"{dataset_path}_{split}", + search_label="image", + ) + + # Display selected image with annotations + if selected_img_path: + selected_img_name = os.path.basename(selected_img_path) + try: + img = Image.open(selected_img_path).convert("RGB") + img_np = np.array(img) + + if dataset_type == "yolo": + ann_row = dataset.dataset[ + dataset.dataset["image"].str.endswith(selected_img_name) + ] + else: + ann_row = dataset.dataset[dataset.dataset["image"] == selected_img_name] + + if not ann_row.empty: + annotation_id = ann_row.iloc[0]["annotation"] + if dataset_type == "yolo": + annotation_id = os.path.join(dataset.dataset_dir, annotation_id) + + boxes, category_indices = dataset.read_annotation(annotation_id) + + # Get class names from ontology + ontology = getattr(dataset, "ontology", None) + if ontology is None and hasattr(dataset.dataset, "attrs"): + ontology = dataset.dataset.attrs.get("ontology", None) + + if ontology: + catid_to_name = {v["idx"]: k for k, v in ontology.items()} + class_names = [ + catid_to_name.get(cat_id, str(cat_id)) + for cat_id in category_indices + ] + else: + class_names = [str(cat_id) for cat_id in category_indices] + + # Annotate image + palette = ColorPalette.default() + detections = Detections( + xyxy=np.array(boxes), class_id=np.array(category_indices) + ) + annotator = BoxAnnotator( + color=palette, text_scale=0.7, text_thickness=1, text_padding=2 + ) + annotated_img = annotator.annotate( + scene=img_np, detections=detections, labels=class_names + ) + + # Resize for display + annotated_pil = Image.fromarray(annotated_img) + try: + resample = getattr(Image, "Resampling", Image).LANCZOS + except AttributeError: + resample = Image.LANCZOS + annotated_pil.thumbnail((500, 500), resample) + st.image(annotated_pil, width="content") + else: + st.warning("No annotation found for this image.") + except Exception as e: + st.error(f"Error displaying image: {e}") + else: + st.info("Select an image to view with annotations.") diff --git a/tabs/tasks/image_detection/evaluator.py b/tabs/tasks/image_detection/evaluator.py new file mode 100644 index 00000000..a0e788d0 --- /dev/null +++ b/tabs/tasks/image_detection/evaluator.py @@ -0,0 +1,508 @@ +import streamlit as st +import os +import tempfile +import json +from perceptionmetrics.datasets.coco import CocoDataset + + +from tabs.tasks.utils import browse_folder +from perceptionmetrics.datasets.coco import find_img_dir_and_ann_file + + +def browse_predictions_outdir(): + folder = browse_folder() + if folder: + st.session_state.predictions_outdir = folder + + +def render_image_detection_evaluator(): + """Render the image detection evaluator tab in Streamlit.""" + st.header("Evaluator") + st.markdown("Evaluate your model on the loaded dataset using PerceptionMetrics.") + + # Check if we have the required objects from sidebar inputs + dataset_available = False + model_available = False + dataset = None + model = None + + # Check for dataset from sidebar inputs + dataset_path = st.session_state.get("dataset_path", "") + 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): + 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." + ) + except Exception as e: + st.error(f"❌ Error loading dataset: {e}") + else: + st.warning( + "⚠️ No dataset path provided. Please set the dataset path in the sidebar." + ) + + # Check for model from sidebar (loaded via Load Model button) + if ( + "detection_model" in st.session_state + and st.session_state.detection_model is not None + ): + model = st.session_state.detection_model + model_available = True + st.success("✅ Model loaded and ready for evaluation") + else: + st.warning( + "⚠️ No model loaded. Please load a model using the 'Load Model' button in the sidebar." + ) + + # Evaluation configuration + st.markdown("### Evaluation Configuration") + + save_predictions = st.checkbox( + "Save Predictions", + value=False, + help="Save individual predictions and metrics per sample", + ) + + save_visualizations = st.checkbox( + "Save Visualizations", + value=False, + help="Save visualized qualitative results (Image + GT + Preds)", + ) + + predictions_outdir_input = None + if save_predictions or save_visualizations: + col1, col2 = st.columns([3, 1]) + with col1: + st.text_input("Predictions Output Directory", key="predictions_outdir") + with col2: + st.markdown( + "
", unsafe_allow_html=True + ) + st.button( + "Browse", + on_click=browse_predictions_outdir, + key="browse_preds_outdir", + ) + + predictions_outdir_input = st.session_state.get("predictions_outdir") + + ontology_translation = st.file_uploader( + "Ontology Translation (Optional)", + type=["json"], + help="JSON file for translating between dataset and model ontologies", + ) + + # Disable Run Evaluation button if output dir is missing + output_dir_required = save_predictions or save_visualizations + output_dir_missing = output_dir_required and not ( + predictions_outdir_input and predictions_outdir_input.strip() + ) + + if output_dir_missing: + st.warning( + "⚠️ Please provide a Predictions Output Directory to enable evaluation " + "when 'Save Predictions' or 'Save Visualizations' is turned on." + ) + + # Run evaluation button + if st.button( + "🚀 Run Evaluation", + type="primary", + disabled=not (dataset_available and model_available) or output_dir_missing, + ): + if not dataset_available or not model_available: + st.error( + "Please ensure both dataset and model are loaded before running evaluation." + ) + return + + # Prepare evaluation + with st.spinner("Running evaluation..."): + try: + # Validate dataset and model + if len(dataset.dataset) == 0: + st.error( + "Dataset has no samples. Please check the dataset configuration." + ) + return + + if not hasattr(model, "model_cfg") or model.model_cfg is None: + st.error( + "Model configuration is missing. Please reload the model in the Inference tab." + ) + return + + # Handle ontology translation if provided + ontology_translation_path = None + if ontology_translation is not None: + with tempfile.NamedTemporaryFile( + delete=False, suffix=".json", mode="w" + ) as tmp_trans: + json.dump(json.load(ontology_translation), tmp_trans) + ontology_translation_path = tmp_trans.name + + # Prepare predictions output directory if needed + predictions_outdir = None + if save_predictions or save_visualizations: + if predictions_outdir_input and predictions_outdir_input.strip(): + predictions_outdir = predictions_outdir_input.strip() + os.makedirs(predictions_outdir, exist_ok=True) + else: + predictions_outdir = tempfile.mkdtemp( + prefix="eval_predictions_" + ) + + # Create progress bar for evaluation + progress_bar = st.progress(0) + status_text = st.empty() + + # Create placeholders for intermediate metrics that will be updated in place + intermediate_metrics_placeholder = st.empty() + intermediate_table_placeholder = st.empty() + + def progress_callback(processed, total): + """Progress callback for Streamlit UI""" + try: + progress = processed / total if total > 0 else 0 + progress_bar.progress(progress) + status_text.text( + f"Processing: {processed}/{total} images ({progress:.1%})" + ) + except Exception as e: + st.error(f"Progress callback error: {e}") + + def metrics_callback(metrics_df, processed, total): + """Metrics callback for intermediate results display""" + try: + # Update the metrics placeholder with current summary metrics + if "mean" in metrics_df.columns: + mean_metrics = metrics_df["mean"] + + with intermediate_metrics_placeholder.container(): + st.markdown( + f"#### 📊 Intermediate Results (after {processed} images)" + ) + + col1, col2, col3 = st.columns(3) + with col1: + st.metric("mAP", f"{mean_metrics.get('AP', 0):.3f}") + with col2: + st.metric( + "Mean Precision", + f"{mean_metrics.get('Precision', 0):.3f}", + ) + with col3: + st.metric( + "Mean Recall", + f"{mean_metrics.get('Recall', 0):.3f}", + ) + + # Update the table placeholder with current per-class results + per_class_results = ( + metrics_df.drop(columns=["mean"]) + if "mean" in metrics_df.columns + else metrics_df + ) + per_class_results = per_class_results.drop( + ["AUC-PR", "mAP@[0.5:0.95]"], errors="ignore" + ) + + # Round for display + display_df = per_class_results.copy() + numeric_columns = display_df.select_dtypes( + include=["float64", "int64"] + ).columns + for col in numeric_columns: + if col in display_df.columns: + display_df[col] = display_df[col].round(3) + + with intermediate_table_placeholder.container(): + st.markdown("#### Per-Class Metrics (Intermediate)") + st.dataframe(display_df, width="stretch") + + except Exception as e: + st.error(f"Metrics callback error: {e}") + + # Run evaluation with progress tracking + # Use full dataset for evaluation + + try: + # Use the full dataset for evaluation + results = model.eval( + dataset=dataset, + split=split, + ontology_translation=ontology_translation_path, + predictions_outdir=predictions_outdir, + results_per_sample=save_predictions, + save_visualizations=save_visualizations, + progress_callback=progress_callback, + metrics_callback=metrics_callback, + ) + except Exception as e: + st.error(f"Error in model.eval(): {e}") + return + + # Results ready + + # Clear progress elements and intermediate results + progress_bar.empty() + status_text.empty() + intermediate_metrics_placeholder.empty() + intermediate_table_placeholder.empty() + + # Store results in session state + st.session_state["evaluation_results"] = results + st.session_state["evaluation_config"] = { + "split": split, + "predictions_saved": save_predictions, + "visualizations_saved": save_visualizations, + } + + st.success("✅ Evaluation completed successfully!") + + except Exception as e: + st.error(f"❌ Evaluation failed: {e}") + import traceback + + st.code(traceback.format_exc()) + + # Display results (either from current evaluation or previous) + if "evaluation_results" in st.session_state: + display_evaluation_results(st.session_state["evaluation_results"]) + + +def display_evaluation_results(results): + """Display evaluation results in a comprehensive format""" + + if results is None: + st.warning("No evaluation results to display.") + return + + # Handle new results format (dictionary with metrics_df and metrics_factory) + if isinstance(results, dict): + metrics_df = results.get("metrics_df") + metrics_factory = results.get("metrics_factory") + else: + # Fallback for old format + metrics_df = results + metrics_factory = None + + if metrics_df is None or metrics_df.empty: + st.warning("No evaluation results to display.") + return + + # Display summary metrics + st.markdown("#### Summary Metrics") + + # Get mean metrics - mean is a column + if "mean" in metrics_df.columns: + mean_metrics = metrics_df["mean"] + + col1, col2, col3, col4, col5 = st.columns(5) + with col1: + st.metric("mAP", f"{mean_metrics.get('AP', 0):.3f}") + with col2: + st.metric("Mean Precision", f"{mean_metrics.get('Precision', 0):.3f}") + with col3: + st.metric("Mean Recall", f"{mean_metrics.get('Recall', 0):.3f}") + with col4: + coco_map = mean_metrics.get("mAP@[0.5:0.95]", 0) + st.metric("mAP@[0.5:0.95]", f"{coco_map:.3f}") + with col5: + auc_pr = mean_metrics.get("AUC-PR", 0) + st.metric("AUC-PR", f"{auc_pr:.3f}") + + # Display per-class metrics first + st.markdown("#### Per-Class Metrics") + + # Filter out the 'mean' column for per-class display + per_class_results = ( + metrics_df.drop(columns=["mean"]) + if "mean" in metrics_df.columns + else metrics_df + ) + + # Remove overall metrics rows (AUC-PR and mAP@[0.5:0.95]) from per-class display + per_class_results = per_class_results.drop( + ["AUC-PR", "mAP@[0.5:0.95]"], errors="ignore" + ) + + # Create a more readable display + display_df = per_class_results.copy() + + # Round numeric columns for better display + numeric_columns = display_df.select_dtypes(include=["float64", "int64"]).columns + for col in numeric_columns: + if col in display_df.columns: + display_df[col] = display_df[col].round(3) + + st.dataframe(display_df, width="stretch") + + # Now display Precision-Recall Curve + if metrics_factory is not None: + st.markdown("#### Precision-Recall Curve") + + try: + # Get the precision-recall curve data + curve_data = metrics_factory.get_overall_precision_recall_curve() + auc_pr = metrics_factory.compute_auc_pr() + + # Create the plot using streamlit's plotly integration + import plotly.graph_objects as go + + # Create the precision-recall curve + fig = go.Figure() + + # Add the curve + fig.add_trace( + go.Scatter( + x=curve_data["recall"], + y=curve_data["precision"], + mode="lines", + name="Precision-Recall Curve", + line=dict(color="blue", width=2), + fill="tonexty", + fillcolor="rgba(0, 0, 255, 0.1)", + ) + ) + + # Add AUC-PR annotation + fig.add_annotation( + x=0.6, + y=0.2, + text=f"AUC-PR: {auc_pr:.3f}", + showarrow=False, + font=dict(size=12), + bgcolor="white", + bordercolor="black", + borderwidth=1, + ) + + # Update layout + fig.update_layout( + # title='Overall Precision-Recall Curve', + xaxis_title="Recall", + yaxis_title="Precision", + xaxis=dict(range=[0, 1]), + yaxis=dict(range=[0, 1]), + showlegend=True, + height=500, + ) + + # Add grid + fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor="lightgray") + fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor="lightgray") + + st.plotly_chart(fig, width="stretch") + + except Exception as e: + st.error(f"Error plotting precision-recall curve: {e}") + st.info("Precision-recall curve data not available.") + + # Download results + st.markdown("#### Download Results") + + # Convert to CSV for download + csv = metrics_df.to_csv(index=True) + st.download_button( + label="📥 Download per class metrics", + data=csv, + file_name="evaluation_results.csv", + mime="text/csv", + ) + try: + curve_data = ( + metrics_factory.get_overall_precision_recall_curve() + if metrics_factory is not None + else None + ) + if curve_data is not None: + import pandas as pd + + pr_points_df = pd.DataFrame( + {"recall": curve_data["recall"], "precision": curve_data["precision"]} + ) + pr_csv = pr_points_df.to_csv(index=False) + st.download_button( + label="📈 Download precision-recall points", + data=pr_csv, + file_name="precision_recall_points.csv", + mime="text/csv", + ) + else: + st.write("No precision-recall data available.") + except Exception as e: + st.write(f"Error preparing precision-recall points: {e}") + + # Show detailed statistics + with st.expander("📊 Detailed Statistics"): + st.markdown("**Results Shape:**") + st.write(f"Rows: {metrics_df.shape[0]}, Columns: {metrics_df.shape[1]}") + + st.markdown("**Available Metrics:**") + st.write(list(metrics_df.columns)) + + st.markdown("**Class Names:**") + st.write( + list(metrics_df.index) if len(metrics_df.index) > 0 else "No classes found" + ) + + st.markdown("**DataFrame Info:**") + st.write("Index:", metrics_df.index.tolist()) + st.write("Columns:", metrics_df.columns.tolist()) + + st.markdown("**Sample Data:**") + st.dataframe(metrics_df.head(), width="stretch") + + if "evaluation_config" in st.session_state: + st.markdown("**Evaluation Configuration:**") + config = st.session_state["evaluation_config"] + for key, value in config.items(): + st.write(f"- {key}: {value}") + + # Show precision-recall curve data if available + if metrics_factory is not None: + st.markdown("**Precision-Recall Curve Data:**") + try: + curve_data = metrics_factory.get_overall_precision_recall_curve() + st.write(f"Number of points: {len(curve_data['precision'])}") + st.write( + f"Precision range: {min(curve_data['precision']):.3f} - {max(curve_data['precision']):.3f}" + ) + st.write( + f"Recall range: {min(curve_data['recall']):.3f} - {max(curve_data['recall']):.3f}" + ) + st.write(f"AUC-PR: {metrics_factory.compute_auc_pr():.3f}") + except Exception as e: + st.write(f"Error accessing curve data: {e}") diff --git a/tabs/tasks/image_detection/inference.py b/tabs/tasks/image_detection/inference.py new file mode 100644 index 00000000..a884d23c --- /dev/null +++ b/tabs/tasks/image_detection/inference.py @@ -0,0 +1,153 @@ +from typing import Optional + +import streamlit as st +import json +from PIL import Image + +try: + import torch +except ImportError: + raise ImportError("PyTorch is required for GUI-based inference and evaluation. ") + + +def draw_detections(image: Image, predictions: dict, label_map: Optional[dict] = None): + """Draw color-coded bounding boxes and labels on the image using supervision. + + :param image: PIL Image + :type image: Image.Image + :param predictions: dict with 'boxes', 'labels', 'scores' (torch tensors) + :type predictions: dict + :param label_map: dict mapping label indices to class names (optional) + :type label_map: dict + :return: np.ndarray with detections drawn (for st.image) + :rtype: np.ndarray + """ + from perceptionmetrics.utils import image as ui + + boxes = predictions.get("boxes", torch.empty(0)).cpu().numpy() + class_ids = predictions.get("labels", torch.empty(0)).cpu().numpy().astype(int) + + scores_tensor = predictions.get("scores") + if scores_tensor is not None and len(scores_tensor) > 0: + scores = scores_tensor.cpu().numpy() + else: + scores = None + + if label_map: + class_names = [label_map.get(int(label), str(label)) for label in class_ids] + else: + class_names = [str(label) for label in class_ids] + + return ui.draw_detections( + image=image, + boxes=boxes, + class_ids=class_ids, + class_names=class_names, + scores=scores, + ) + + +def render_image_detection_inference(): + """Render the image detection inference tab in Streamlit.""" + + st.header("Model Inference") + st.markdown("Select an image and run inference using the loaded model.") + + # Check if a model has been loaded and saved in session + if ( + "detection_model" not in st.session_state + or st.session_state.detection_model is None + ): + st.warning("⚠️ Load a model from the sidebar to start inference") + return + + st.success("Model loaded and saved. You can now select an image.") + + # Image picker in the tab + image_file = st.file_uploader( + "Choose an image", + type=["jpg", "jpeg", "png"], + key="inference_image_file", + help="Upload an image to run inference", + ) + + if image_file is not None: + with st.spinner("Running inference..."): + try: + image = Image.open(image_file).convert("RGB") + predictions, sample_tensor = st.session_state.detection_model.predict( + image, return_sample=True + ) + from torchvision.transforms import v2 as transforms + + img_to_draw = transforms.ToPILImage()(sample_tensor[0]) + label_map = getattr( + st.session_state.detection_model, "idx_to_class_name", None + ) + result_img = draw_detections(img_to_draw, predictions, label_map) + + st.markdown("#### Detection Results") + st.image(result_img, caption="Detection Results", width="stretch") + + # Display detection statistics + if ( + predictions.get("scores") is not None + and len(predictions["scores"]) > 0 + ): + st.markdown("#### Detection Statistics") + col1, col2, col3 = st.columns(3) + with col1: + st.metric("Total Detections", len(predictions["scores"])) + with col2: + avg_confidence = float(predictions["scores"].mean()) + st.metric("Avg Confidence", f"{avg_confidence:.3f}") + with col3: + max_confidence = float(predictions["scores"].max()) + st.metric("Max Confidence", f"{max_confidence:.3f}") + + # Display and download detection results + st.markdown("#### Detection Details") + + # Convert predictions to JSON format + detection_results = [] + boxes = predictions.get("boxes", torch.empty(0)).cpu().numpy() + labels = predictions.get("labels", torch.empty(0)).cpu().numpy() + scores = predictions.get("scores", torch.empty(0)).cpu().numpy() + + for i in range(len(scores)): + class_name = ( + label_map.get(int(labels[i]), f"class_{labels[i]}") + if label_map + else f"class_{labels[i]}" + ) + detection_results.append( + { + "detection_id": i, + "class_id": int(labels[i]), + "class_name": class_name, + "confidence": float(scores[i]), + "bbox": { + "x1": float(boxes[i][0]), + "y1": float(boxes[i][1]), + "x2": float(boxes[i][2]), + "y2": float(boxes[i][3]), + }, + "bbox_xyxy": boxes[i].tolist(), + } + ) + + with st.expander(" View Detection Results (JSON)", expanded=False): + st.json(detection_results) + + json_str = json.dumps(detection_results, indent=2) + st.download_button( + label="Download Detection Results as JSON", + data=json_str, + file_name="detection_results.json", + mime="application/json", + help="Download the detection results as a JSON file", + ) + else: + st.info("No detections found in the image.") + except Exception as e: + st.error(f"Failed to run inference: {e}") diff --git a/tabs/tasks/image_detection/sidebar.py b/tabs/tasks/image_detection/sidebar.py new file mode 100644 index 00000000..a68845a6 --- /dev/null +++ b/tabs/tasks/image_detection/sidebar.py @@ -0,0 +1,370 @@ +import json +import os +import tempfile +from typing import Optional + +import streamlit as st + +from tabs.tasks.utils import browse_folder + + +def browse_dataset_path(): + folder = browse_folder() + if folder: + st.session_state.dataset_path = folder + + +def render_image_detection_sidebar(available_devices): + """Render the sidebar for the image detection task in Streamlit. + :param available_devices: List of available devices for model inference + :type available_devices: list + """ + + with st.expander("Image Detection Dataset", expanded=True): + col1, col2 = st.columns(2) + with col1: + st.selectbox( + "Type", ["COCO", "YOLO"], key="dataset_type" + ) # split into two columns + with col2: + st.selectbox("Split", ["train", "val", "test"], key="split") + + col1, col2 = st.columns([3, 1]) + with col1: + st.text_input("Dataset Folder", key="dataset_path") + with col2: + st.markdown( + "
", # add some spacing to align with the text input + unsafe_allow_html=True, + ) + st.button( + "Browse", + on_click=browse_dataset_path, + key="browse_detection_dataset_path", + ) + + if st.session_state.get("dataset_type", "COCO") == "YOLO": + st.file_uploader( + "Dataset Configuration (.yaml)", + type=["yaml"], + key="dataset_config_file", + help="Upload a YAML dataset configuration file.", + ) + + with st.expander("Image Detection Model", expanded=False): + st.file_uploader( + "Model File (.pt, .onnx, .h5, .pb, .pth, .torchscript)", + type=["pt", "onnx", "h5", "pb", "pth", "torchscript"], + key="model_file", + help="Upload your trained model file.", + max_upload_size=1024, + ) + st.file_uploader( + "Ontology File (.json)", + type=["json"], + key="ontology_file", + help="Upload a JSON file with class labels.", + ) + st.radio( + "Configuration Method:", + ["Manual Configuration", "Upload Config File"], + key="config_option", + horizontal=True, + ) # radio button to select between manual configuration and uploading a config file + if ( + st.session_state.get("config_option", "Manual Configuration") + == "Upload Config File" + ): + st.file_uploader( + "Configuration File (.json)", + type=["json"], + key="config_file", + help="Upload a JSON configuration file.", + ) + else: + _render_manual_detection_model_config(available_devices) + + if st.button( + "Load Model", + type="primary", + width="stretch", + help="Load and save the model for use in the Inference tab", + key="sidebar_load_model_btn", + ): + load_image_detection_model() + + +def _render_manual_detection_model_config(available_devices): + """Render the manual configuration options for the image detection model in the sidebar. + :param available_devices: List of available devices for model inference + :type available_devices: list + """ + col1, col2 = st.columns(2) + with col1: + st.slider( + "Confidence Threshold", + min_value=0.0, + max_value=1.0, + step=0.01, + key="confidence_threshold", + help="Minimum confidence score for detections", + ) + st.slider( + "NMS Threshold", + min_value=0.0, + max_value=1.0, + step=0.01, + key="nms_threshold", + help="Non-maximum suppression threshold", + ) + st.number_input( + "Max Detections/Image", + min_value=1, + max_value=1000, + step=1, + key="max_detections", + ) + with col2: + st.selectbox("Device", available_devices, key="device") + st.selectbox( + "Model Format", + ["torchvision", "YOLO"], + index=( + 0 + if st.session_state.get("model_format", "torchvision") == "torchvision" + else 1 + ), + key="model_format", + ) + st.number_input( + "Batch Size", + min_value=1, + max_value=256, + step=1, + key="batch_size", + ) + st.number_input( + "Evaluation Step", + min_value=0, + max_value=1000, + step=1, + key="evaluation_step", + help="Update UI with intermediate metrics every N images (0 = disable intermediate updates)", + ) + + st.write("---") + st.write("**Image Size Configuration**") + + enable_resize = st.checkbox("Enable Resize", value=True, key="enable_resize") + + if enable_resize: + resize_strategy = st.radio( + "Resize Strategy", + ["Fixed Dimensions", "Min Side"], + key="resize_strategy", + horizontal=True, + label_visibility="collapsed", + ) + + if resize_strategy == "Fixed Dimensions": + c1, c2 = st.columns(2) + with c1: + st.number_input( + "Image Resize Height", + min_value=1, + max_value=4096, + value=640, + step=1, + key="resize_height", + help="Height to resize images for inference", + ) + with c2: + st.number_input( + "Image Resize Width", + min_value=1, + max_value=4096, + value=640, + step=1, + key="resize_width", + help="Width to resize images for inference", + ) + else: + st.number_input( + "Min Side", + min_value=1, + max_value=4096, + value=640, + step=1, + key="min_side", + help="Minimum size of the shorter side of the image", + ) + + enable_crop = st.checkbox("Enable Center Crop", key="enable_crop") + + if enable_crop: + c1, c2 = st.columns(2) + with c1: + st.number_input( + "Crop Height", + min_value=1, + max_value=4096, + value=640, + step=1, + key="crop_height", + help="Center crop height", + ) + with c2: + st.number_input( + "Crop Width", + min_value=1, + max_value=4096, + value=640, + step=1, + key="crop_width", + help="Center crop width", + ) + + +def load_image_detection_model(): + """Load the image detection model based on the provided configuration and ontology files.""" + from perceptionmetrics.models.torch_detection import TorchImageDetectionModel + + model_file = st.session_state.get("model_file") + ontology_file = st.session_state.get("ontology_file") + config_path = _write_detection_config() + + if model_file is None: + st.error("Please upload a model file") + elif config_path is None: + st.error("Please provide a valid model configuration") + elif ontology_file is None: + st.error("Please upload an ontology file") + else: + with st.spinner("Loading model..."): + ontology_path = _uploaded_json_to_tempfile(ontology_file) + model_temp_path = _uploaded_model_to_tempfile(model_file) + + if ontology_path and model_temp_path: + try: + model = TorchImageDetectionModel( + model=model_temp_path, + model_cfg=config_path, + ontology_fname=ontology_path, + device=st.session_state.get("device", "cpu"), + ) + st.session_state.detection_model = model + st.session_state.detection_model_loaded = True + st.success("Model loaded and saved for inference") + except Exception as e: + st.session_state.detection_model = None + st.session_state.detection_model_loaded = False + st.error(f"Failed to load model: {e}") + + +def _write_detection_config() -> Optional[str]: + """Write the detection configuration to a temporary JSON file based on the selected configuration method. + :return: Path to the temporary JSON configuration file or None if an error occurred + :rtype: Optional[str] + """ + config_option = st.session_state.get("config_option", "Manual Configuration") + config_file = ( + st.session_state.get("config_file") + if config_option == "Upload Config File" + else None + ) + + try: + if config_option == "Upload Config File": + if config_file is None: + st.error("Please upload a configuration file") + return None + config_data = json.load(config_file) + else: + config_data = _manual_detection_config() + + with tempfile.NamedTemporaryFile( + delete=False, suffix=".json", mode="w" + ) as tmp_cfg: + json.dump(config_data, tmp_cfg) + return tmp_cfg.name + except Exception as e: + st.error(f"Failed to prepare configuration: {e}") + return None + + +def _manual_detection_config() -> dict: + """Generate a configuration dictionary based on the manual configuration options in the sidebar. + :return: Configuration dictionary + :rtype: dict + """ + resize_cfg = None + if st.session_state.get("enable_resize", True): + resize_strategy = st.session_state.get("resize_strategy", "Fixed Dimensions") + if resize_strategy == "Fixed Dimensions": + resize_cfg = { + "height": int(st.session_state.get("resize_height", 640)), + "width": int(st.session_state.get("resize_width", 640)), + } + else: + resize_cfg = {"min_side": int(st.session_state.get("min_side", 640))} + + config_data = { + "confidence_threshold": float( + st.session_state.get("confidence_threshold", 0.5) + ), + "nms_threshold": float(st.session_state.get("nms_threshold", 0.5)), + "max_detections_per_image": int(st.session_state.get("max_detections", 100)), + "device": st.session_state.get("device", "cpu"), + "batch_size": int(st.session_state.get("batch_size", 1)), + "evaluation_step": int(st.session_state.get("evaluation_step", 5)), + "model_format": st.session_state.get("model_format", "torchvision").lower(), + } + if resize_cfg is not None: + config_data["resize"] = resize_cfg + + if st.session_state.get("enable_crop", False): + config_data["crop"] = { + "height": int(st.session_state.get("crop_height", 640)), + "width": int(st.session_state.get("crop_width", 640)), + } + + return config_data + + +def _uploaded_json_to_tempfile(uploaded_file) -> Optional[str]: + """Save the uploaded JSON file to a temporary file and return its path. + :param uploaded_file: Uploaded JSON file + :type uploaded_file: UploadedFile + :return: Path to the temporary JSON file or None if an error occurred + :rtype: Optional[str] + """ + + try: + data = json.load(uploaded_file) + with tempfile.NamedTemporaryFile( + delete=False, suffix=".json", mode="w" + ) as tmp_file: + json.dump(data, tmp_file) + return tmp_file.name + except Exception as e: + st.error(f"Failed to load JSON file: {e}") + return None + + +def _uploaded_model_to_tempfile(uploaded_file) -> Optional[str]: + """Save the uploaded model file to a temporary file and return its path. + :param uploaded_file: Uploaded model file + :type uploaded_file: UploadedFile + :return: Path to the temporary model file or None if an error occurred + :rtype: Optional[str] + """ + try: + suffix = os.path.splitext(uploaded_file.name)[1] or ".pt" + with tempfile.NamedTemporaryFile( + delete=False, suffix=suffix, mode="wb" + ) as tmp_model: + tmp_model.write(uploaded_file.read()) + return tmp_model.name + except Exception as e: + st.error(f"Failed to save model file: {e}") + return None diff --git a/tabs/tasks/image_segmentation/dataset_viewer.py b/tabs/tasks/image_segmentation/dataset_viewer.py new file mode 100644 index 00000000..911b11ea --- /dev/null +++ b/tabs/tasks/image_segmentation/dataset_viewer.py @@ -0,0 +1,227 @@ +import os + +import numpy as np +import pandas as pd +import streamlit as st +from PIL import Image + +from perceptionmetrics.datasets.cityscapes import CityscapesImageSegmentationDataset +from perceptionmetrics.datasets.nuimages import NuImagesSegmentationDataset +from tabs.tasks.utils import render_image_grid + + + + +def _overlay_mask(image, label, ontology, opacity): + """Overlay a segmentation mask on an image. + param image: PIL Image object of the original image. + param label: 2D numpy array of the segmentation mask. + param ontology: Dictionary mapping class names to their properties, including 'idx' and 'rgb'. + param opacity: Float value between 0 and 1 indicating the opacity of the overlay. + return: PIL Image object of the image with the overlay applied. + + """ + image_np = np.array(image) + color_mask = np.zeros((*label.shape, 3), dtype=np.uint8) + for class_data in ontology.values(): + class_idx = int(class_data["idx"]) + rgb = class_data.get("rgb") + if rgb is None: + rng = np.random.default_rng(abs(class_idx)) + rgb = tuple(int(value) for value in rng.integers(0, 255, size=3)) + color_mask[label == class_idx] = rgb + + resampling = getattr(Image, "Resampling", Image) + color_mask_image = Image.fromarray(color_mask).resize( + image.size, resampling.NEAREST + ) + color_mask_np = np.array(color_mask_image) + + overlay = ((1.0 - opacity) * image_np + opacity * color_mask_np).astype(np.uint8) + return Image.fromarray(overlay) + + +def render_image_segmentation_viewer(): + """Render the image segmentation dataset viewer tab in Streamlit.""" + dataset_type = st.session_state.get("segmentation_dataset_type", "Cityscapes") + dataset_path = st.session_state.get("dataset_path", "") + split = st.session_state.get("split", "val") + + st.header("Dataset Viewer") + + if not dataset_path or not os.path.isdir(dataset_path): + st.warning("Please select a valid image segmentation dataset folder.") + return + + try: + dataset = load_image_segmentation_dataset(dataset_type, dataset_path, split) + except Exception as exc: + st.error(f"Failed to load image segmentation dataset: {exc}") + return + + render_segmentation_dataset_viewer( + dataset=dataset, + dataset_type=dataset_type, + split=split, + state_prefix="image_segmentation", + context=f"{dataset_path}_{split}", + ) + + + +def render_segmentation_dataset_viewer( + dataset, + dataset_type, + split, + state_prefix, + context, +): + """Render a loaded image segmentation dataset.""" + split_df = dataset.dataset[dataset.dataset["split"] == split] + if split_df.empty: + st.warning(f"No {dataset_type} samples found for split '{split}'.") + return + + sample_names = split_df.index.astype(str).tolist() + image_paths = split_df["image"].tolist() + selected_img_path, sample_name = render_image_grid( + item_names=sample_names, + image_paths=image_paths, + state_prefix=state_prefix, + context=context, + search_label="sample", + ) + + if not selected_img_path: + st.info("Select an image to view the ground truth mask.") + return + + mask_opacity = st.slider( + "Mask Opacity", + min_value=0.0, + max_value=1.0, + value=0.45, + step=0.05, + key=f"{state_prefix}_mask_opacity", + ) + + row_key = sample_name + if row_key not in split_df.index and sample_name.isdigit(): + row_key = int(sample_name) + + row = split_df.loc[row_key] + image_fname = row["image"] + label_fname = row["label"] + + try: + image = Image.open(image_fname).convert("RGB") + label = dataset.read_label(label_fname) + except Exception as exc: + st.error(f"Failed to read sample '{sample_name}': {exc}") + return + + overlay = _overlay_mask(image, label, dataset.ontology, mask_opacity) + + image_col, overlay_col = st.columns(2) + with image_col: + st.image(image, caption="Image", use_container_width=True) + with overlay_col: + st.image(overlay, caption="Ground Truth Overlay", use_container_width=True) + + with st.expander("Classes", expanded=False): + st.dataframe(_classes_dataframe(dataset.ontology), use_container_width=True) + + +def load_image_segmentation_dataset(dataset_type, dataset_path, split): + if dataset_type == "Cityscapes": + return load_cityscapes_dataset(dataset_path, split) + if dataset_type == "NuImages": + return load_nuimages_dataset(dataset_path, split) + raise ValueError(f"{dataset_type} image segmentation dataset is not wired yet.") + + +def load_cityscapes_dataset(dataset_path, split): + """Load the Cityscapes dataset based on the provided path and split. + param dataset_path: Path to the Cityscapes dataset directory. + param split: Dataset split to load (e.g., "train", "val", "test"). + return: Instance of CityscapesImageSegmentationDataset as a session state variable. + """ + roots = {"train": None, "val": None, "test": None} + roots[split] = dataset_path + + dataset_key = ( + "cityscapes_segmentation_dataset", + os.path.abspath(dataset_path), + split, + st.session_state.get( + "segmentation_image_dir", "leftImg8bit_trainvaltest/leftImg8bit" + ), + st.session_state.get("segmentation_label_dir", "gtFine"), + st.session_state.get("segmentation_image_suffix", "_leftImg8bit.png"), + st.session_state.get("segmentation_label_suffix", "_gtFine_labelIds.png"), + st.session_state.get("segmentation_use_train_id", False), + ) + + if dataset_key not in st.session_state: + st.session_state[dataset_key] = CityscapesImageSegmentationDataset( + train_dataset_root=roots["train"], + val_dataset_root=roots["val"], + test_dataset_root=roots["test"], + image_dir=dataset_key[3], + label_dir=dataset_key[4], + image_suffix=dataset_key[5], + label_suffix=dataset_key[6], + use_train_id=dataset_key[7], + ) + + return st.session_state[dataset_key] + + +def load_nuimages_dataset(dataset_path, split): + """Load the NuImages dataset based on the provided path and split. + param dataset_path: Path to the NuImages dataset directory. + param split: Dataset split to load (e.g., "train", "val"). + return: Instance of NuImagesSegmentationDataset as a session state variable. + """ + + version = st.session_state.get("nuimages_segmentation_version", "v1.0-mini") + labels_rel_dir = st.session_state.get( + "nuimages_segmentation_labels_dir", + "generated/nuimages_segmentation_labels", + ) + dataset_key = ( + "nuimages_segmentation_dataset", + os.path.abspath(dataset_path), + version, + split, + labels_rel_dir, + ) + + if dataset_key not in st.session_state: + st.session_state[dataset_key] = NuImagesSegmentationDataset( + dataset_dir=dataset_path, + version=version, + split=split, + labels_rel_dir=labels_rel_dir, + ) + + return st.session_state[dataset_key] + + +def _classes_dataframe(ontology): + """Convert the ontology dictionary to a pandas DataFrame for display. + param ontology: Dictionary mapping class names to their properties, including 'idx', 'train_id', 'category', and 'rgb'. + return: pandas DataFrame with columns ['class', 'id', 'train_id', 'category', 'rgb']. + """ + rows = [] + for class_name, class_data in ontology.items(): + rows.append( + { + "class": class_name, + "id": class_data["idx"], + "train_id": class_data.get("train_id"), + "category": class_data.get("category"), + "rgb": class_data.get("rgb"), + } + ) + return pd.DataFrame(rows) diff --git a/tabs/tasks/image_segmentation/evaluator.py b/tabs/tasks/image_segmentation/evaluator.py new file mode 100644 index 00000000..fd05f78f --- /dev/null +++ b/tabs/tasks/image_segmentation/evaluator.py @@ -0,0 +1,209 @@ +import json +import os +import tempfile + +import streamlit as st + +from tabs.tasks.image_segmentation.dataset_viewer import ( + load_image_segmentation_dataset, +) +from tabs.tasks.utils import browse_folder + + +def browse_segmentation_predictions_outdir(): + folder = browse_folder() + if folder: + st.session_state.segmentation_predictions_outdir = folder + + +def render_image_segmentation_evaluator(): + """Render the image segmentation evaluator tab in the Streamlit app.""" + st.header("Evaluator") + st.markdown("Evaluate your model on the loaded dataset using PerceptionMetrics.") + + dataset_type = st.session_state.get("segmentation_dataset_type", "Cityscapes") + if dataset_type not in ["Cityscapes", "NuImages"]: + st.info(f"{dataset_type} image segmentation evaluation is not wired yet.") + return + + dataset = None + model = st.session_state.get("segmentation_model") + dataset_path = st.session_state.get("dataset_path", "") + split = st.session_state.get("split", "val") + + if not dataset_path or not os.path.isdir(dataset_path): + st.warning( + "No dataset path provided. Please set the dataset path in the sidebar." + ) + elif dataset_type == "NuImages" and split == "test": + st.warning("NuImages segmentation evaluation supports train and val splits.") + else: + try: + dataset = load_image_segmentation_dataset(dataset_type, dataset_path, split) + st.success( + f"✅ Dataset loaded: {dataset_path} ({split} split) - {len(dataset.dataset)} samples" + ) + except Exception as exc: + st.error(f"Error loading dataset: {exc}") + + if model is not None: + st.success("✅ Model loaded and ready for evaluation") + else: + st.warning( + "No model loaded. Please load a model using the " + "'Load Segmentation Model' button in the sidebar." + ) + + st.markdown("### Evaluation Configuration") + + save_predictions = st.checkbox( + "Save Predictions", + value=False, + help="Save predicted label images to an output directory.", + key="segmentation_save_predictions", + ) + + predictions_outdir = None + if save_predictions: + col1, col2 = st.columns([3, 1]) + with col1: + st.text_input( + "Predictions Output Directory", + key="segmentation_predictions_outdir", + ) + with col2: + st.markdown( + "
", + unsafe_allow_html=True, + ) + st.button( + "Browse", + on_click=browse_segmentation_predictions_outdir, + key="browse_segmentation_predictions_outdir", + ) + predictions_outdir = st.session_state.get("segmentation_predictions_outdir") + + ontology_translation = st.file_uploader( + "Ontology Translation (Optional)", + type=["json"], + key="segmentation_ontology_translation", + help="JSON file for translating between dataset and model ontologies.", + ) + if dataset_type == "Cityscapes": + st.info( + "For Cityscapes SegFormer models, use train-ID labels or provide a " + "label-ID to train-ID ontology translation." + ) + elif dataset_type == "NuImages": + st.info( + "For NuImages SegFormer models, upload " + "nuimages_to_model_ontology_translation.json and use dataset_to_model." + ) + + translation_direction = st.selectbox( + "Translation Direction", + ["dataset_to_model", "model_to_dataset"], + key="segmentation_translation_direction", + help=( + "dataset_to_model maps GT labels to model IDs. " + "model_to_dataset maps predictions to dataset IDs." + ), + ) + + output_dir_missing = save_predictions and not ( + predictions_outdir and predictions_outdir.strip() + ) + if output_dir_missing: + st.warning("Please provide a predictions output directory.") + + if st.button( + "🚀 Run Evaluation", + type="primary", + disabled=dataset is None or model is None or output_dir_missing, + key="run_segmentation_evaluation", + ): + with st.spinner("Running evaluation..."): + try: + ontology_translation_path = None + if ontology_translation is not None: + with tempfile.NamedTemporaryFile( + delete=False, suffix=".json", mode="w" + ) as tmp_trans: + json.dump(json.load(ontology_translation), tmp_trans) + ontology_translation_path = tmp_trans.name + + predictions_outdir = ( + predictions_outdir.strip() + if (save_predictions and predictions_outdir) + else None + ) + if predictions_outdir is not None: + os.makedirs(predictions_outdir, exist_ok=True) + + progress_bar = st.progress(0) + status_text = st.empty() + intermediate_metrics_placeholder = st.empty() + + def progress_callback(processed, total): + progress = processed / total if total > 0 else 0 + progress_bar.progress(progress) + status_text.text( + f"Processing: {processed}/{total} images ({progress:.1%})" + ) + + def metrics_callback(metrics_df, processed, total): + with intermediate_metrics_placeholder.container(): + st.markdown(f"#### Results (after {processed}/{total} images)") + display_segmentation_evaluation_results( + metrics_df, show_download=False + ) + + results = model.eval( + dataset=dataset, + split=split, + ontology_translation=ontology_translation_path, + translation_direction=translation_direction, + predictions_outdir=predictions_outdir, + results_per_sample=save_predictions, + progress_callback=progress_callback, + metrics_callback=metrics_callback, + ) + + progress_bar.empty() + status_text.empty() + intermediate_metrics_placeholder.empty() + + st.session_state["segmentation_evaluation_results"] = results + st.success("✅ Evaluation completed successfully!") + except Exception as exc: + st.error(f"Error in model.eval(): {exc}") + + if "segmentation_evaluation_results" in st.session_state: + display_segmentation_evaluation_results( + st.session_state["segmentation_evaluation_results"] + ) + + +def display_segmentation_evaluation_results(results, show_download=True): + """Display the evaluation results in a Streamlit dataframe and provide a download button. + Param results: pd.DataFrame, the evaluation results to display + Param show_download: bool, whether to show the download button for the results""" + if results is None or results.empty: + st.warning("No evaluation results to display.") + return + + st.markdown("#### Metrics") + display_df = results.copy() + numeric_columns = display_df.select_dtypes(include=["float64", "int64"]).columns + for col in numeric_columns: + display_df[col] = display_df[col].round(3) + st.dataframe(display_df, width="stretch") + + if show_download: + csv = results.to_csv(index=True) + st.download_button( + label="📥 Download segmentation metrics", + data=csv, + file_name="segmentation_evaluation_results.csv", + mime="text/csv", + ) diff --git a/tabs/tasks/image_segmentation/inference.py b/tabs/tasks/image_segmentation/inference.py new file mode 100644 index 00000000..8fbbb5cc --- /dev/null +++ b/tabs/tasks/image_segmentation/inference.py @@ -0,0 +1,63 @@ +import numpy as np +import streamlit as st +from PIL import Image + +from perceptionmetrics.utils import conversion as uc + + +def render_image_segmentation_inference(): + """Render the image segmentation inference tab in the Streamlit app.""" + st.header("Model Inference") + st.markdown("Select an image and run inference using the loaded model.") + + model = st.session_state.get("segmentation_model") + if model is None: + st.warning("Load a segmentation model from the sidebar to start inference.") + return + + image_file = st.file_uploader( + "Choose an image", + type=["jpg", "jpeg", "png"], + key="segmentation_inference_image_file", + help="Upload an image to run segmentation inference.", + ) + + if image_file is None: + return + + with st.spinner("Running segmentation inference..."): + try: + image = Image.open(image_file).convert("RGB") + pred = model.predict(image) + except Exception as exc: + st.error(f"Failed to run inference: {exc}") + return + + pred = uc.label_to_rgb(pred, model.ontology) + pred = pred.resize(image.size) + prediction_overlay = _overlay_mask(image, pred, opacity=0.45) + + cols = st.columns(3) + with cols[0]: + st.image(image, caption="Image", use_container_width=True) + with cols[1]: + st.image(pred, caption="Prediction", use_container_width=True) + with cols[2]: + st.image( + prediction_overlay, + caption="Prediction Overlay", + use_container_width=True, + ) + + +def _overlay_mask(image, mask_rgb, opacity): + """Overlay a segmentation mask on an image with a given opacity. + Param image: PIL.Image, the original image + Param mask_rgb: PIL.Image, the segmentation mask in RGB format + Param opacity: float, the opacity of the mask overlay (0.0 to 1.0) + Return: PIL.Image, the image with the mask overlay + """ + image_np = np.array(image) + mask_np = np.array(mask_rgb) + overlay = ((1.0 - opacity) * image_np + opacity * mask_np).astype(np.uint8) + return Image.fromarray(overlay) diff --git a/tabs/tasks/image_segmentation/sidebar.py b/tabs/tasks/image_segmentation/sidebar.py new file mode 100644 index 00000000..6679abf3 --- /dev/null +++ b/tabs/tasks/image_segmentation/sidebar.py @@ -0,0 +1,280 @@ +import os + +import streamlit as st + +from tabs.tasks.utils import browse_file, browse_folder + + +IMAGE_SEGMENTATION_DATASETS = [ + "Cityscapes", + "NuImages", + "GAIA", + "Generic", + "Wildscenes", + "RUGD", + "Rellis3D", + "GOOSE", +] + +MODEL_INPUT_DATASETS = ["Cityscapes", "NuImages"] + + +def browse_dataset_path(): + folder = browse_folder() + if folder: + st.session_state.dataset_path = folder + + +def browse_segmentation_model_path(): + if st.session_state.get("segmentation_model_type") == "Hugging Face SegFormer": + path = browse_folder() + else: + path = browse_file() + + if path: + st.session_state.segmentation_model_path = path + + +def browse_segmentation_config_path(): + path = browse_file() + if path: + st.session_state.segmentation_config_path = path + + +def browse_segmentation_ontology_path(): + path = browse_file() + if path: + st.session_state.segmentation_ontology_path = path + + +def render_image_segmentation_sidebar(_available_devices): + """Render the sidebar for the image segmentation task in Streamlit. + param _available_devices: List of available devices for model inference. + """ + with st.expander("Image Segmentation Dataset", expanded=True): + col1, col2 = st.columns(2) + with col1: + dataset_type = st.selectbox( + "Type", + IMAGE_SEGMENTATION_DATASETS, + key="segmentation_dataset_type", + ) + with col2: + st.selectbox("Split", ["train", "val", "test"], key="split") + + col1, col2 = st.columns([3, 1]) + with col1: + st.text_input("Dataset Folder", key="dataset_path") + with col2: + st.markdown( + "
", + unsafe_allow_html=True, + ) + st.button( + "Browse", + on_click=browse_dataset_path, + key="browse_segmentation_dataset_path", + ) + + if dataset_type == "Cityscapes": + render_cityscapes_dataset_inputs() + elif dataset_type == "NuImages": + render_nuimages_dataset_inputs() + else: + st.info(f"{dataset_type} image segmentation inputs are not wired yet.") + + with st.expander("Image Segmentation Model", expanded=False): + dataset_type = st.session_state.get("segmentation_dataset_type", "Cityscapes") + if dataset_type not in MODEL_INPUT_DATASETS: + st.info( + f"Image segmentation model loading is not wired for {dataset_type} yet." + ) + return + + render_segmentation_model_inputs() + + if st.button( + "Load Segmentation Model", + type="primary", + width="stretch", + key="sidebar_load_segmentation_model_btn", + ): + load_image_segmentation_model() + + +def render_cityscapes_dataset_inputs(): + """Render the input fields for the Cityscapes dataset in the sidebar.""" + st.text_input( + "Image Directory", + value="leftImg8bit_trainvaltest/leftImg8bit", + key="segmentation_image_dir", + ) + st.text_input( + "Label Directory", + value="gtFine", + key="segmentation_label_dir", + ) + st.text_input( + "Image Suffix", + value="_leftImg8bit.png", + key="segmentation_image_suffix", + ) + st.text_input( + "Label Suffix", + value="_gtFine_labelIds.png", + key="segmentation_label_suffix", + ) + st.checkbox( + "Use Train IDs", + value=False, + key="segmentation_use_train_id", + help="Enable when labels use _gtFine_labelTrainIds.png.", + ) + + +def render_nuimages_dataset_inputs(): + """Render the input fields for the NuImages dataset in the sidebar.""" + st.text_input( + "Version", + value="v1.0-mini", + key="nuimages_segmentation_version", + help="nuImages version, for example v1.0-mini, v1.0-train, or v1.0-val.", + ) + st.text_input( + "Generated Labels Directory", + value="generated/nuimages_segmentation_labels", + key="nuimages_segmentation_labels_dir", + help="Relative directory where generated segmentation masks are stored.", + ) + +def render_segmentation_model_inputs(): + """Render the input fields for the image segmentation model in the sidebar.""" + model_type = st.selectbox( + "Model Type", + ["Torch Model File", "Hugging Face SegFormer"], + key="segmentation_model_type", + ) + + col1, col2 = st.columns([3, 1]) + with col1: + st.text_input( + ( + "Model Path" + if model_type == "Torch Model File" + else "Model Name or Folder" + ), + key="segmentation_model_path", + help=( + "Path to a TorchScript model or saved PyTorch model file." + if model_type == "Torch Model File" + else "Hugging Face model name or local folder downloaded with save_pretrained." + ), + ) + with col2: + st.markdown( + "
", + unsafe_allow_html=True, + ) + st.button( + "Browse", + on_click=browse_segmentation_model_path, + key="browse_segmentation_model_path", + ) + + col1, col2 = st.columns([3, 1]) + with col1: + st.text_input( + "Config File", + key="segmentation_config_path", + help="JSON model configuration file.", + ) + with col2: + st.markdown( + "
", + unsafe_allow_html=True, + ) + st.button( + "Browse", + on_click=browse_segmentation_config_path, + key="browse_segmentation_config_path", + ) + + col1, col2 = st.columns([3, 1]) + with col1: + st.text_input( + "Ontology File", + key="segmentation_ontology_path", + help="JSON file containing the model output ontology.", + ) + with col2: + st.markdown( + "
", + unsafe_allow_html=True, + ) + st.button( + "Browse", + on_click=browse_segmentation_ontology_path, + key="browse_segmentation_ontology_path", + ) + + +def load_image_segmentation_model(): + """Render the image segmentation model in the sidebar.""" + from perceptionmetrics.models.torch_segmentation import TorchImageSegmentationModel + + model_type = st.session_state.get("segmentation_model_type", "Torch Model File") + model_path = st.session_state.get("segmentation_model_path", "") + config_path = st.session_state.get("segmentation_config_path", "") + ontology_path = st.session_state.get("segmentation_ontology_path", "") + + if not model_path: + st.error("Please provide a model path or model name.") + return + if not config_path or not os.path.isfile(config_path): + st.error("Please provide a valid config JSON path.") + return + if not ontology_path or not os.path.isfile(ontology_path): + st.error("Please provide a valid ontology JSON path.") + return + + with st.spinner("Loading image segmentation model..."): + try: + model = load_model_for_type(model_type, model_path) + segmentation_model = TorchImageSegmentationModel( + model=model, + model_cfg=config_path, + ontology_fname=ontology_path, + ) + st.session_state.segmentation_model = segmentation_model + st.session_state.segmentation_model_loaded = True + st.success("Segmentation model loaded and saved for inference") + except Exception as exc: + st.session_state.segmentation_model = None + st.session_state.segmentation_model_loaded = False + st.error(f"Failed to load segmentation model: {exc}") + + +def load_model_for_type(model_type, model_path): + """Load a model based on the specified type and path. + param model_type: Type of the model to load (e.g., "Torch Model File", "Hugging Face SegFormer"). + param model_path: Path to the model file or directory. + """ + if model_type == "Torch Model File": + if not os.path.isfile(model_path): + raise ValueError( + "Torch Model File expects a .pt/.pth/.torchscript file saved with " + "torch.save or torch.jit.save. For a Hugging Face model folder, " + "select 'Hugging Face SegFormer'." + ) + return model_path + + if model_type == "Hugging Face SegFormer": + try: + from transformers import SegformerForSemanticSegmentation + except ImportError as exc: + raise ImportError( + "transformers is required for Hugging Face SegFormer models." + ) from exc + return SegformerForSemanticSegmentation.from_pretrained(model_path) + + raise ValueError(f"Unsupported segmentation model type: {model_type}") diff --git a/tabs/tasks/lidar_segmentation/dataset_viewer.py b/tabs/tasks/lidar_segmentation/dataset_viewer.py new file mode 100644 index 00000000..dc1d0fc9 --- /dev/null +++ b/tabs/tasks/lidar_segmentation/dataset_viewer.py @@ -0,0 +1,2 @@ +def render_lidar_segmentation_viewer(): + return \ No newline at end of file diff --git a/tabs/tasks/lidar_segmentation/evaluator.py b/tabs/tasks/lidar_segmentation/evaluator.py new file mode 100644 index 00000000..b7972914 --- /dev/null +++ b/tabs/tasks/lidar_segmentation/evaluator.py @@ -0,0 +1,2 @@ +def render_lidar_segmentation_evaluator(): + return \ No newline at end of file diff --git a/tabs/tasks/lidar_segmentation/inference.py b/tabs/tasks/lidar_segmentation/inference.py new file mode 100644 index 00000000..4ee9aa7c --- /dev/null +++ b/tabs/tasks/lidar_segmentation/inference.py @@ -0,0 +1,2 @@ +def render_lidar_segmentation_inference(): + return \ No newline at end of file diff --git a/tabs/tasks/lidar_segmentation/sidebar.py b/tabs/tasks/lidar_segmentation/sidebar.py new file mode 100644 index 00000000..89cd20ca --- /dev/null +++ b/tabs/tasks/lidar_segmentation/sidebar.py @@ -0,0 +1,2 @@ +def render_lidar_segmentation_sidebar(_available_devices): + return \ No newline at end of file diff --git a/tabs/tasks/utils.py b/tabs/tasks/utils.py new file mode 100644 index 00000000..d14f213e --- /dev/null +++ b/tabs/tasks/utils.py @@ -0,0 +1,261 @@ +import platform +import subprocess +import sys + +import streamlit as st +from streamlit_image_select import image_select + + +def is_wsl(): + """ + Detect if running in Windows Subsystem for Linux (WSL). + Returns True if WSL is detected, False otherwise. + """ + return ( + "wsl" in platform.release().lower() or "microsoft" in platform.release().lower() + ) + + +def browse_folder(): + """ + Opens a native folder selection dialog and returns the selected folder path. + Works on Windows, macOS, and Linux (with zenity or kdialog). + Returns None if cancelled or error. + """ + try: + is_windows = sys.platform.startswith("win") + is_wsl_env = is_wsl() + if is_windows or is_wsl_env: + script = ( + "Add-Type -AssemblyName System.windows.forms;" + "$f=New-Object System.Windows.Forms.FolderBrowserDialog;" + 'if($f.ShowDialog() -eq "OK"){Write-Output $f.SelectedPath}' + ) + result = subprocess.run( + ["powershell.exe", "-NoProfile", "-Command", script], + capture_output=True, + text=True, + timeout=30, + ) + folder = result.stdout.strip() + if folder and is_wsl_env: + result = subprocess.run( + ["wslpath", "-u", folder], + capture_output=True, + text=True, + timeout=30, + ) + folder = result.stdout.strip() + return folder if folder else None + elif sys.platform == "darwin": + script = 'POSIX path of (choose folder with prompt "Select folder:")' + result = subprocess.run( + ["osascript", "-e", script], capture_output=True, text=True, timeout=30 + ) + folder = result.stdout.strip() + return folder if folder else None + else: + for cmd in [ + [ + "zenity", + "--file-selection", + "--directory", + "--title=Select folder", + ], + [ + "kdialog", + "--getexistingdirectory", + "--title", + "Select folder", + ], + ]: + try: + result = subprocess.run( + cmd, capture_output=True, text=True, timeout=30 + ) + if result.returncode == 0 or result.returncode == 1: + folder = result.stdout.strip() + return folder if folder else None + except subprocess.TimeoutExpired: + return None + except (FileNotFoundError, Exception): + continue + return None + except Exception: + return None + + +def browse_file(): + """ + Opens a native file selection dialog and returns the selected file path. + Works on Windows, macOS, and Linux (with zenity or kdialog). + Returns None if cancelled or error. + """ + try: + is_windows = sys.platform.startswith("win") + is_wsl_env = is_wsl() + if is_windows or is_wsl_env: + script = ( + "Add-Type -AssemblyName System.windows.forms;" + "$f=New-Object System.Windows.Forms.OpenFileDialog;" + 'if($f.ShowDialog() -eq "OK"){Write-Output $f.FileName}' + ) + result = subprocess.run( + ["powershell.exe", "-NoProfile", "-Command", script], + capture_output=True, + text=True, + timeout=30, + ) + file_path = result.stdout.strip() + if file_path and is_wsl_env: + result = subprocess.run( + ["wslpath", "-u", file_path], + capture_output=True, + text=True, + timeout=30, + ) + file_path = result.stdout.strip() + return file_path if file_path else None + elif sys.platform == "darwin": + script = 'POSIX path of (choose file with prompt "Select file:")' + result = subprocess.run( + ["osascript", "-e", script], capture_output=True, text=True, timeout=30 + ) + file_path = result.stdout.strip() + return file_path if file_path else None + else: + for cmd in [ + ["zenity", "--file-selection", "--title=Select file"], + ["kdialog", "--getopenfilename", "--title", "Select file"], + ]: + try: + result = subprocess.run( + cmd, capture_output=True, text=True, timeout=30 + ) + if result.returncode == 0 or result.returncode == 1: + file_path = result.stdout.strip() + return file_path if file_path else None + except subprocess.TimeoutExpired: + return None + except (FileNotFoundError, Exception): + continue + return None + except Exception: + return None + + +def render_image_grid( + item_names, + image_paths, + state_prefix, + context, + search_label="image", + images_per_page=12, +): + total_pages = (len(item_names) + images_per_page - 1) // images_per_page + page_key = f"{state_prefix}_page_{context}" + + if page_key not in st.session_state: + st.session_state[page_key] = 0 + + current_page = max(0, min(st.session_state[page_key], total_pages - 1)) + st.session_state[page_key] = current_page + + start_idx = current_page * images_per_page + page_item_names = item_names[start_idx : start_idx + images_per_page] + page_image_paths = image_paths[start_idx : start_idx + images_per_page] + + col1, col2, col3, col4 = st.columns([0.5, 9.5, 0.5, 0.5]) + with col1: + if st.button( + "⟨", + key=f"{state_prefix}_prev_page_btn", + disabled=(current_page == 0), + ): + st.session_state[page_key] = current_page - 1 + st.rerun() + with col2: + st.markdown( + f"
Page {current_page + 1} of {total_pages}
", + unsafe_allow_html=True, + ) + with col3: + if st.button( + "⟩", + key=f"{state_prefix}_next_page_btn", + disabled=(current_page >= total_pages - 1), + ): + st.session_state[page_key] = current_page + 1 + st.rerun() + with col4: + if st.button( + "🔍", + key=f"{state_prefix}_search_icon_btn", + help=f"Search for a {search_label} by name", + ): + st.session_state[f"show_{state_prefix}_search"] = True + + if st.session_state.get(f"show_{state_prefix}_search", False): + col1, col2, col3 = st.columns([4, 1, 1]) + with col1: + selected_item = st.selectbox( + f"Search {search_label}:", + options=item_names, + key=f"{state_prefix}_search_item", + ) + with col2: + st.markdown( + "
", + unsafe_allow_html=True, + ) + if st.button(f"Go to {search_label}", key=f"{state_prefix}_go_to_item"): + selected_idx = item_names.index(selected_item) + new_page = selected_idx // images_per_page + st.session_state[page_key] = new_page + st.session_state[f"{state_prefix}_select_{context}_{new_page}"] = ( + selected_idx % images_per_page + ) + st.session_state[f"show_{state_prefix}_search"] = False + st.rerun() + with col3: + st.markdown( + "
", + unsafe_allow_html=True, + ) + if st.button("Cancel", key=f"{state_prefix}_cancel_search"): + st.session_state[f"show_{state_prefix}_search"] = False + st.rerun() + + caption_len_limit = 17 + captions = [ + ( + (name[:caption_len_limit] + "..." + name[-3:]) + if len(name) > caption_len_limit + else name + ) + for name in page_item_names + ] + + select_key = f"{state_prefix}_select_{context}_{current_page}" + select_index = st.session_state.get(select_key) + if select_index is None or not isinstance(select_index, int): + select_index = 0 + + selected_image_path = ( + image_select( + label="", + images=page_image_paths, + captions=captions, + use_container_width=False, + key=select_key, + index=select_index, + ) + if page_image_paths + else None + ) + + if not selected_image_path: + return None, None + + selected_index = page_image_paths.index(selected_image_path) + return selected_image_path, page_item_names[selected_index]