dataset/isic.py

from pathlib import Path
from typing import Optional
from torchvision import transforms
from dataset.definition_dataset import URURHRLRDataset

class ISICHRLRDataset(URURHRLRDataset):
    def _postprocess_masks(self, msk_hr, msk_lr):
        msk_hr[msk_hr == 255] = 1
        msk_lr[msk_lr == 255] = 1
        return msk_hr, msk_lr

def build(image_dir: Path, mask_dir: Path, num_classes: int, mode: str, ignore_index: int,
          transform: Optional[transforms.Compose], augment=None):
    return ISICHRLRDataset(
        image_dir=image_dir,
        mask_dir=mask_dir,
        num_classes=num_classes,
        mode=mode,
        ignore_index=ignore_index,
        transform=transform,
        augment=augment,
    )