File size: 779 Bytes
482ab8a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
from typing import Dict
import albumentations as A
from .dataset import ImageDataset, crop_to_smallest_collate_fn
def get_dataset(datalist: Dict, subset, transform, opt):
datasets = {}
for k, v in datalist.items():
# val_transform = transform
if k in ["imd2020", "nist16"]:
val_transform = A.Compose([A.SmallestMaxSize(opt.tile_size)])
else:
val_transform = transform
datasets[k] = ImageDataset(
k,
v,
subset,
val_transform,
opt.uncorrect_label,
opt.mvc_spixel
if subset == "train"
else opt.crf_postproc or opt.convcrf_postproc or opt.spixel_postproc,
opt.mvc_num_spixel,
)
return datasets
|