def main(dataset_path, workers):
transform = T.Compose([
ApplyTo(
['image'],
T.Compose([
SplitInSites(),
T.Lambda(
lambda xs: torch.stack([ToTensor()(x) for x in xs], 0)),
])),
Extract(['image']),
])
train_data = pd.read_csv(os.path.join(dataset_path, 'train.csv'))
train_data['root'] = os.path.join(dataset_path, 'train')
test_data = pd.read_csv(os.path.join(dataset_path, 'test.csv'))
test_data['root'] = os.path.join(dataset_path, 'test')
data = pd.concat([train_data, test_data])
stats = {}
for (exp, plate), group in tqdm(data.groupby(['experiment', 'plate'])):
dataset = TestDataset(group, transform=transform)
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=32,
num_workers=workers)
with torch.no_grad():
images = [images for images, in data_loader]
images = torch.cat(images, 0)
mean = images.mean((0, 1, 3, 4))
std = images.std((0, 1, 3, 4))
stats[(exp, plate)] = mean, std
del images, mean, std
gc.collect()
torch.save(stats, 'plate_stats.pth')
def dist_sharpen(dist, temp):
assert dist.dim() == 2
dist = dist**(1 / temp)
dist = dist / dist.sum(1, keepdim=True)
return dist
random_crop = Resettable(RandomCrop)
center_crop = Resettable(CenterCrop)
infer_image_transform = Resettable(lambda tta: test_image_transform
if tta else eval_image_transform)
to_tensor = ToTensor()
if config.normalize is None:
normalize = T.Compose([])
elif config.normalize == 'experiment':
normalize = NormalizeByExperimentStats(
torch.load('./experiment_stats.pth'))
elif config.normalize == 'plate':
normalize = NormalizeByPlateStats(torch.load('./plate_stats.pth'))
else:
raise AssertionError('invalid normalization {}'.format(config.normalize))
eval_image_transform = T.Compose([
RandomSite(),
Resize(config.resize_size),
center_crop,
random_crop = Resetable(RandomCrop)
center_crop = Resetable(CenterCrop)
train_transform = T.Compose([
ApplyTo(
['image'],
T.Compose([
RandomSite(),
Resize(config.resize_size),
random_crop,
RandomFlip(),
RandomTranspose(),
RandomRotation(180), # FIXME:
ToTensor(),
ChannelReweight(config.aug.channel_weight),
])),
# NormalizeByRefStats(),
Extract(['image', 'feat', 'label', 'id']),
])
eval_transform = T.Compose([
ApplyTo(
['image'],
T.Compose([
RandomSite(), # FIXME:
Resize(config.resize_size),
center_crop,
ToTensor(),
])),
# NormalizeByRefStats(),