| """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 |
|
|
|
|
| |
| |
| |
|
|
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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 = [] |
|
|
