def create_template_dict(self, classes, to_tensor=False):
template_dict = {}
for class_name in classes:
if to_tensor:
template_dict[class_name] = load(
join_path(self.root, f'templates/{class_name}.pt'))
else:
template_dict[class_name] = open_image(
join_path(self.root, f'templates/{class_name}.jpg'))
return template_dict
def __init__(self, root, image_shape=None, train=True, transform=None):
self.root = root
self.template_path = join_path(self.root, 'human_template.png')
self.image_shape = image_shape
self.transform = transform
if train:
self.root = join_path(self.root, 'train_data')
else:
self.root = join_path(self.root, 'test_data')
self.files_list = list_files(join_path(self.root, 'images'))
self.size = len(self.files_list)
def __init__(self, root, train=True, transform=None):
self.root = root
self.class_names = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
self.templates_path = join_path(self.root, 'templates')
self.transform = transform
if train:
self.root = join_path(self.root, 'train')
else:
self.root = join_path(self.root, 'test')
self.files_list = list_files(join_path(self.root, 'images'))
self.size = len(self.files_list)
with open(join_path(self.root, 'counts', 'counts.txt'), 'r') as counts_file:
self.count_lines = counts_file.readlines()
def __init__(self,
model,
criterion,
optimizer,
run_name,
device=torch.device('cpu')):
self.model = model
self.criterion = criterion
self.optimizer = optimizer
self.device = device
self.run_name = run_name
logs_path = create_dirs(f'logs/{run_name}')
self.train_writer = SummaryWriter(join_path(logs_path, 'train'))
self.val_writer = SummaryWriter(join_path(logs_path, 'val'))
def open_templates(self):
templates = []
for class_name in self.class_names:
template_name = class_name + '.jpg'
template = Image.open(join_path(self.templates_path, template_name))
if self.transform is not None:
template = self.transform(template)
templates.append(template)
return templates
def __init__(self,
root,
image_shape=None,
data_percentage=1,
train=True,
transform=None):
self.subset_folder = 'train' if train else 'test'
self.image_root = join_path(root, 'Images')
self.anno_root = join_path(root, 'Annotations')
self.image_sets_file = join_path(root, 'ImageSets',
self.subset_folder + '.txt')
self.image_shape = image_shape
self.transform = transform
self.size = 0
self.files = []
with open(self.image_sets_file, 'r') as image_sets_file:
self.files = image_sets_file.read().splitlines()
self.size = len(self.files)
def __init__(self,
root,
image_shape=None,
data_percentage=1,
train=True,
transform=None):
self.subset_folder = 'train' if train else 'val'
self.data_root = join_path(root, 'Data', 'VID', self.subset_folder)
self.anno_root = join_path(root, 'Annotations', 'VID',
self.subset_folder)
self.image_shape = image_shape
self.transform = transform
self.size = 0
self.folder_size = []
for folder in list_files(self.data_root):
images = list_files(join_path(self.data_root, folder))
self.size += len(images)
self.folder_size.append(self.size)
self.size = int(self.size * data_percentage)
def __getitem__(self, index):
im = Image.open(join_path(self.root, 'images', self.files_list[index]))
templates = self.open_templates()
counts = self.count_lines[index].split(' ')
counts = np.asarray([int(string) for string in counts], dtype=np.float32)
counts = counts.reshape((counts.shape[0], 1))
if self.transform is not None:
im = self.transform(im)
return im, templates, counts
def __getitem__(self, index):
im = Image.open(join_path(self.root, 'images', self.files_list[index]))
template = Image.open(self.template_path)
if self.image_shape is not None:
im = thumbnail_image(im, self.image_shape)
if self.transform is not None:
im = self.transform(im)
template = self.transform(template)
return im, template
def __getitem__(self, index):
file = self.files[index]
image_file_path = join_path(self.image_root, file + '.png')
anno_file_path = join_path(self.anno_root, file + '.txt')
im = Image.open(image_file_path)
annotations = Annotation(anno_file_path)
if self.image_shape is None:
size = im.size
else:
size = self.image_shape
size = im.size
# size = (size[0] // 4, size[1] // 4)
ground_truth = Image.new('L', size)
template_box = None
max_area = 0
for gaussian_center, bounding_box in zip(annotations.centers,
annotations.bounding_boxes):
current_area = ((bounding_box[1][0] - bounding_box[0][0]) *
(bounding_box[1][1] - bounding_box[0][1]))
if max_area < current_area:
max_area = ((bounding_box[1][0] - bounding_box[0][0]) *
(bounding_box[1][1] - bounding_box[0][1]))
template_box = bounding_box
box = [[
int(bounding_box[0][0] * size[0] / im.size[0]),
int(bounding_box[0][1] * size[1] / im.size[1])
],
[
int(bounding_box[1][0] * size[0] / im.size[0]),
int(bounding_box[1][1] * size[1] / im.size[1])
]]
shape = [(box[1][0] - box[0][0]), (box[1][1] - box[0][1])]
gaussian = gkern(shape)
gaussian *= 255
gaussian = gaussian.astype(np.uint8)
gaussian = Image.fromarray(gaussian, mode='L')
gaussian_center = (
int((gaussian_center[0] * size[0] / im.size[0]) -
gaussian.size[0] // 2),
int((gaussian_center[1] * size[1] / im.size[1]) -
gaussian.size[1] // 2))
ground_truth = paste_image(ground_truth, gaussian, gaussian_center)
# ground_truth.show()
if template_box is not None:
coords = self.square_template(template_box, im)
template = im.crop(tuple(coords))
template = thumbnail_image(template, (63, 63))
else:
template = Image.new('RGB', (63, 63))
if self.image_shape is not None:
im = thumbnail_image(im, self.image_shape)
padding = Image.new('RGBA', self.image_shape)
x = int(padding.size[0] / 2 - im.size[0] / 2)
y = int(padding.size[1] / 2 - im.size[1] / 2)
im = im.convert('RGBA')
im = paste_image(padding, im, (x, y)).convert('RGB')
size = (self.image_shape[0] // 4, self.image_shape[1] // 4)
ground_truth = thumbnail_image(ground_truth, size)
padding = Image.new('RGBA', size)
x = int(padding.size[0] / 2 - ground_truth.size[0] / 2)
y = int(padding.size[1] / 2 - ground_truth.size[1] / 2)
ground_truth = ground_truth.convert('RGBA')
ground_truth = paste_image(padding, ground_truth,
(x, y)).convert('L')
# im.show()
# template.show()
count = len(annotations.centers)
resized_template = resize_image(template, (96, 96))
if self.transform is not None:
im = self.transform(im)
template = self.transform(template)
ground_truth = self.transform(ground_truth)
resized_template = self.transform(resized_template)
return im, template, ground_truth, count, resized_template