dataset/crag.py

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

class CRAGHRLRDataset(URURHRLRDataset):
    pass

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