baby-yoda-segmentor/src/dataset.py

import torch
import torchvision
import transforms as T
from coco_utils import preview_img_and_target

torch.manual_seed(42)


# Adapt to maskrcnn
def get_transforms(train, preview=False):
    t = [
        T.FilterAndRemapCocoCategories(categories=[0, 3]),
        T.ConvertCocoPolysToMask(),
        T.ToTensor(),
    ]
    if train:
        t += [
            T.RandomHorizontalFlip(.5),
            T.RandomVerticalFlip(.5),
            T.ColorJitter((.8, 1.8), (.8, 1.8), .3, .1),
        ]

    if preview:
        t.append(T.PreviewTransforms())
    transforms = T.Compose(t)

    return transforms


class CocoDetection(torchvision.datasets.CocoDetection):
    def __init__(self, img_folder, ann_file, transforms):
        super(CocoDetection, self).__init__(img_folder, ann_file)
        self._transforms = transforms

    def __getitem__(self, idx):
        img, target = super(CocoDetection, self).__getitem__(idx)
        image_id = self.ids[idx]
        target = dict(image_id=image_id, annotations=target)
        if self._transforms is not None:
            img, target = self._transforms(img, target)
        return img, target