File size: 7,488 Bytes
1197f7d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 |
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()
|