import cv2 import numpy as np import torch from PIL import Image from torch.utils.data import Dataset from torch.utils.data.dataset import Subset from torchvision.datasets.sbd import SBDataset from torchvision.datasets.voc import VOCSegmentation import ignite.distributed as idist from ignite.utils import convert_tensor class TransformedDataset(Dataset): def __init__(self, ds, transform_fn): assert isinstance(ds, Dataset) assert callable(transform_fn) self.ds = ds self.transform_fn = transform_fn def __len__(self): return len(self.ds) def __getitem__(self, index): dp = self.ds[index] return self.transform_fn(**dp) class VOCSegmentationOpencv(VOCSegmentation): target_names = [ "background", "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "plant", "sheep", "sofa", "train", "tv/monitor", ] def __init__(self, *args, return_meta=False, **kwargs): super(VOCSegmentationOpencv, self).__init__(*args, **kwargs) self.return_meta = return_meta def __getitem__(self, index): img = cv2.imread(self.images[index]) assert img is not None, f"Image at '{self.images[index]}' has a problem" img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask = np.asarray(Image.open(self.masks[index])) if self.return_meta: return { "image": img, "mask": mask, "meta": {"index": index, "image_path": self.images[index], "mask_path": self.masks[index]}, } return {"image": img, "mask": mask} class SBDatasetOpencv(SBDataset): def __init__(self, *args, return_meta=False, **kwargs): super(SBDatasetOpencv, self).__init__(*args, **kwargs) assert self.mode == "segmentation", "SBDatasetOpencv should be in segmentation mode only" self.return_meta = return_meta def _get_segmentation_target(self, filepath): mat = self._loadmat(filepath) return mat["GTcls"][0]["Segmentation"][0] def __getitem__(self, index): img = cv2.imread(self.images[index]) assert img is not None, f"Image at '{self.images[index]}' has a problem" img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask = self._get_target(self.masks[index]) if self.return_meta: return { "image": img, "mask": mask, "meta": {"index": index, "image_path": self.images[index], "mask_path": self.masks[index]}, } return {"image": img, "mask": mask} def get_train_dataset(root_path, return_meta=False): return VOCSegmentationOpencv( root=root_path, year="2012", image_set="train", download=False, return_meta=return_meta ) def get_val_dataset(root_path, return_meta=False): return VOCSegmentationOpencv(root=root_path, year="2012", image_set="val", download=False, return_meta=return_meta) def get_train_noval_sbdataset(root_path, return_meta=False): return SBDatasetOpencv(root_path, image_set="train_noval", mode="segmentation", return_meta=return_meta) def get_dataloader(dataset, sampler=None, shuffle=False, limit_num_samples=None, **kwargs): if limit_num_samples is not None: g = torch.Generator().manual_seed(limit_num_samples) indices = torch.randperm(len(dataset), generator=g)[:limit_num_samples] dataset = Subset(dataset, indices) return idist.auto_dataloader(dataset, sampler=sampler, shuffle=(sampler is None) and shuffle, **kwargs) def get_train_val_loaders( root_path, train_transforms, val_transforms, batch_size=16, num_workers=8, train_sampler=None, val_batch_size=None, sbd_path=None, limit_train_num_samples=None, limit_val_num_samples=None, ): train_ds = get_train_dataset(root_path) val_ds = get_val_dataset(root_path) if sbd_path is not None: sbd_train_ds = get_train_noval_sbdataset(sbd_path) train_ds = train_ds + sbd_train_ds if len(val_ds) < len(train_ds): g = torch.Generator().manual_seed(len(train_ds)) train_eval_indices = torch.randperm(len(train_ds), generator=g)[: len(val_ds)] train_eval_ds = Subset(train_ds, train_eval_indices) else: train_eval_ds = train_ds train_ds = TransformedDataset(train_ds, transform_fn=train_transforms) val_ds = TransformedDataset(val_ds, transform_fn=val_transforms) train_eval_ds = TransformedDataset(train_eval_ds, transform_fn=val_transforms) val_batch_size = batch_size * 4 if val_batch_size is None else val_batch_size train_loader = get_dataloader( train_ds, shuffle=True, sampler=train_sampler, batch_size=batch_size, num_workers=num_workers, drop_last=True, limit_num_samples=limit_train_num_samples, ) val_loader = get_dataloader( val_ds, shuffle=False, batch_size=val_batch_size, num_workers=num_workers, drop_last=False, limit_num_samples=limit_val_num_samples, ) train_eval_loader = get_dataloader( train_eval_ds, shuffle=False, batch_size=val_batch_size, num_workers=num_workers, drop_last=False, limit_num_samples=limit_val_num_samples, ) return train_loader, val_loader, train_eval_loader def get_inference_dataloader( root_path, mode, transforms, batch_size=16, num_workers=8, pin_memory=True, limit_num_samples=None ): assert mode in ("train", "test"), "Mode should be 'train' or 'test'" get_dataset_fn = get_train_dataset if mode == "train" else get_val_dataset dataset = get_dataset_fn(root_path, return_meta=True) dataset = TransformedDataset(dataset, transform_fn=transforms) return get_dataloader( dataset, limit_num_samples=limit_num_samples, shuffle=False, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, drop_last=False, ) def ignore_mask_boundaries(**kwargs): assert "mask" in kwargs, "Input should contain 'mask'" mask = kwargs["mask"] mask[mask == 255] = 0 kwargs["mask"] = mask return kwargs def denormalize(t, mean, std, max_pixel_value=255): assert isinstance(t, torch.Tensor), f"{type(t)}" assert t.ndim == 3 d = t.device mean = torch.tensor(mean, device=d).unsqueeze(-1).unsqueeze(-1) std = torch.tensor(std, device=d).unsqueeze(-1).unsqueeze(-1) tensor = std * t + mean tensor *= max_pixel_value return tensor def prepare_image_mask(batch, device, non_blocking): x, y = batch["image"], batch["mask"] x = convert_tensor(x, device, non_blocking=non_blocking) y = convert_tensor(y, device, non_blocking=non_blocking).long() return x, y