|
|
r""" Dataloader builder for few-shot semantic segmentation dataset """ |
|
|
from torchvision import transforms |
|
|
from torch.utils.data import DataLoader |
|
|
|
|
|
from fewshot_data.data.pascal import DatasetPASCAL |
|
|
from fewshot_data.data.coco import DatasetCOCO |
|
|
from fewshot_data.data.fss import DatasetFSS |
|
|
|
|
|
|
|
|
class FSSDataset: |
|
|
@classmethod |
|
|
def initialize(cls, img_size, datapath, use_original_imgsize, imagenet_norm=False): |
|
|
cls.datasets = { |
|
|
'pascal': DatasetPASCAL, |
|
|
'coco': DatasetCOCO, |
|
|
'fss': DatasetFSS, |
|
|
} |
|
|
|
|
|
if imagenet_norm: |
|
|
cls.img_mean = [0.485, 0.456, 0.406] |
|
|
cls.img_std = [0.229, 0.224, 0.225] |
|
|
print('use norm: {}, {}'.format(cls.img_mean, cls.img_std)) |
|
|
else: |
|
|
cls.img_mean = [0.5] * 3 |
|
|
cls.img_std = [0.5] * 3 |
|
|
print('use norm: {}, {}'.format(cls.img_mean, cls.img_std)) |
|
|
|
|
|
cls.datapath = datapath |
|
|
cls.use_original_imgsize = use_original_imgsize |
|
|
|
|
|
cls.transform = transforms.Compose([transforms.Resize(size=(img_size, img_size)), |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize(cls.img_mean, cls.img_std)]) |
|
|
|
|
|
@classmethod |
|
|
def build_dataloader(cls, benchmark, bsz, nworker, fold, split, shot=1): |
|
|
shuffle = split == 'trn' |
|
|
nworker = nworker if split == 'trn' else 0 |
|
|
dataset = cls.datasets[benchmark](cls.datapath, fold=fold, transform=cls.transform, split=split, shot=shot, use_original_imgsize=cls.use_original_imgsize) |
|
|
dataloader = DataLoader(dataset, batch_size=bsz, shuffle=shuffle, num_workers=nworker) |
|
|
|
|
|
return dataloader |
|
|
|
