""" Data loading utilities for IVUS annotation tool. Handles case discovery, image loading, and ground truth labels. """ import os import glob from typing import List, Dict, Tuple, Union from pathlib import Path from .excel import Excel def get_all_case_ids(data_root: str) -> List[Union[int, float]]: """ Scan data directory and extract case numbers from folder names. Handles both integer case IDs (134) and decimal case IDs (134.1, 134.2). Args: data_root: Root directory containing case folders Returns: List of case numbers (int or float) sorted in ascending order Example: >>> get_all_case_ids("/mnt/d/Research/data/2025_10_30_updated") [134, 134.1, 134.2, 135, 135.1, ..., 795] """ if not os.path.exists(data_root): return [] try: case_dirs = [ d for d in os.listdir(data_root) if os.path.isdir(os.path.join(data_root, d)) and d.startswith("CHIBAMI_") and d.endswith("_pre") ] case_ids = [] for case_dir in case_dirs: try: # Extract number from "CHIBAMI_134_pre" or "CHIBAMI_134.1_pre" # Split by underscore: ["CHIBAMI", "134", "pre"] or ["CHIBAMI", "134.1", "pre"] parts = case_dir.split("_") if len(parts) >= 3: # Parse as float to preserve decimals, then convert to int if whole number case_num_float = float(parts[1]) # If it's a whole number, store as int, otherwise as float if case_num_float.is_integer(): case_ids.append(int(case_num_float)) else: case_ids.append(case_num_float) except (IndexError, ValueError): # Skip folders that don't match expected pattern continue return sorted(case_ids) except Exception as e: print(f"Error scanning directory {data_root}: {e}") return [] def get_case_images(data_root: str, case_id: Union[int, float]) -> List[str]: """ Get all image paths for a given case, sorted by frame index. Args: data_root: Root directory containing case folders case_id: Case number (int or float, e.g., 134 or 134.1) Returns: List of absolute paths to PNG files, sorted by frame index Example: >>> get_case_images("/mnt/d/Research/data/2025_10_30_updated", 134) ['/mnt/d/.../CHIBAMI_134_pre/images/frame_134_4440.png', ...] >>> get_case_images("/mnt/d/Research/data/2025_10_30_updated", 134.1) ['/mnt/d/.../CHIBAMI_134.1_pre/images/frame_134_4440.png', ...] """ # Construct folder name based on case_id type if isinstance(case_id, float) and not case_id.is_integer(): # Decimal case: CHIBAMI_134.1_pre case_dir = os.path.join(data_root, f"CHIBAMI_{case_id}_pre") else: # Integer case: CHIBAMI_134_pre case_dir = os.path.join(data_root, f"CHIBAMI_{int(case_id)}_pre") if not os.path.exists(case_dir): return [] images_dir = os.path.join(case_dir, "images") if not os.path.exists(images_dir): return [] # Get all PNG files matching the frame pattern image_files = glob.glob(os.path.join(images_dir, "frame_*.png")) # Sort by frame index (extract from filename) def extract_frame_idx(filepath): """Extract frame index from filename like 'frame_134_4440.png'.""" filename = os.path.basename(filepath) try: # Remove extension: "frame_134_4440.png" -> "frame_134_4440" rest_name = filename.replace(".png", "") # Split: ["frame", "134", "4440"] parts = rest_name.split("_") # Return the last part as integer return int(parts[-1]) except (IndexError, ValueError): return 0 return sorted(image_files, key=extract_frame_idx) def load_ground_truth_labels(excel_path: str) -> Dict[Union[int, float], bool]: """ Load ground truth labels from Excel file. Args: excel_path: Path to Excel file containing ground truth labels Returns: Dictionary mapping case_id to has_complication boolean - True: Complication present (No/Slow flow) - False: No complication Example: >>> load_ground_truth_labels("/path/to/labels.xlsx") {134: True, 134.1: False, 135: False, ...} """ try: excel = Excel(excel_path) return excel.extract_data() except FileNotFoundError: print(f"Warning: Excel file not found: {excel_path}") return {} except Exception as e: print(f"Warning: Could not load Excel file: {e}") return {} def validate_case_directory(data_root: str, case_id: Union[int, float]) -> Tuple[bool, str]: """ Validate that a case has accessible images. Args: data_root: Root directory containing case folders case_id: Case number to validate (int or float) Returns: Tuple of (is_valid, error_message) - If valid: (True, "") - If invalid: (False, "Error description") Example: >>> validate_case_directory("/mnt/d/Research/data/2025_10_30_updated", 134) (True, "") >>> validate_case_directory("/mnt/d/Research/data/2025_10_30_updated", 999) (False, "Case directory not found: ...") """ # Construct folder name based on case_id type if isinstance(case_id, float) and not case_id.is_integer(): case_dir = os.path.join(data_root, f"CHIBAMI_{case_id}_pre") case_name = f"CHIBAMI_{case_id}_pre" else: case_dir = os.path.join(data_root, f"CHIBAMI_{int(case_id)}_pre") case_name = f"CHIBAMI_{int(case_id)}_pre" if not os.path.exists(case_dir): return False, f"Case directory not found: {case_name}" if not os.path.isdir(case_dir): return False, f"Path is not a directory: {case_dir}" images_dir = os.path.join(case_dir, "images") if not os.path.exists(images_dir): return False, f"Images directory not found: {images_dir}" if not os.path.isdir(images_dir): return False, f"Images path is not a directory: {images_dir}" images = get_case_images(data_root, case_id) if len(images) == 0: return False, f"No PNG images found in: {images_dir}" return True, ""