| import bisect |
| import numpy as np |
| import albumentations |
| from PIL import Image |
| from torch.utils.data import Dataset, ConcatDataset |
|
|
|
|
| class ConcatDatasetWithIndex(ConcatDataset): |
| """Modified from original pytorch code to return dataset idx""" |
| def __getitem__(self, idx): |
| if idx < 0: |
| if -idx > len(self): |
| raise ValueError("absolute value of index should not exceed dataset length") |
| idx = len(self) + idx |
| dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx) |
| if dataset_idx == 0: |
| sample_idx = idx |
| else: |
| sample_idx = idx - self.cumulative_sizes[dataset_idx - 1] |
| return self.datasets[dataset_idx][sample_idx], dataset_idx |
|
|
|
|
| class ImagePaths(Dataset): |
| def __init__(self, paths, size=None, random_crop=False, labels=None): |
| self.size = size |
| self.random_crop = random_crop |
|
|
| self.labels = dict() if labels is None else labels |
| self.labels["file_path_"] = paths |
| self._length = len(paths) |
|
|
| if self.size is not None and self.size > 0: |
| self.rescaler = albumentations.SmallestMaxSize(max_size = self.size) |
| if not self.random_crop: |
| self.cropper = albumentations.CenterCrop(height=self.size,width=self.size) |
| else: |
| self.cropper = albumentations.RandomCrop(height=self.size,width=self.size) |
| self.preprocessor = albumentations.Compose([self.rescaler, self.cropper]) |
| else: |
| self.preprocessor = lambda **kwargs: kwargs |
|
|
| def __len__(self): |
| return self._length |
|
|
| def preprocess_image(self, image_path): |
| image = Image.open(image_path) |
| if not image.mode == "RGB": |
| image = image.convert("RGB") |
| image = np.array(image).astype(np.uint8) |
| image = self.preprocessor(image=image)["image"] |
| image = (image/127.5 - 1.0).astype(np.float32) |
| return image |
|
|
| def __getitem__(self, i): |
| example = dict() |
| example["image"] = self.preprocess_image(self.labels["file_path_"][i]) |
| for k in self.labels: |
| example[k] = self.labels[k][i] |
| return example |
|
|
|
|
| class NumpyPaths(ImagePaths): |
| def preprocess_image(self, image_path): |
| image = np.load(image_path).squeeze(0) |
| image = np.transpose(image, (1,2,0)) |
| image = Image.fromarray(image, mode="RGB") |
| image = np.array(image).astype(np.uint8) |
| image = self.preprocessor(image=image)["image"] |
| image = (image/127.5 - 1.0).astype(np.float32) |
| return image |
|
|
