def test_dataloader(self):
augmentations = Compose([
A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
ToTensorV2(),
])
test_ds = MelanomaDataset(
df=self.test_df,
images_path=self.test_images_path,
augmentations=augmentations, # TODO: add TTA
train_or_valid=False,
)
return DataLoader(
test_ds,
batch_size=self.hparams.bs,
shuffle=False,
num_workers=os.cpu_count(),
pin_memory=True,
)
def train_dataloader(self):
augmentations = Compose(
[
A.RandomResizedCrop(
height=self.hparams.sz,
width=self.hparams.sz,
scale=(0.7, 1.0),
),
# AdvancedHairAugmentation(),
A.GridDistortion(),
A.RandomBrightnessContrast(),
A.ShiftScaleRotate(),
A.Flip(p=0.5),
A.CoarseDropout(
max_height=int(self.hparams.sz / 10),
max_width=int(self.hparams.sz / 10),
),
# A.HueSaturationValue(),
A.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
max_pixel_value=255,
),
ToTensorV2(),
]
)
train_ds = MelanomaDataset(
df=self.train_df,
images_path=self.train_images_path,
augmentations=augmentations,
train_or_valid=True,
)
return DataLoader(
train_ds,
# sampler=sampler,
batch_size=self.hparams.bs,
shuffle=True,
num_workers=os.cpu_count(),
pin_memory=True,
)