-
Notifications
You must be signed in to change notification settings - Fork 2
/
generate_augmenter_list.py
executable file
·112 lines (92 loc) · 4.72 KB
/
generate_augmenter_list.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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
from arguments import generator_options
from get_backgrounds_and_data import fetch_background_images
import numpy as np
import random
from scipy.spatial.distance import cdist
def remove_clutter(objects_list, augmenter_list, regenerate_count):
"""
This function removes elements from the augmenter list which have
objects occupy an area in image space larger than "max_occupied_area"
or have objects which are too close to each other determined my
"min_distance".
:param objects_list: List containing details of all objects.
:param augmenter_list: Augmenter list containing details regarding the
artificial image.
:param regenerate_count: Count of attempts already made to create new
elements for augmenter list.
:return: No returns.
"""
removed_vectors = 0
for index, vector in enumerate(augmenter_list):
vector_area = 0
for i in range(vector['num_objects_to_place']):
vector_area += objects_list[vector['what_objects'][i]]['obj_area']
dist_btw_locations = cdist(vector['locations'], vector['locations'])
np.fill_diagonal(dist_btw_locations, np.inf)
if (vector_area > np.product(generator_options.image_dimension)
* generator_options.max_occupied_area or
np.any(dist_btw_locations < generator_options.min_distance)):
del augmenter_list[index]
removed_vectors += 1
if (regenerate_count < generator_options.num_regenerate
and removed_vectors is not 0):
regenerate_count += 1
create_augmenter_list(objects_list, is_regeneration=True,
removed_elements=removed_vectors,
regenerate_count=regenerate_count,
augmenter_list=augmenter_list)
def get_random_locations(num_random_locations):
"""
:param num_random_locations: Number of random locations in pixel space.
:return: A list of random (x,y) locations in pixel space.
"""
location = [[random.randrange(0, 440, 120), random.randrange(0, 600, 120)]
for _ in range(num_random_locations)]
return np.array(location)
def create_augmenter_list(objects_list, is_regeneration=False, removed_elements=None,
regenerate_count=None, augmenter_list=None):
"""
:param objects_list: List containing details of all objects.
:param is_regeneration: Currently regenerate elements or create new
list of elements.
:param removed_elements: How many elements have been removed and needs
to be regenerated.
:param regenerate_count: How many regeneration attempts have already
been made.
:param augmenter_list: A list of elements where each element is a dictionary
containing the following details regarding the artificial
image:
1. 'background_image': The background image to use.
2. 'num_objects_to_place': Number of objects to be placed.
3. 'what_objects': A list of indexes indication what objects
need to be placed.
4. 'locations': A list of coordinates in the image where
each object needs to be placed.
:return: The generated augmenter list.
"""
objects_index = np.arange(0, len(objects_list))
background_images = fetch_background_images()
if is_regeneration:
augmenter_list = augmenter_list
num_images = removed_elements
regenerate_count = regenerate_count
else:
augmenter_list = []
num_images = generator_options.num_images
regenerate_count = 0
for i in range(num_images):
num_objects_to_place = np.random.randint(1,
high=generator_options.max_objects)
what_objects = [objects_index[i] for i in range(num_objects_to_place)]
if i % len(background_images) == 0:
np.random.shuffle(background_images)
np.random.shuffle(objects_index)
augmenter_list.append({'background_image': background_images[
i % len(background_images)],
'num_objects_to_place': num_objects_to_place,
'what_objects': what_objects,
'locations': get_random_locations(num_objects_to_place)})
if generator_options.remove_clutter:
remove_clutter(objects_list, augmenter_list,
regenerate_count)
return augmenter_list