"""Utility functions for dataset loading, visualization, and timing.""" import os from functools import wraps from time import perf_counter from typing import List import cv2 import numpy as np import pandas as pd import torch import albumentations as aug from torch.utils.data import Dataset # ============================================================================== # VISUALIZATION PALETTE # ============================================================================== ID2COLOR = { 0: [128, 0, 128], 1: [204, 163, 72], 2: [128, 0, 0], 3: [192, 192, 192], 4: [0, 255, 0], 5: [112, 148, 32], 6: [64, 64, 0], 7: [255, 255, 0], 8: [0, 128, 128], 9: [0, 0, 255], 10: [255, 0, 0], 11: [64, 160, 120], 12: [128, 64, 128], 13: [240, 120, 120], 14: [128, 128, 64], 255: [0, 0, 0], } def mask2label(mask, palette): """Convert a class-index mask to a colorized RGB image.""" if mask.ndim == 3 and mask.shape[0] == 1: mask = mask.squeeze(0) if isinstance(mask, torch.Tensor): mask = mask.detach().cpu().numpy() color_seg = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8) for label, color in palette.items(): color_seg[mask == label, :] = color return color_seg # ============================================================================== # DATASET CLASSES # ============================================================================== class InferenceDataset(Dataset): """Dataset for RGB inference with class remapping and dropping.""" def __init__( self, _root_dir: str, config: dict, _ignore_index: int = 255, _original_num_classes: int = 22, ): """ Args: _root_dir (str): Path to the root directory. Walked recursively; images are collected from any subfolder named `src/` and masks from any subfolder named `gt/`. No specific top-level structure is required. config (dict): Dictionary containing the configuration. _ignore_index (int): Index used for dropped classes. _original_num_classes (int): Number of original classes. """ self.root_dir = _root_dir self.transforms = aug.Compose( [aug.Resize(config["image_size"], config["image_size"], interpolation=cv2.INTER_NEAREST)] ) self.ignore_index = _ignore_index self.mapping_lut = np.full(256, self.ignore_index, dtype=np.uint8) df = pd.read_csv(config["class_dict_path"]) new_class_id = 0 raw_to_new_dict = {} for _, row in df.iterrows(): raw_id = int(row['id']) if raw_id == self.ignore_index: continue self.mapping_lut[raw_id] = new_class_id raw_to_new_dict[raw_id] = new_class_id new_class_id += 1 if 9 in raw_to_new_dict: target_id_for_merges = raw_to_new_dict[9] self.mapping_lut[5] = target_id_for_merges self.mapping_lut[22] = target_id_for_merges self.num_new_classes = new_class_id image_file_names = [] annotation_file_names = [] for root, _, files in os.walk(self.root_dir): for name in files: if ".ipynb_checkpoints" in root: continue if "src" in root and name.lower().endswith((".png", ".jpeg", ".jpg")): image_file_names.append(os.path.join(root, name)) elif "gt" in root and name.lower().endswith(".png"): if "vis_rem" not in os.path.join(root, name): annotation_file_names.append(os.path.join(root, name)) self.images = sorted(image_file_names) self.annotations = sorted(annotation_file_names) assert len(self.images) == len(self.annotations), ( f"There must be as many images as there are segmentation maps. " f"Image = {len(self.images)}, Masks = {len(self.annotations)}" ) def __len__(self): return len(self.images) def __getitem__(self, idx): image = cv2.imread(self.images[idx]) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) segmentation_map = cv2.imread(self.annotations[idx], cv2.IMREAD_GRAYSCALE) segmentation_map = self.mapping_lut[segmentation_map] if self.transforms is not None: augmented = self.transforms(image=image, mask=segmentation_map) return augmented["image"], augmented["mask"] return image, segmentation_map class InferenceDatasetThermal(Dataset): """Dataset for thermal inference with class remapping and dropping.""" def __init__( self, _root_dir: str, config: dict, _ignore_index: int = 255, ): """ Args: _root_dir (str): Path to the root directory. Walked recursively; images are collected from any subfolder named `thermal_src/` and masks from any subfolder named `gt/`. No specific top-level structure is required. config (dict): Dictionary containing the configuration. _ignore_index (int): Index used for dropped classes. """ self.root_dir = _root_dir self.transforms = aug.Compose( [aug.Resize(config["image_size"], config["image_size"], interpolation=cv2.INTER_NEAREST)] ) self.ignore_index = _ignore_index self.mapping_lut = np.full(256, self.ignore_index, dtype=np.uint8) df = pd.read_csv(config["class_dict_path"]) new_class_id = 0 raw_to_new_dict = {} for _, row in df.iterrows(): raw_id = int(row['id']) self.mapping_lut[raw_id] = new_class_id raw_to_new_dict[raw_id] = new_class_id new_class_id += 1 if 9 in raw_to_new_dict: target_id_for_merges = raw_to_new_dict[9] self.mapping_lut[5] = target_id_for_merges self.mapping_lut[22] = target_id_for_merges self.num_new_classes = new_class_id image_file_names = [] annotation_file_names = [] for root, _, files in os.walk(self.root_dir): for name in files: if ".ipynb_checkpoints" in root: continue if "thermal_src" in root and name.lower().endswith((".png", ".jpeg", ".jpg")): image_file_names.append(os.path.join(root, name)) elif "gt" in root and name.lower().endswith(".png"): annotation_file_names.append(os.path.join(root, name)) self.images = sorted(image_file_names) self.annotations = sorted(annotation_file_names) assert len(self.images) == len(self.annotations), ( f"There must be as many images as there are segmentation maps." f"Image = {len(self.images)}, Masks = {len(self.annotations)}" ) def __len__(self): return len(self.images) def __getitem__(self, idx): image = cv2.imread(self.images[idx]) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) segmentation_map = cv2.imread(self.annotations[idx], cv2.IMREAD_GRAYSCALE) segmentation_map = self.mapping_lut[segmentation_map] if self.transforms is not None: augmented = self.transforms(image=image, mask=segmentation_map) return augmented["image"], augmented["mask"] return image, segmentation_map # ============================================================================== # TIMING UTILITY # ============================================================================== class Timing: """Class to time functions and methods.""" def __init__(self): self.inf_time = [] def __call__(self, func): @wraps(func) def wrapper(*args, **kwargs): start = perf_counter() result = func(*args, **kwargs) end = perf_counter() self.inf_time.append(end - start) return result return wrapper def print_average(self): """Print the average timing of the function.""" print(f"---\tFunction took {np.mean(self.inf_time):.4f} seconds to run!") def get_average(self) -> float: """Return the average timing of the function.""" return float(np.mean(self.inf_time)) def reset(self): """Reset the timing of the function.""" self.inf_time = []