I don’t have a public repo for this, but I’m happy to share those classes. I’m using the Unet model directly from this repo without changing it: https://github.com/milesial/Pytorch-UNet/blob/master/unet/unet_parts.py 4 and https://github.com/milesial/Pytorch-UNet/blob/master/unet/unet_model.py 2
My dataset class:
class medImDataset(Dataset):
    def __init__(self, image_path, mask_path, dist_path, contrast = True, hflip=True, train=True):
        self.image_path = image_path
        self.mask_path = mask_path
        self.dist_path = dist_path
        self.image_files = os.listdir(self.image_path)
        self.mask_files = os.listdir(self.mask_path)
        self.dist_files = os.listdir(self.dist_path)
        self.contrast = contrast
        self.hflip = hflip

    def transform(self, image, mask, dist_map, contrast, hflip):
        resize = transforms.Resize(size=(img_height,img_width))
        image = resize(image)
        mask = resize(mask)
        dist_map = resize(dist_map)

        image = TF.to_grayscale(image)
        mask = TF.to_grayscale(mask)
        dist_map = TF.to_grayscale(dist_map)

        if hflip is True:
            if random.random() > 0.5:
                image = TF.hflip(image)
                mask = TF.hflip(mask)
                dist_map = TF.hflip(dist_map)

        if contrast is True:
            if random.random() > 0.5:
                image = TF.adjust_contrast(image, contrast_factor=((random.random() * 1.5) + 0.25))

        image = TF.to_tensor(image)
        mask = TF.to_tensor(mask)
        dist_map = TF.to_tensor(dist_map)

        image = TF.normalize(image, mean=[0], std=[1])

        return image, mask, dist_map

    def __getitem__(self, index):
        image_name = self.image_files[index]
        mask_name = self.mask_files[index]
        dist_name = self.dist_files[index]
        image = Image.open(os.path.join(self.image_path, image_name))
        mask = Image.open(os.path.join(self.mask_path, mask_name))
        dist_map = Image.open(os.path.join(self.dist_path, dist_name))
        x, y, z = self.transform(image, mask, dist_map, contrast=self.contrast, hflip=self.hflip)
        return x, y, z

    def __len__(self):
        return len(self.image_files)

My datasets, dataloaders for training and validation:
train_dataset = medImDataset(image_path=train_img_path, mask_path=train_mask_path, dist_path=train_dist_path)

batch_size = 8

train_dataset_loader = DataLoader(
    train_dataset,
    batch_size=batch_size, shuffle=True,
    num_workers=8
)

val_dataset = medImDataset(image_path=val_img_path, mask_path=val_mask_path, dist_path=val_dist_path, contrast=False, hflip=False)

val_batch_size = 8 #I vary this between 1-8, which gives me different validation dice scores

val_dataset_loader = DataLoader(
    us_val_dataset,
    batch_size=val_batch_size, shuffle=False,
    num_workers=8
)