-
Notifications
You must be signed in to change notification settings - Fork 47
/
JigsawImageLoader.py
97 lines (79 loc) · 3 KB
/
JigsawImageLoader.py
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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
# -*- coding: utf-8 -*-
"""
Created on Fri Aug 18 11:58:07 2017
@author: Biagio Brattoli
"""
import numpy as np
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
from PIL import Image
class DataLoader(data.Dataset):
def __init__(self, data_path, txt_list, classes=1000):
self.data_path = data_path
self.names, _ = self.__dataset_info(txt_list)
self.N = len(self.names)
self.permutations = self.__retrive_permutations(classes)
self.__image_transformer = transforms.Compose([
transforms.Resize(256, Image.BILINEAR),
transforms.CenterCrop(255)])
self.__augment_tile = transforms.Compose([
transforms.RandomCrop(64),
transforms.Resize((75, 75), Image.BILINEAR),
transforms.Lambda(rgb_jittering),
transforms.ToTensor(),
# transforms.Normalize(mean=[0.485, 0.456, 0.406],
# std =[0.229, 0.224, 0.225])
])
def __getitem__(self, index):
framename = self.data_path + '/' + self.names[index]
img = Image.open(framename).convert('RGB')
if np.random.rand() < 0.30:
img = img.convert('LA').convert('RGB')
if img.size[0] != 255:
img = self.__image_transformer(img)
s = float(img.size[0]) / 3
a = s / 2
tiles = [None] * 9
for n in range(9):
i = n / 3
j = n % 3
c = [a * i * 2 + a, a * j * 2 + a]
c = np.array([c[1] - a, c[0] - a, c[1] + a + 1, c[0] + a + 1]).astype(int)
tile = img.crop(c.tolist())
tile = self.__augment_tile(tile)
# Normalize the patches indipendently to avoid low level features shortcut
m, s = tile.view(3, -1).mean(dim=1).numpy(), tile.view(3, -1).std(dim=1).numpy()
s[s == 0] = 1
norm = transforms.Normalize(mean=m.tolist(), std=s.tolist())
tile = norm(tile)
tiles[n] = tile
order = np.random.randint(len(self.permutations))
data = [tiles[self.permutations[order][t]] for t in range(9)]
data = torch.stack(data, 0)
return data, int(order), tiles
def __len__(self):
return len(self.names)
def __dataset_info(self, txt_labels):
with open(txt_labels, 'r') as f:
images_list = f.readlines()
file_names = []
labels = []
for row in images_list:
row = row.split(' ')
file_names.append(row[0])
labels.append(int(row[1]))
return file_names, labels
def __retrive_permutations(self, classes):
all_perm = np.load('permutations_%d.npy' % (classes))
# from range [1,9] to [0,8]
if all_perm.min() == 1:
all_perm = all_perm - 1
return all_perm
def rgb_jittering(im):
im = np.array(im, 'int32')
for ch in range(3):
im[:, :, ch] += np.random.randint(-2, 2)
im[im > 255] = 255
im[im < 0] = 0
return im.astype('uint8')