import os from os import path from typing import List, Tuple, Union import diskcache as dc import hydra import numpy as np import torch from loguru import logger from PIL import Image from torch.utils.data import DataLoader, Dataset from torchvision.transforms import functional as TF from tqdm.rich import tqdm from yolo.tools.dataset_helper import ( create_image_info_dict, find_labels_path, get_scaled_segmentation, ) from yolo.utils.data_augment import Compose, HorizontalFlip, MixUp, Mosaic, VerticalFlip from yolo.utils.drawer import draw_bboxes class YoloDataset(Dataset): def __init__(self, config: dict, phase: str = "train2017", image_size: int = 640): dataset_cfg = config.data augment_cfg = config.augmentation phase_name = dataset_cfg.get(phase, phase) self.image_size = image_size transforms = [eval(aug)(prob) for aug, prob in augment_cfg.items()] self.transform = Compose(transforms, self.image_size) self.transform.get_more_data = self.get_more_data self.data = self.load_data(dataset_cfg.path, phase_name) def load_data(self, dataset_path, phase_name): """ Loads data from a cache or generates a new cache for a specific dataset phase. Parameters: dataset_path (str): The root path to the dataset directory. phase_name (str): The specific phase of the dataset (e.g., 'train', 'test') to load or generate data for. Returns: dict: The loaded data from the cache for the specified phase. """ cache_path = path.join(dataset_path, ".cache") cache = dc.Cache(cache_path) data = cache.get(phase_name) if data is None: logger.info("Generating {} cache", phase_name) data = self.filter_data(dataset_path, phase_name) cache[phase_name] = data cache.close() logger.info("📦 Loaded {} cache", phase_name) data = cache[phase_name] return data def filter_data(self, dataset_path: str, phase_name: str) -> list: """ Filters and collects dataset information by pairing images with their corresponding labels. Parameters: images_path (str): Path to the directory containing image files. labels_path (str): Path to the directory containing label files. Returns: list: A list of tuples, each containing the path to an image file and its associated segmentation as a tensor. """ images_path = path.join(dataset_path, "images", phase_name) labels_path, data_type = find_labels_path(dataset_path, phase_name) images_list = sorted(os.listdir(images_path)) if data_type == "json": annotations_index, image_info_dict = create_image_info_dict(labels_path) data = [] valid_inputs = 0 for image_name in tqdm(images_list, desc="Filtering data"): if not image_name.lower().endswith((".jpg", ".jpeg", ".png")): continue image_id, _ = path.splitext(image_name) if data_type == "json": image_info = image_info_dict.get(image_id, None) if image_info is None: continue annotations = annotations_index.get(image_info["id"], []) image_seg_annotations = get_scaled_segmentation(annotations, image_info) if not image_seg_annotations: continue elif data_type == "txt": label_path = path.join(labels_path, f"{image_id}.txt") if not path.isfile(label_path): continue with open(label_path, "r") as file: image_seg_annotations = [list(map(float, line.strip().split())) for line in file] labels = self.load_valid_labels(image_id, image_seg_annotations) if labels is not None: img_path = path.join(images_path, image_name) data.append((img_path, labels)) valid_inputs += 1 logger.info("Recorded {}/{} valid inputs", valid_inputs, len(images_list)) return data def load_valid_labels(self, label_path, seg_data_one_img) -> Union[torch.Tensor, None]: """ Loads and validates bounding box data is [0, 1] from a label file. Parameters: label_path (str): The filepath to the label file containing bounding box data. Returns: torch.Tensor or None: A tensor of all valid bounding boxes if any are found; otherwise, None. """ bboxes = [] for seg_data in seg_data_one_img: cls = seg_data[0] points = np.array(seg_data[1:]).reshape(-1, 2) valid_points = points[(points >= 0) & (points <= 1)].reshape(-1, 2) if valid_points.size > 1: bbox = torch.tensor([cls, *valid_points.min(axis=0), *valid_points.max(axis=0)]) bboxes.append(bbox) if bboxes: return torch.stack(bboxes) else: logger.warning("No valid BBox in {}", label_path) return None def get_data(self, idx): img_path, bboxes = self.data[idx] img = Image.open(img_path).convert("RGB") return img, bboxes def get_more_data(self, num: int = 1): indices = torch.randint(0, len(self), (num,)) return [self.get_data(idx) for idx in indices] def __getitem__(self, idx) -> Union[Image.Image, torch.Tensor]: img, bboxes = self.get_data(idx) if self.transform: img, bboxes = self.transform(img, bboxes) img = TF.to_tensor(img) return img, bboxes def __len__(self) -> int: return len(self.data) class YoloDataLoader(DataLoader): def __init__(self, config: dict): """Initializes the YoloDataLoader with hydra-config files.""" hyper = config.hyper.data dataset = YoloDataset(config) super().__init__( dataset, batch_size=hyper.batch_size, shuffle=hyper.shuffle, num_workers=hyper.num_workers, pin_memory=hyper.pin_memory, collate_fn=self.collate_fn, ) def collate_fn(self, batch: List[Tuple[torch.Tensor, torch.Tensor]]) -> Tuple[torch.Tensor, List[torch.Tensor]]: """ A collate function to handle batching of images and their corresponding targets. Args: batch (list of tuples): Each tuple contains: - image (torch.Tensor): The image tensor. - labels (torch.Tensor): The tensor of labels for the image. Returns: Tuple[torch.Tensor, List[torch.Tensor]]: A tuple containing: - A tensor of batched images. - A list of tensors, each corresponding to bboxes for each image in the batch. """ images = torch.stack([item[0] for item in batch]) targets = [item[1] for item in batch] return images, targets def get_dataloader(config): return YoloDataLoader(config) @hydra.main(config_path="../config", config_name="config", version_base=None) def main(cfg): dataloader = get_dataloader(cfg) draw_bboxes(*next(iter(dataloader))) if __name__ == "__main__": import sys sys.path.append("./") from tools.log_helper import custom_logger custom_logger() main()