dataset/swissimage_inference.py

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

def build(image_dir: Path, num_classes: int, mode: str, ignore_index: int, transform: Optional[transforms.Compose],
          hr_size: int = 1024, lr_side: int = 2048, stride: Optional[int] = None):
    return SWISSIMAGEINFERENCEDataset(
        image_dir=image_dir,
        num_classes=num_classes,
        mode=mode,
        ignore_index=ignore_index,
        hr_size=hr_size,
        lr_side=lr_side,
        transform=transform,
        stride=stride,
    )