| import os |
| import numpy as np |
| import cv2 |
| import albumentations |
| from PIL import Image |
| from torch.utils.data import Dataset |
|
|
|
|
| class SegmentationBase(Dataset): |
| def __init__(self, |
| data_csv, data_root, segmentation_root, |
| size=None, random_crop=False, interpolation="bicubic", |
| n_labels=182, shift_segmentation=False, |
| ): |
| self.n_labels = n_labels |
| self.shift_segmentation = shift_segmentation |
| self.data_csv = data_csv |
| self.data_root = data_root |
| self.segmentation_root = segmentation_root |
| with open(self.data_csv, "r") as f: |
| self.image_paths = f.read().splitlines() |
| self._length = len(self.image_paths) |
| self.labels = { |
| "relative_file_path_": [l for l in self.image_paths], |
| "file_path_": [os.path.join(self.data_root, l) |
| for l in self.image_paths], |
| "segmentation_path_": [os.path.join(self.segmentation_root, l.replace(".jpg", ".png")) |
| for l in self.image_paths] |
| } |
|
|
| size = None if size is not None and size<=0 else size |
| self.size = size |
| if self.size is not None: |
| self.interpolation = interpolation |
| self.interpolation = { |
| "nearest": cv2.INTER_NEAREST, |
| "bilinear": cv2.INTER_LINEAR, |
| "bicubic": cv2.INTER_CUBIC, |
| "area": cv2.INTER_AREA, |
| "lanczos": cv2.INTER_LANCZOS4}[self.interpolation] |
| self.image_rescaler = albumentations.SmallestMaxSize(max_size=self.size, |
| interpolation=self.interpolation) |
| self.segmentation_rescaler = albumentations.SmallestMaxSize(max_size=self.size, |
| interpolation=cv2.INTER_NEAREST) |
| self.center_crop = not random_crop |
| if self.center_crop: |
| self.cropper = albumentations.CenterCrop(height=self.size, width=self.size) |
| else: |
| self.cropper = albumentations.RandomCrop(height=self.size, width=self.size) |
| self.preprocessor = self.cropper |
|
|
| def __len__(self): |
| return self._length |
|
|
| def __getitem__(self, i): |
| example = dict((k, self.labels[k][i]) for k in self.labels) |
| image = Image.open(example["file_path_"]) |
| if not image.mode == "RGB": |
| image = image.convert("RGB") |
| image = np.array(image).astype(np.uint8) |
| if self.size is not None: |
| image = self.image_rescaler(image=image)["image"] |
| segmentation = Image.open(example["segmentation_path_"]) |
| assert segmentation.mode == "L", segmentation.mode |
| segmentation = np.array(segmentation).astype(np.uint8) |
| if self.shift_segmentation: |
| |
| segmentation = segmentation+1 |
| if self.size is not None: |
| segmentation = self.segmentation_rescaler(image=segmentation)["image"] |
| if self.size is not None: |
| processed = self.preprocessor(image=image, |
| mask=segmentation |
| ) |
| else: |
| processed = {"image": image, |
| "mask": segmentation |
| } |
| example["image"] = (processed["image"]/127.5 - 1.0).astype(np.float32) |
| segmentation = processed["mask"] |
| onehot = np.eye(self.n_labels)[segmentation] |
| example["segmentation"] = onehot |
| return example |
|
|
|
|
| class Examples(SegmentationBase): |
| def __init__(self, size=None, random_crop=False, interpolation="bicubic"): |
| super().__init__(data_csv="data/sflckr_examples.txt", |
| data_root="data/sflckr_images", |
| segmentation_root="data/sflckr_segmentations", |
| size=size, random_crop=random_crop, interpolation=interpolation) |
|
|
