class ParserTestCase(unittest.TestCase):
def setUp(self):
self.parser = Parser()
def test_parse_args(self):
parsed_args = self.parser.parse_args(['-m', 'HDPE', '-p', '100', '-c', '1',
'-i', '1', '-b', 'background', '-o', 'images'])
self.assertEqual(parsed_args.materials, ['HDPE'])
def test_exception_1(self):
self.assertRaises(OSError, lambda: self.parser.parse_args(['-m', 'HDPE', 'PET', '-p', '100',
'-c', '1', '-i', '1', '-b',
'background', '-o', 'images']))
def test_exception_2(self):
self.assertRaises(OSError, lambda: self.parser.parse_args(['-m', 'HDPE', '-p', '80', '-c',
'1', '-i', '1', '-b',
'background', '-o', 'images']))
def test_yaml(self):
data = self.parser.parse_long_term_configuration(pathlib.Path(
src_dir + r"/test/util/test.yaml"))
self.assertEqual(data, {'key1': {'key1.1': 'stringvalue'},
'key2': {'key2.1': 1, 'key2.2': [1, 2, 3]}})
def main(args): # noqa: CFQ001
"""
Main method of our code.
Sets the paths of background and object.
:param args: arguments used for parsing to values.
:return: time_data: a dictionary with timestamps of the different stages of the pipeline
"""
total_start_time = time.time()
parser = Parser()
# Load the config file for camera and lights
configuration = parser.parse_long_term_configuration(pathlib.Path(
src_dir + r"/configuration.yaml"))
# Parse arguments provided as input by user
args = parser.parse_args(args)
# Load objects
if args.reuse_crushes:
material_dirs = list(map(lambda x: list(pathlib.Path(src_dir + r'/Crushed Models/' +
x).glob('**/*.obj')), args.materials))
else:
material_dirs = list(map(lambda x: list(pathlib.Path(src_dir + r'/Models/' + x).glob(
'**/*.obj')), args.materials))
# List of the models
models = [list(map(lambda path: Object(str(
path), 'random', 'random', 'random', None), material)) for material in material_dirs]
# Calculate number of objects per material based on proportions
number_of_objects = list(map(lambda x: int(round(
args.objects_per_image * (x / 100))), args.proportions))
objects = make_object_selection(args, number_of_objects, models)
if not args.only_crush:
# Setup scene
scene = Scene(configuration['camera']['location'], configuration['camera']['rotation'],
configuration['light']['location'], configuration['light']['energy'],
configuration['light']['type'])
scene.add_background(args.background, None)
# Render images
image_object_bboxes = render(args, configuration['render'], scene, objects)
# Write info file
write_file(configuration['info_json'],
(configuration['render']['res_width'], configuration['render']['res_height']),
'info', args.image_count, image_object_bboxes)
time_data['total'] = time.time() - total_start_time
# Print time data
print('Total Time: ' + str(time_data['total']))
print('Object Creation Time: ' + str(time_data['object_creation_time']))
if not args.only_crush:
print('Object Setup Time: ' + str(time_data['object_setup_time']))
for i in range(args.image_count):
print('Image ' + str(i) + ' Time: ' + str(time_data['image_' + str(i)]))
# To be displayed on the server page
return time_data