class ImageResize(object):
"""Resize image and targets to specified size.
"""
# --------------------------------------------------------------------------
def __init__(self, output_size):
assert isinstance(output_size, (int, tuple))
if isinstance(output_size, int):
self.output_size = (output_size, output_size)
else:
assert len(output_size) == 2
self.output_size = output_size
self.tform = Resize(self.output_size)
# --------------------------------------------------------------------------
def __call__(self, inputs):
"""
Random contextual image crop to desired image size.
"""
image = inputs['image']
H, W = image.shape[:2]
rois = inputs['ROI']
image = self.tform.__call__(Image.fromarray(image))
image = np.array(image)
# Resize ROIs:
scale = self.output_size / np.array([H, W], dtype=np.float)
scale = np.tile(scale, 2)
rois[:, :4] = scale * rois[:, :4]
rois = np.round(rois)
return {'image': image, 'ROI': rois}
class SiameseNetworkDataset:
def __init__(self, imageFile, inputSize, minibatchSize):
self.inputSize = inputSize
self.imageFile = imageFile
self.minibatchSize = minibatchSize
#self.resize = Resize((self.inputSize, self.inputSize))
self.resize = Resize(self.inputSize)
self.randomcrop = RandomCrop(self.inputSize)
self.totensor = ToTensor()
self.transforms = Compose([
ToTensor(),
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
self.IOErrorLog = open(
'/data/srajpal2/AmazonDataset/TrainingHistory/IOErrorLog.txt', 'w')
self.RuntimeErrorLog = open(
'/data/srajpal2/AmazonDataset/TrainingHistory/RuntimeErrorLog.txt',
'w')
def __getitem__(self):
with open(self.imageFile) as f:
self.compatibility_buffer = {
x: [[], []]
for x in ['bt', 'tb', 'sb', 'bs', 'st', 'ts']
}
self.compatibility_labels = {
x: []
for x in ['bt', 'tb', 'sb', 'bs', 'st', 'ts']
}
self.similarity_buffer = [[], []]
self.similarity_labels = []
self.compatibility_lines = {
x: []
for x in ['bt', 'tb', 'sb', 'bs', 'st', 'ts']
}
self.similarity_lines = []
line_num = 0
for line in f:
line_num += 1
eg_info = line.rstrip().split()
try:
im1 = Image.open(eg_info[1])
except IOError:
self.IOErrorLog.write(str(line_num) + '\n')
continue
try:
im2 = Image.open(eg_info[2])
except:
self.IOErrorLog.write(str(line_num) + '\n')
continue
im1 = self.randomcrop.__call__(self.resize.__call__(im1))
im2 = self.randomcrop.__call__(self.resize.__call__(im2))
objective = eg_info[0]
if objective == 'C':
category_pair = eg_info[4]
self.compatibility_buffer[category_pair][0].append(im1)
self.compatibility_buffer[category_pair][1].append(im2)
self.compatibility_labels[category_pair].append(
int(eg_info[3]))
self.compatibility_lines[category_pair].append(line_num)
if len(self.compatibility_labels[category_pair]
) == self.minibatchSize:
eg = self.yieldFunction(objective,
category_pair=category_pair)
if eg == None:
continue
yield eg
else:
self.similarity_buffer[0].append(im1)
self.similarity_buffer[1].append(im2)
self.similarity_labels.append(int(eg_info[3]))
self.similarity_lines.append(line_num)
if len(self.similarity_labels) == self.minibatchSize:
eg = self.yieldFunction(objective)
if eg == None:
continue
yield eg
for category_pair in self.compatibility_buffer.keys():
if len(self.compatibility_labels[category_pair]) == 0:
continue
eg = self.yieldFunction('C', category_pair=category_pair)
yield eg
if len(self.similarity_labels) > 0:
eg = self.yieldFunction('S')
yield eg
def yieldFunction(self, objective, category_pair=None):
if objective == 'C':
try:
im1Tensor = stack([
self.transforms.__call__(x)
for x in self.compatibility_buffer[category_pair][0]
], 0)
im2Tensor = stack([
self.transforms.__call__(x)
for x in self.compatibility_buffer[category_pair][1]
], 0)
labelTensor = from_numpy(
np.array(
self.compatibility_labels[category_pair])).float()
except RuntimeError:
self.RuntimeErrorLog.write(','.join(
[str(i)
for i in self.compatibility_lines[category_pair]]) + '\n')
return None
self.compatibility_buffer[category_pair][0] = []
self.compatibility_buffer[category_pair][1] = []
self.compatibility_labels[category_pair] = []
self.compatibility_lines[category_pair] = []
return [
objective, im1Tensor, im2Tensor, labelTensor, category_pair
]
else:
try:
im1Tensor = stack([
self.transforms.__call__(x)
for x in self.similarity_buffer[0]
], 0)
im2Tensor = stack([
self.transforms.__call__(x)
for x in self.similarity_buffer[1]
], 0)
labelTensor = from_numpy(np.array(
self.similarity_labels)).float()
except RuntimeError:
self.RuntimeErrorLog.write(
','.join([str(i) for i in self.similarity_lines]) + '\n')
return None
self.similarity_buffer[0] = []
self.similarity_buffer[1] = []
self.similarity_labels = []
return ['S', im1Tensor, im2Tensor, labelTensor]
