diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3d1c133 --- /dev/null +++ b/.gitignore @@ -0,0 +1,43 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + + +.DS_Store +datasets/ +*.zip. +.idea/ +Models/ +Logs/ +cls/ +slurm-*.out +node_modules/ +neptune_config.yml +neptune_session.json +.neptune/ +.vscode/ +*.egg_inf +lightning_logs/ +stability_results.json +dataset_processing/model-converter-python/ diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/common/__init__.py b/common/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/common/bpy_util.py b/common/bpy_util.py new file mode 100644 index 0000000..c67e0ef --- /dev/null +++ b/common/bpy_util.py @@ -0,0 +1,165 @@ +import bpy +import math +from mathutils import Vector +from typing import Union +from pathlib import Path + + +def save_obj(target_obj_file_path: Union[Path, str], additional_objs_to_save=None, simplification_ratio=None): + """ + save the object and returns a mesh duplicate version of it + """ + obj = select_shape() + refresh_obj_in_viewport(obj) + dup_obj = copy(obj) + # set active + bpy.ops.object.select_all(action='DESELECT') + dup_obj.select_set(True) + bpy.context.view_layer.objects.active = dup_obj + # apply the modifier to turn the geometry node to a mesh + bpy.ops.object.modifier_apply(modifier="GeometryNodes") + if simplification_ratio and simplification_ratio < 1.0: + bpy.ops.object.modifier_add(type='DECIMATE') + dup_obj.modifiers["Decimate"].decimate_type = 'COLLAPSE' + dup_obj.modifiers["Decimate"].ratio = simplification_ratio + bpy.ops.object.modifier_apply(modifier="Decimate") + assert dup_obj.type == 'MESH' + bpy.ops.object.transform_apply(location=True, rotation=True, scale=True) + # set origin to center of bounding box + bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY', center='BOUNDS') + dup_obj.location.x = dup_obj.location.y = dup_obj.location.z = 0 + normalize_scale(dup_obj) + if additional_objs_to_save: + for additional_obj in additional_objs_to_save: + additional_obj.select_set(True) + # save + bpy.ops.export_scene.obj(filepath=str(target_obj_file_path), use_selection=True, use_materials=False, use_triangles=True) + return dup_obj + + +def get_geometric_nodes_modifier(obj): + # loop through all modifiers of the given object + gnodes_mod = None + for modifier in obj.modifiers: + # check if current modifier is the geometry nodes modifier + if modifier.type == "NODES": + gnodes_mod = modifier + break + return gnodes_mod + + +def normalize_scale(obj): + obj.select_set(True) + bpy.context.view_layer.objects.active = obj + bpy.ops.object.transform_apply(location=True, rotation=True, scale=True) + # set origin to the center of the bounding box + bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY', center='BOUNDS') + + obj.location.x = 0 + obj.location.y = 0 + obj.location.z = 0 + + bpy.ops.object.transform_apply(location=True, rotation=True, scale=True) + max_vert_dist = math.sqrt(max([v.co.dot(v.co) for v in obj.data.vertices])) + + for v in obj.data.vertices: + v.co /= max_vert_dist + + bpy.ops.object.transform_apply(location=True, rotation=True, scale=True) + + # verify that the shape is normalized + # max_vert_dist = math.sqrt(max([v.co.dot(v.co) for v in obj.data.vertices])) + # assert abs(max_vert_dist - 1.0) < 0.01 + + +def setup_lights(): + """ + setup lights for rendering + used for visualization of 3D objects as images + """ + scene = bpy.context.scene + # light 1 + light_data_1 = bpy.data.lights.new(name="light_data_1", type='POINT') + light_data_1.energy = 300 + light_object_1 = bpy.data.objects.new(name="Light_1", object_data=light_data_1) + light_object_1.location = Vector((10, -10, 10)) + scene.collection.objects.link(light_object_1) + # light 2 + light_data_2 = bpy.data.lights.new(name="light_data_2", type='POINT') + light_data_2.energy = 300 + light_object_2 = bpy.data.objects.new(name="Light_2", object_data=light_data_2) + light_object_2.location = Vector((-10, -10, 10)) + scene.collection.objects.link(light_object_2) + # light 3 + light_data_3 = bpy.data.lights.new(name="light_data_3", type='POINT') + light_data_3.energy = 300 + light_object_3 = bpy.data.objects.new(name="Light_3", object_data=light_data_3) + light_object_3.location = Vector((10, 0, 10)) + scene.collection.objects.link(light_object_3) + + +def look_at(obj_camera, point): + """ + orient the given camera with a fixed position to loot at a given point in space + """ + loc_camera = obj_camera.matrix_world.to_translation() + direction = point - loc_camera + # point the cameras '-Z' and use its 'Y' as up + rot_quat = direction.to_track_quat('-Z', 'Y') + obj_camera.rotation_euler = rot_quat.to_euler() + + +def clean_scene(start_with_strings=["Camera", "procedural", "Light"]): + """ + delete all object of which the name's prefix is matching any of the given strings + """ + scene = bpy.context.scene + bpy.ops.object.select_all(action='DESELECT') + for obj in scene.objects: + if any([obj.name.startswith(starts_with_string) for starts_with_string in start_with_strings]): + # select the object + if obj.visible_get(): + obj.select_set(True) + bpy.ops.object.delete() + + +def del_obj(obj): + bpy.ops.object.select_all(action='DESELECT') + obj.select_set(True) + bpy.ops.object.delete() + + +def refresh_obj_in_viewport(obj): + # the following two line cause the object to update according to the new geometric nodes input + obj.show_bounds = not obj.show_bounds + obj.show_bounds = not obj.show_bounds + + +def select_objs(*objs): + bpy.ops.object.select_all(action='DESELECT') + for i, obj in enumerate(objs): + if i == 0: + bpy.context.view_layer.objects.active = obj + obj.select_set(True) + + +def select_obj(obj): + select_objs(obj) + + +def select_shape(): + """ + select the procedural shape in the blend file + note that in all our domains, the procedural shape is named "procedural shape" within the blend file + """ + obj = bpy.data.objects["procedural shape"] + select_obj(obj) + return obj + + +def copy(obj): + dup_obj = obj.copy() + dup_obj.data = obj.data.copy() + dup_obj.animation_data_clear() + bpy.context.collection.objects.link(dup_obj) + return dup_obj diff --git a/common/domain.py b/common/domain.py new file mode 100644 index 0000000..ec2dc22 --- /dev/null +++ b/common/domain.py @@ -0,0 +1,9 @@ +from enum import Enum + +class Domain(Enum): + chair = 'chair' + vase = 'vase' + table = 'table' + + def __str__(self): + return self.value diff --git a/common/file_util.py b/common/file_util.py new file mode 100644 index 0000000..73a3d67 --- /dev/null +++ b/common/file_util.py @@ -0,0 +1,50 @@ +import yaml +import hashlib +import numpy as np +from pathlib import Path +from typing import Union + + +def save_yml(yml_obj, target_yml_file_path): + with open(target_yml_file_path, 'w') as target_yml_file: + yaml.dump(yml_obj, target_yml_file, sort_keys=False, width=1000) + + +def get_source_recipe_file_path(domain): + """ + get the path to the recipe file path that is found in the source code under the directory "recipe_files" + """ + return Path(__file__).parent.joinpath('..', 'dataset_generator', 'recipe_files', f'recipe_{domain}.yml').resolve() + + +def hash_file_name(file_name): + return int(hashlib.sha1(file_name.encode("utf-8")).hexdigest(), 16) % (10 ** 8) + + +def get_recipe_yml_obj(recipe_file_path: Union[str, Path]): + with open(recipe_file_path, 'r') as recipe_file: + recipe_yml_obj = yaml.load(recipe_file, Loader=yaml.FullLoader) + return recipe_yml_obj + + +def load_obj(file: str): + vs, faces = [], [] + f = open(file) + for line in f: + line = line.strip() + split_line = line.split() + if not split_line: + continue + elif split_line[0] == 'v': + vs.append([float(v) for v in split_line[1:4]]) + elif split_line[0] == 'f': + face_vertex_ids = [int(c.split('/')[0]) for c in split_line[1:]] + assert len(face_vertex_ids) == 3 + face_vertex_ids = [(ind - 1) if (ind >= 0) else (len(vs) + ind) + for ind in face_vertex_ids] + faces.append(face_vertex_ids) + f.close() + vs = np.asarray(vs) + faces = np.asarray(faces, dtype=np.int64) + assert np.logical_and(faces >= 0, faces < len(vs)).all() + return vs, faces diff --git a/common/input_param_map.py b/common/input_param_map.py new file mode 100644 index 0000000..1673b49 --- /dev/null +++ b/common/input_param_map.py @@ -0,0 +1,170 @@ +import yaml +import random +import traceback +import numpy as np +from typing import List, Optional +from dataclasses import dataclass +from bpy.types import NodeInputs, Modifier +from common.bpy_util import select_shape, refresh_obj_in_viewport, get_geometric_nodes_modifier + + +@dataclass +class InputParam: + gnodes_mod: Modifier + input: NodeInputs + axis: Optional[str] # None indicates that this is not a vector + possible_values: List + + def assign_random_value(self): + self.assign_value(random.choice(self.possible_values)) + + def assign_value(self, val): + assert val in self.possible_values + input_type = self.input.bl_label + identifier = self.input.identifier + if input_type == "Float": + self.gnodes_mod[identifier] = val + + if input_type == "Integer": + self.gnodes_mod[identifier] = int(val) + + if input_type == "Boolean": + self.gnodes_mod[identifier] = int(val) + + if input_type == "Vector": + axis_idx = ['x', 'y', 'z'].index(self.axis) + self.gnodes_mod[identifier][axis_idx] = val + + def get_value(self): + identifier = self.input.identifier + if self.axis: + axis_idx = ['x', 'y', 'z'].index(self.axis) + return self.gnodes_mod[identifier][axis_idx] + return self.gnodes_mod[identifier] + + def get_name_for_file(self): + res = str(self.input.name) + ("" if not self.axis else "_" + self.axis) + return res.replace(" ", "_") + + +def get_input_values(input, yml_gen_rule): + min_value = None + max_value = None + if input.bl_label != 'Boolean': + min_value = input.min_value + max_value = input.max_value + # override min and max with requested values from recipe yml file + if 'min' in yml_gen_rule: + requested_min_value = yml_gen_rule['min'] + if min_value and requested_min_value < min_value: + if abs(min_value - requested_min_value) > 1e-6: + raise Exception( + f'Requested a min value of [{requested_min_value}] for parameter [{input.name}], but min allowed is [{min_value}]') + # otherwise min_value should remain input.min_value + else: + min_value = requested_min_value + if 'max' in yml_gen_rule: + requested_max_value = yml_gen_rule['max'] + if max_value and requested_max_value > max_value: + if abs(max_value - requested_max_value) > 1e-6: + raise Exception( + f'Requested a max value of [{requested_max_value}] for parameter [{input.name}], but max allowed is [{max_value}]') + # otherwise max_value should remain input.max_value + max_value = requested_max_value + step = 1 if 'samples' not in yml_gen_rule else calculate_step(min_value, max_value, yml_gen_rule['samples']) + res = np.arange(min_value, max_value + 1e-6, step) + + # convert to integers if needed + if input.bl_label in ['Boolean', 'Integer']: + res = list(res.astype(int)) + else: + res = [round(x, 4) for x in list(res)] + + return res + + +def calculate_step(min_value, max_value, samples): + return (max_value - min_value) / (samples - 1) + + +def get_input_param_map(gnodes_mod, yml): + input_params_map = {} + # loops through all the inputs in the geometric node group + for param_name in yml['dataset_generation']: + if param_name not in gnodes_mod.node_group.inputs: + raise Exception(f"Parameter named [{param_name}] was not found in geometry nodes input group.") + for input in gnodes_mod.node_group.inputs: + param_name = str(input.name) + + # we only change inputs that are explicitly noted in the yaml object + if param_name in yml['dataset_generation']: + param_gen_rule = yml['dataset_generation'][param_name] + if 'x' in param_gen_rule or 'y' in param_gen_rule or 'z' in param_gen_rule: + # vector handling + for idx, axis in enumerate(['x', 'y', 'z']): + if not axis in param_gen_rule: + continue + curr_param_values = get_input_values(input, param_gen_rule[axis]) + input_params_map[f"{param_name} {axis}"] = InputParam(gnodes_mod, input, axis, curr_param_values) + else: + curr_param_values = get_input_values(input, param_gen_rule) + input_params_map[param_name] = InputParam(gnodes_mod, input, None, curr_param_values) + return input_params_map + + +def yml_to_shape(shape_yml_obj, input_params_map, ignore_sanity_check=False): + try: + # select the object in blender + obj = select_shape() + # get the geometric nodes modifier fo the object + gnodes_mod = get_geometric_nodes_modifier(obj) + + # loops through all the inputs in the geometric node group + for input in gnodes_mod.node_group.inputs: + param_name = str(input.name) + if param_name not in shape_yml_obj: + continue + param_val = shape_yml_obj[param_name] + if hasattr(param_val, '__iter__'): + # vector handling + for axis_idx, axis in enumerate(['x', 'y', 'z']): + val = param_val[axis] + val = round(val, 4) + param_name_with_axis = f'{param_name} {axis}' + gnodes_mod[input.identifier][axis_idx] = val if abs(val + 1.0) > 0.1 else input_params_map[param_name_with_axis].possible_values[0].item() + assert gnodes_mod[input.identifier][axis_idx] >= 0.0 + else: + param_val = round(param_val, 4) + if not ignore_sanity_check: + err_msg = f'param_name [{param_name}] param_val [{param_val}] possible_values {input_params_map[param_name].possible_values}' + assert param_val == -1 or (param_val in input_params_map[param_name].possible_values), err_msg + gnodes_mod[input.identifier] = param_val if (abs(param_val + 1.0) > 0.1) else (input_params_map[param_name].possible_values[0].item()) + # we assume that all input values are non-negative + assert gnodes_mod[input.identifier] >= 0.0 + + refresh_obj_in_viewport(obj) + except Exception as e: + print(repr(e)) + print(traceback.format_exc()) + + +def load_shape_from_yml(yml_file_path, input_params_map, ignore_sanity_check=False): + with open(yml_file_path, 'r') as f: + yml_obj = yaml.load(f, Loader=yaml.FullLoader) + yml_to_shape(yml_obj, input_params_map, ignore_sanity_check=ignore_sanity_check) + + +def load_base_shape_from_yml(recipe_file_path, input_params_map): + print(f'Loading the base shape from the YML file [{recipe_file_path}]') + + with open(recipe_file_path, 'r') as f: + yml_obj = yaml.load(f, Loader=yaml.FullLoader) + + yml_to_shape(yml_obj['base'], input_params_map) + + +def randomize_all_params(input_params_map): + param_values_map = {} + for param_name, input_param in input_params_map.items(): + param_values_map[param_name] = random.choice(input_param.possible_values) + return param_values_map diff --git a/common/intersection_util.py b/common/intersection_util.py new file mode 100644 index 0000000..bee5045 --- /dev/null +++ b/common/intersection_util.py @@ -0,0 +1,161 @@ +import bpy +import array +import mathutils +from object_print3d_utils import mesh_helpers +from common.bpy_util import select_objs, select_shape, refresh_obj_in_viewport + + +def isolate_node_as_final_geometry(obj, node_label): + gm = obj.modifiers.get("GeometryNodes") + group_output_node = None + node_to_isolate = None + for n in gm.node_group.nodes: + # print(n.name), print(n.type), print(dir(n)) + if n.type == 'GROUP_OUTPUT': + group_output_node = n + elif n.label == node_label: + node_to_isolate = n + if not node_to_isolate: + raise Exception(f"Did not find any node with the label [{node_label}]") + + realize_instances_node = group_output_node.inputs[0].links[0].from_node + third_to_last_node = realize_instances_node.inputs[0].links[0].from_node + third_to_last_node_socket = None + # to later revert this operation, we need to find the socket which is currently connected + # this happens since the SWITCH node has multiple options, and each option translates to + # a different output socket in the node (so there isn't just one socket as you would think) + for i, socket in enumerate(third_to_last_node.outputs): + if socket.is_linked: + third_to_last_node_socket = i + break + node_group = next(m for m in obj.modifiers if m.type == 'NODES').node_group + # find the output socket that actually is connected to something, + # we do this since some nodes have multiple output sockets + out_socket_idx = 0 + for out_socket_idx, out_socket in enumerate(node_to_isolate.outputs): + if out_socket.is_linked: + break + node_group.links.new(node_to_isolate.outputs[out_socket_idx], realize_instances_node.inputs[0]) + def revert(): + node_group.links.new(third_to_last_node.outputs[third_to_last_node_socket], realize_instances_node.inputs[0]) + refresh_obj_in_viewport(obj) + return revert + + +def detect_self_intersection(obj): + """ + refer to: + https://blenderartists.org/t/self-intersection-detection/671080 + documentation of the intersection detection method + https://docs.blender.org/api/current/mathutils.bvhtree.html + """ + if not obj.data.polygons: + return array.array('i', ()) + + bm = mesh_helpers.bmesh_copy_from_object(obj, transform=False, triangulate=False) + tree = mathutils.bvhtree.BVHTree.FromBMesh(bm, epsilon=0.00001) + + overlap = tree.overlap(tree) + faces_error = {i for i_pair in overlap for i in i_pair} + return array.array('i', faces_error) + + +def find_self_intersections(node_label): + # intersection detection + chair = select_shape() + revert_isolation = isolate_node_as_final_geometry(chair, node_label) + + dup_obj = chair.copy() + dup_obj.data = chair.data.copy() + dup_obj.animation_data_clear() + bpy.context.collection.objects.link(dup_obj) + # move for clarity + dup_obj.location.x += 2.0 + + # set active + bpy.ops.object.select_all(action='DESELECT') + dup_obj.select_set(True) + bpy.context.view_layer.objects.active = dup_obj + # apply the modifier to turn the geometry node to a mesh + bpy.ops.object.modifier_apply(modifier="GeometryNodes") + assert dup_obj.type == 'MESH' + + intersections = detect_self_intersection(dup_obj) + + # delete the duplicate + bpy.ops.object.delete() + + revert_isolation() + + # reselect the original object + select_shape() + + return len(intersections) + + +def detect_cross_intersection(obj1, obj2): + if not obj1.data.polygons or not obj2.data.polygons: + return array.array('i', ()) + + bm1 = mesh_helpers.bmesh_copy_from_object(obj1, transform=False, triangulate=False) + tree1 = mathutils.bvhtree.BVHTree.FromBMesh(bm1, epsilon=0.00001) + bm2 = mesh_helpers.bmesh_copy_from_object(obj2, transform=False, triangulate=False) + tree2 = mathutils.bvhtree.BVHTree.FromBMesh(bm2, epsilon=0.00001) + + overlap = tree1.overlap(tree2) + faces_error = {i for i_pair in overlap for i in i_pair} + return array.array('i', faces_error) + + +def find_cross_intersections(node_label1, node_label2): + # intersection detection + chair = select_shape() + revert_isolation = isolate_node_as_final_geometry(chair, node_label1) + + dup_obj1 = chair.copy() + dup_obj1.data = chair.data.copy() + dup_obj1.animation_data_clear() + bpy.context.collection.objects.link(dup_obj1) + # move for clarity + dup_obj1.location.x += 2.0 + # set active + bpy.ops.object.select_all(action='DESELECT') + dup_obj1.select_set(True) + bpy.context.view_layer.objects.active = dup_obj1 + # apply the modifier to turn the geometry node to a mesh + bpy.ops.object.modifier_apply(modifier="GeometryNodes") + # export the object + assert dup_obj1.type == 'MESH' + + revert_isolation() + + chair = select_shape() + revert_isolation = isolate_node_as_final_geometry(chair, node_label2) + + dup_obj2 = chair.copy() + dup_obj2.data = chair.data.copy() + dup_obj2.animation_data_clear() + bpy.context.collection.objects.link(dup_obj2) + # move for clarity + dup_obj2.location.x += 2.0 + # set active + bpy.ops.object.select_all(action='DESELECT') + dup_obj2.select_set(True) + bpy.context.view_layer.objects.active = dup_obj2 + # apply the modifier to turn the geometry node to a mesh + bpy.ops.object.modifier_apply(modifier="GeometryNodes") + # export the object + assert dup_obj2.type == 'MESH' + + revert_isolation() + + intersections = detect_cross_intersection(dup_obj1, dup_obj2) + + # delete the duplicate + select_objs(dup_obj1, dup_obj2) + bpy.ops.object.delete() + + # reselect the original object + select_shape() + + return len(intersections) diff --git a/common/param_descriptors.py b/common/param_descriptors.py new file mode 100644 index 0000000..229a7ae --- /dev/null +++ b/common/param_descriptors.py @@ -0,0 +1,230 @@ +import numpy as np +from dataclasses import dataclass +from collections import OrderedDict +from typing import List, Optional, Dict + + +arithmetic_symbols = ['and', 'or', 'not', '(', ')', '<', '<=' , '>', '>=', '==', '-', '+', '/', '*'] + +def isfloat(num): + try: + float(num) + return True + except ValueError: + return False + + +@dataclass +class ParamDescriptor: + input_type: str + num_classes: int + step: float + classes: np.ndarray + normalized_classes: np.ndarray + min_val: float + max_val: float + visibility_condition: str + is_regression: bool + normalized_acc_threshold: float + + def is_visible(self, param_values_map): + assert param_values_map + if not self.visibility_condition: + return True + is_visible_cond = " ".join([(word if (word in arithmetic_symbols or isfloat(word) or word.isnumeric()) else (f"param_values_map[\"{word}\"] == 1" if 'is_' in word else f"param_values_map[\"{word}\"]")) for word in self.visibility_condition.split(" ")]) + return eval(is_visible_cond) + + +class ParamDescriptors: + def __init__(self, recipe_yml_obj, inputs_to_eval, use_regression=False, train_with_visibility_label=True): + self.epsilon = 1e-6 + self.recipe_yml_obj = recipe_yml_obj + self.inputs_to_eval = inputs_to_eval + self.use_regression = use_regression + self.train_with_visibility_label = train_with_visibility_label + self.__overall_num_of_classes_without_visibility_label = 0 + self.param_descriptors_map: Optional[Dict[str, ParamDescriptor]] = None + self.__constraints: Optional[List[str]] = None + + def check_constraints(self, param_values_map): + assert param_values_map + for constraint in self.get_constraints(): + is_fulfilled = " ".join([(word if (word in arithmetic_symbols or isfloat(word) or word.isnumeric()) else (f"param_values_map[\"{word}\"] == 1" if 'is_' in word else f"param_values_map[\"{word}\"]")) for word in constraint.split(" ")]) + if not eval(is_fulfilled): + return False + return True + + def get_constraints(self): + if self.__constraints: + return self.__constraints + self.__constraints = [] + if 'constraints' in self.recipe_yml_obj: + for constraint_name, constraint in self.recipe_yml_obj['constraints'].items(): + self.__constraints.append(constraint) + return self.__constraints + + def get_param_descriptors_map(self): + if self.param_descriptors_map: + return self.param_descriptors_map + recipe_yml_obj = self.recipe_yml_obj # for readability + param_descriptors_map = OrderedDict() + visibility_conditions = {} + if 'visibility_conditions' in recipe_yml_obj: + visibility_conditions = recipe_yml_obj['visibility_conditions'] + + for i, param_name in enumerate(self.inputs_to_eval): + is_regression = False + normalized_acc_threshold = None + if " x" in param_name or " y" in param_name or " z" in param_name: + input_type = recipe_yml_obj['data_types'][param_name[:-2]]['type'] + else: + input_type = recipe_yml_obj['data_types'][param_name]['type'] + if input_type == 'Integer' or input_type == 'Boolean': + max_val = recipe_yml_obj['dataset_generation'][param_name]['max'] + min_val = recipe_yml_obj['dataset_generation'][param_name]['min'] + step = 1 + num_classes = max_val - min_val + step + self.__overall_num_of_classes_without_visibility_label += num_classes + classes = np.arange(min_val, max_val + self.epsilon, step) + normalized_classes = classes - min_val + # visibility label adjustment + if self.train_with_visibility_label: + for vis_cond_name, vis_cond in visibility_conditions.items(): + if vis_cond_name in param_name: + num_classes += 1 + classes = np.concatenate((np.array([-1.0]), classes)) + normalized_classes = np.concatenate((np.array([-1.0]), normalized_classes)) + break + elif input_type == 'Float': + max_val = recipe_yml_obj['dataset_generation'][param_name]['max'] + min_val = recipe_yml_obj['dataset_generation'][param_name]['min'] + samples = recipe_yml_obj['dataset_generation'][param_name]['samples'] + step, num_classes, classes, normalized_classes, is_regression, normalized_acc_threshold \ + = self._handle_float(param_name, samples, min_val, max_val, visibility_conditions) + elif input_type == 'Vector': + axis = param_name[-1] + param_name_no_axis = param_name[:-2] + max_val = recipe_yml_obj['dataset_generation'][param_name_no_axis][axis]['max'] + min_val = recipe_yml_obj['dataset_generation'][param_name_no_axis][axis]['min'] + samples = recipe_yml_obj['dataset_generation'][param_name_no_axis][axis]['samples'] + step, num_classes, classes, normalized_classes, is_regression, normalized_acc_threshold \ + = self._handle_float(param_name_no_axis, samples, min_val, max_val, visibility_conditions) + else: + raise Exception(f'Input type [{input_type}] is not supported yet') + + visibility_condition = None + for vis_cond_name, vis_cond in visibility_conditions.items(): + if vis_cond_name in param_name: + visibility_condition = vis_cond + break + param_descriptors_map[param_name] = ParamDescriptor(input_type, num_classes, step, classes, + normalized_classes, min_val, max_val, + visibility_condition, is_regression, + normalized_acc_threshold) + self.param_descriptors_map = param_descriptors_map + return self.param_descriptors_map + + def _handle_float(self, param_name, samples, min_val, max_val, visibility_conditions): + """ + :param param_name: the parameter name, if the parameter is a vector, the axis should be omitted + :param samples: the number of samples requested in the recipe file + :param min_val: the min value allowed in the recipe file + :param max_val: the max value allowed in the recipe file + :param visibility_conditions: visibility conditions from the recipe file + :return: step, num_classes, classes, normalized_classes, is_regression, normalized_acc_threshold + """ + is_regression = False + normalized_acc_threshold = None + if not self.use_regression: + step = (max_val - min_val) / (samples - 1) + classes = np.arange(min_val, max_val + self.epsilon, step) + classes = classes.astype(np.float64) + normalized_classes = (classes - min_val) / (max_val - min_val) + normalized_classes = normalized_classes.astype(np.float64) + num_classes = classes.shape[0] + self.__overall_num_of_classes_without_visibility_label += num_classes + # visibility label adjustment + if self.train_with_visibility_label: + for vis_cond_name, vis_cond in visibility_conditions.items(): + if vis_cond_name in param_name: + num_classes += 1 + classes = np.concatenate((np.array([-1.0]), classes)) + normalized_classes = np.concatenate((np.array([-1.0]), normalized_classes)) + break + else: + step = 0 + num_classes = 2 # one for prediction and one for visibility label + classes = None + normalized_classes = None + is_regression = True + normalized_acc_threshold = 1 / (2 * (samples - 1)) + return step, num_classes, classes, normalized_classes, is_regression, normalized_acc_threshold + + + def convert_prediction_vector_to_map(self, pred_vector, use_regression=False): + """ + :param pred_vector: predicted vector from the network + :param use_regression: whether we use regression for float values + :return: map object representing the shape + """ + pred_vector = pred_vector.squeeze() + assert len(pred_vector.shape) == 1 + shape_map = {} + idx = 0 + param_descriptors_map = self.get_param_descriptors_map() + for param_name in self.inputs_to_eval: + param_descriptor = param_descriptors_map[param_name] + input_type = param_descriptor.input_type + classes = param_descriptor.classes + num_classes = param_descriptor.num_classes + if input_type == 'Float' or input_type == 'Vector': + if not use_regression: + normalized_pred_class = int(np.argmax(pred_vector[idx:idx + num_classes])) + pred_val = float(classes[normalized_pred_class]) + else: + min_val = param_descriptor.min_val + max_val = param_descriptor.max_val + pred_val = -1.0 + if float(pred_vector[idx + 1]) < 0.5: # visibility class + pred_val = (float(pred_vector[idx]) * (max_val - min_val)) + min_val + else: + # Integer or Boolean + normalized_pred_class = int(np.argmax(pred_vector[idx:idx + num_classes])) + pred_val = int(classes[normalized_pred_class]) + if input_type == 'Vector': + if param_name[:-2] not in shape_map: + shape_map[param_name[:-2]] = {} + shape_map[param_name[:-2]][param_name[-1]] = pred_val + else: + shape_map[param_name] = pred_val + idx += num_classes + return shape_map + + def get_overall_num_of_classes_without_visibility_label(self): + self.get_param_descriptors_map() + return self.__overall_num_of_classes_without_visibility_label + + def expand_target_vector(self, targets): + """ + :param targets: 1-dim target vector which includes a single normalized value for each parameter + :return: 1-dim vector where each parameter prediction is in one-hot representation + """ + targets = targets.squeeze() + assert len(targets.shape) == 1 + res_vector = np.array([]) + param_descriptors = self.get_param_descriptors_map() + for i, param_name in enumerate(self.inputs_to_eval): + param_descriptor = param_descriptors[param_name] + num_classes = param_descriptor.num_classes + if param_descriptor.is_regression: + val = targets[i].reshape(1).item() + if val == -1.0: + res_vector = np.concatenate((res_vector, np.array([0.0, 1.0]))) + else: + res_vector = np.concatenate((res_vector, np.array([val, 0.0]))) + else: + normalized_classes = param_descriptor.normalized_classes + normalized_gt_class_idx = int(np.where(abs(normalized_classes - targets[i].item()) < 1e-3)[0].item()) + one_hot = np.eye(num_classes)[normalized_gt_class_idx] + res_vector = np.concatenate((res_vector, one_hot)) + return res_vector diff --git a/common/point_cloud_util.py b/common/point_cloud_util.py new file mode 100644 index 0000000..1a48244 --- /dev/null +++ b/common/point_cloud_util.py @@ -0,0 +1,16 @@ +import torch + + +def normalize_point_cloud(point_cloud, use_center_of_bounding_box=True): + min_x, max_x = torch.min(point_cloud[:, 0]), torch.max(point_cloud[:, 0]) + min_y, max_y = torch.min(point_cloud[:, 1]), torch.max(point_cloud[:, 1]) + min_z, max_z = torch.min(point_cloud[:, 2]), torch.max(point_cloud[:, 2]) + # center the point cloud + if use_center_of_bounding_box: + center = torch.tensor([(min_x + max_x) / 2, (min_y + max_y) / 2, (min_z + max_z) / 2]) + else: + center = torch.mean(point_cloud, dim=0) + point_cloud = point_cloud - center + dist = torch.max(torch.sqrt(torch.sum((point_cloud ** 2), dim=1))) + point_cloud = point_cloud / dist # scale the point cloud + return point_cloud diff --git a/common/sampling_util.py b/common/sampling_util.py new file mode 100644 index 0000000..bab9b5a --- /dev/null +++ b/common/sampling_util.py @@ -0,0 +1,75 @@ +import torch +from dgl.geometry import farthest_point_sampler + + +def farthest_point_sampling(faces, vertices, num_points=1000): + random_sampling = sample_surface(faces, vertices, num_points=30000) + point_cloud_indices = farthest_point_sampler(random_sampling.unsqueeze(0), npoints=num_points) + point_cloud = random_sampling[point_cloud_indices[0]] + return point_cloud + + +def face_areas_normals(faces, vs): + face_normals = torch.cross(vs[:, faces[:, 1], :] - vs[:, faces[:, 0], :], + vs[:, faces[:, 2], :] - vs[:, faces[:, 1], :], dim=2) + face_areas = torch.norm(face_normals, dim=2) + face_normals = face_normals / face_areas[:, :, None] + face_areas = 0.5 * face_areas + return face_areas, face_normals + + +def sample_surface(faces, vertices, num_points=1000): + """ + sample mesh surface + sample method: + http://mathworld.wolfram.com/TrianglePointPicking.html + Args + --------- + vertices: vertices + faces: triangle faces (torch.long) + num_points: number of samples in the final point cloud + Return + --------- + samples: (count, 3) points in space on the surface of mesh + normals: (count, 3) corresponding face normals for points + """ + bsize, nvs, _ = vertices.shape + weights, normal = face_areas_normals(faces, vertices) + weights_sum = torch.sum(weights, dim=1) + dist = torch.distributions.categorical.Categorical(probs=weights / weights_sum[:, None]) + face_index = dist.sample((num_points,)) + + # pull triangles into the form of an origin + 2 vectors + tri_origins = vertices[:, faces[:, 0], :] + tri_vectors = vertices[:, faces[:, 1:], :].clone() + tri_vectors -= tri_origins.repeat(1, 1, 2).reshape((bsize, len(faces), 2, 3)) + + # pull the vectors for the faces we are going to sample from + face_index = face_index.transpose(0, 1) + face_index = face_index[:, :, None].expand((bsize, num_points, 3)) + tri_origins = torch.gather(tri_origins, dim=1, index=face_index) + face_index2 = face_index[:, :, None, :].expand((bsize, num_points, 2, 3)) + tri_vectors = torch.gather(tri_vectors, dim=1, index=face_index2) + + # randomly generate two 0-1 scalar components to multiply edge vectors by + random_lengths = torch.rand(num_points, 2, 1, device=vertices.device, dtype=tri_vectors.dtype) + + # points will be distributed on a quadrilateral if we use 2x [0-1] samples + # if the two scalar components sum less than 1.0 the point will be + # inside the triangle, so we find vectors longer than 1.0 and + # transform them to be inside the triangle + random_test = random_lengths.sum(dim=1).reshape(-1) > 1.0 + random_lengths[random_test] -= 1.0 + random_lengths = torch.abs(random_lengths) + + # multiply triangle edge vectors by the random lengths and sum + sample_vector = (tri_vectors * random_lengths[None, :]).sum(dim=2) + + # finally, offset by the origin to generate + # (n,3) points in space on the triangle + samples = sample_vector + tri_origins + + # normals = torch.gather(normal, dim=1, index=face_index) + + # return samples, normals + return samples[0] diff --git a/config/neptune_config_example.yml b/config/neptune_config_example.yml new file mode 100644 index 0000000..1b5dfee --- /dev/null +++ b/config/neptune_config_example.yml @@ -0,0 +1,3 @@ +neptune: + api_token = "" + project = "project_dir/project_name" diff --git a/data/__init__.py b/data/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/data/data_processing.py b/data/data_processing.py new file mode 100644 index 0000000..0031e8e --- /dev/null +++ b/data/data_processing.py @@ -0,0 +1,95 @@ +import yaml +import numpy as np +from pathlib import Path +from tqdm import tqdm +import traceback +import torch +from common.file_util import load_obj +from common.point_cloud_util import normalize_point_cloud + + +def generate_point_clouds(data_dir, phase, num_points, num_point_clouds_per_combination, + processed_dataset_dir_name, sampling_method, + gaussian=0.0, apply_point_cloud_normalization=False): + """ + samples point cloud from mesh + """ + dataset_dir = Path(data_dir, phase) + processed_dataset_dir = dataset_dir.joinpath(processed_dataset_dir_name) + processed_dataset_dir.mkdir(exist_ok=True) + files = sorted(dataset_dir.glob('*.obj')) + for file in tqdm(files): + faces = None + vertices = None + should_load_obj = True # this is done to only load the obj if it is actually required + for point_cloud_idx in range(num_point_clouds_per_combination): + new_file_name = Path(processed_dataset_dir, file.with_suffix('.npy').name.replace(".npy", f"_{point_cloud_idx}.npy")) + if new_file_name.is_file(): + continue + # only load the obj (once per all the instances) if it is actually needed + if should_load_obj: + vertices, faces = load_obj(file) + vertices = vertices.reshape(1, vertices.shape[0], vertices.shape[1]) + vertices = torch.from_numpy(vertices) + faces = torch.from_numpy(faces) + should_load_obj = False + + try: + point_cloud = sampling_method(faces, vertices, num_points=num_points) + except Exception as e: + print(traceback.format_exc()) + print(repr(e)) + print(file) + continue + if apply_point_cloud_normalization: + # normalize the point cloud and use center of bounding box + point_cloud = normalize_point_cloud(point_cloud) + if gaussian and gaussian > 0.0: + point_cloud += np.random.normal(0, gaussian, point_cloud.shape) + np.save(str(new_file_name), point_cloud) + + +def normalize_labels(data_dir, phase, processed_dataset_dir_name, params_descriptors, train_with_visibility_label): + dataset_dir = Path(data_dir, phase) + processed_dataset_dir = dataset_dir.joinpath(processed_dataset_dir_name) + processed_dataset_dir.mkdir(exist_ok=True) + + files = sorted(dataset_dir.glob('*.yml')) + for file in files: + if not file.is_file(): + # it is only allowed to not have a gt yml file when we are in test phase + assert phase == "test" + continue + save_path = Path(processed_dataset_dir, file.name) + if save_path.is_file(): + # this will skip normalization if the file exists, but if the recipe file changes, then normalization needs to be performed again + # in that case, disable this if statement to regenerate the normalized labels + continue + with open(file, 'r') as f: + yml_obj = yaml.load(f, Loader=yaml.FullLoader) + normalized_yml_obj = yml_obj.copy() + + # only apply the normalization to the inputs that were changed in this dataset + for param_name, param_descriptor in params_descriptors.items(): + param_input_type = param_descriptor.input_types + min_val = param_descriptor.min_val + max_val = param_descriptor.max_val + if param_input_type == 'Integer': + normalized_yml_obj[param_name] -= min_val + elif param_input_type == 'Float': + value = yml_obj[param_name] + normalized_yml_obj[param_name] = (value - min_val) / (max_val - min_val) + elif param_input_type == 'Boolean': + pass + elif param_input_type == 'Vector': + param_name_no_axis = param_name[:-2] + for axis in ['x', 'y', 'z']: + if param_name[-2:] != f" {axis}": + continue + value = yml_obj[param_name_no_axis][axis] + normalized_yml_obj[param_name_no_axis][axis] = (value - min_val) / (max_val - min_val) + if train_with_visibility_label and not params_descriptors[param_name].is_visible(yml_obj): + normalized_yml_obj[param_name] = -1 + + with open(save_path, 'w') as out_file: + yaml.dump(normalized_yml_obj, out_file) diff --git a/data/dataset_pc.py b/data/dataset_pc.py new file mode 100644 index 0000000..b36d572 --- /dev/null +++ b/data/dataset_pc.py @@ -0,0 +1,123 @@ +import numpy as np +import torch.utils.data as data +import torch +import yaml +from pathlib import Path +from .data_processing import generate_point_clouds, normalize_labels +from common.sampling_util import sample_surface, farthest_point_sampling +from .dataset_util import assemble_targets + + +class DatasetPC(data.Dataset): + def __init__(self, + inputs_to_eval, + dataset_processing_preferred_device, + params_descriptors, + data_dir, + phase, + num_points=1500, + num_point_clouds_per_combination=1, + random_pc=None, + gaussian=0.0, + apply_point_cloud_normalization=False, + scanobjectnn=False, + augment_with_random_points=True, + train_with_visibility_label=True): + self.inputs_to_eval = inputs_to_eval + self.data_dir = data_dir + self.phase = phase + self.random_pc = random_pc + self.gaussian = gaussian + self.apply_point_cloud_normalization = apply_point_cloud_normalization + self.dataset_processing_preferred_device = dataset_processing_preferred_device + self.train_with_visibility_label = train_with_visibility_label + self.yml_gt_normalized_dir_name = 'yml_gt_normalized' + self.point_cloud_fps_dir_name = 'point_cloud_fps' + self.point_cloud_random_dir_name = 'point_cloud_random' + self.num_point_clouds_per_combination = num_point_clouds_per_combination + self.augment_with_random_points = augment_with_random_points + self.ds_path = Path(data_dir, phase) + if not self.ds_path.is_dir(): + raise Exception(f"Could not find a dataset in path [{self.ds_path}]") + + if scanobjectnn: + random_pc_dir = self.ds_path.joinpath(self.point_cloud_random_dir_name) + # [:-2] removes the _0 so that when it is added later it will match the file name + self.file_names = [f.stem[:-2] for f in random_pc_dir.glob("*.npy")] + self.num_files = len(self.file_names) + self.size = self.num_files * self.num_point_clouds_per_combination + return + print(f"Processing dataset [{phase}] with farthest point sampling...") + if not self.random_pc: + generate_point_clouds(data_dir, phase, num_points, self.num_point_clouds_per_combination, + self.point_cloud_fps_dir_name, sampling_method=farthest_point_sampling, gaussian=self.gaussian, + apply_point_cloud_normalization=self.apply_point_cloud_normalization) + else: + num_points = self.random_pc + print(f"Using uniform sampling only with [{num_points}] samples") + normalize_labels(data_dir, phase, self.yml_gt_normalized_dir_name, params_descriptors, self.train_with_visibility_label) + print(f"Processing dataset [{phase}] with uniform sampling (augmentation)...") + generate_point_clouds(data_dir, phase, num_points, self.num_point_clouds_per_combination, + self.point_cloud_random_dir_name, sampling_method=sample_surface, gaussian=self.gaussian, + apply_point_cloud_normalization=self.apply_point_cloud_normalization) + + obj_gt_dir = self.ds_path.joinpath('obj_gt') + self.file_names = [f.stem for f in obj_gt_dir.glob("*.obj")] + + self.num_files = len(self.file_names) + self.size = self.num_files * self.num_point_clouds_per_combination + + + def __getitem__(self, _index): + file_idx = _index // self.num_point_clouds_per_combination + sample_idx = _index % self.num_point_clouds_per_combination + file_name = self.file_names[file_idx] + + pc = [] + random_pc_path = self.ds_path.joinpath(self.point_cloud_random_dir_name, f"{file_name}_{sample_idx}.npy") + fps_pc_path = self.ds_path.joinpath(self.point_cloud_fps_dir_name, f"{file_name}_{sample_idx}.npy") + if self.random_pc: + pc = np.load(str(random_pc_path)) + pc = torch.from_numpy(pc).float() + assert len(pc) == self.random_pc + else: + if fps_pc_path.is_file(): + pc = np.load(str(fps_pc_path)) + pc = torch.from_numpy(pc).float() + + # augment the farthest point sampled point cloud with points from a randomly sampled point cloud + # note that in some tests we did not apply the augmentation + if self.augment_with_random_points: + pc_aug = np.load(str(random_pc_path)) + pc_aug = torch.from_numpy(pc_aug).float() + pc_aug = pc_aug[np.random.choice(pc_aug.shape[0], replace=False, size=800)] + pc = torch.cat((pc, pc_aug), dim=0) + else: + assert self.phase == "real" + + # assert that the point cloud is normalized + max_diff = 0.05 + if self.random_pc: + max_diff = 0.3 + if not self.gaussian or self.gaussian == 0.0: + max_dist_from_center = abs(1.0 - torch.max(torch.sqrt(torch.sum((pc ** 2), dim=1)))) + assert max_dist_from_center < max_diff, f"Point cloud is not normalized [{max_dist_from_center} > {max_diff}] for sample [{file_name}]. If this is an external ds, please consider using prepare_coseg.py script first." + + # load target vectors, for test phase, some examples may not have a yml file attached to the + yml_path = self.ds_path.joinpath(self.yml_gt_normalized_dir_name, f"{file_name}.yml") + yml_obj = None + if yml_path.is_file(): + with open(yml_path, 'r') as f: + yml_obj = yaml.load(f, Loader=yaml.FullLoader) + else: + # for training and validation we must have a yml file for each sample, for certain phases, yml file is not mandatory + assert self.phase == "coseg" or self.phase == "real" + + # assemble the vectors in the requested order of parameters + targets = assemble_targets(yml_obj, self.inputs_to_eval) + + # dataloaders are not allowed to return None, anything empty is converted to [] + return file_name, pc, targets, yml_obj if yml_obj else [] + + def __len__(self): + return self.size diff --git a/data/dataset_sketch.py b/data/dataset_sketch.py new file mode 100644 index 0000000..3469303 --- /dev/null +++ b/data/dataset_sketch.py @@ -0,0 +1,119 @@ +import numpy as np +import torch.utils.data as data +import yaml +from pathlib import Path +from .data_processing import normalize_labels +from torchvision import transforms +from PIL import Image +from skimage.morphology import erosion, dilation +import random +from .dataset_util import assemble_targets + + +class DatasetSketch(data.Dataset): + def __init__(self, + inputs_to_eval, + params_descriptors, + camera_angles_to_process, + pretrained_vgg, + data_dir, + phase, + train_with_visibility_label=True): + self.inputs_to_eval = inputs_to_eval + self.data_dir = data_dir + self.phase = phase + self.pretrained_vgg = pretrained_vgg + self.train_with_visibility_label = train_with_visibility_label + self.camera_angles_to_process = camera_angles_to_process + self.num_sketches_camera_angles = len(self.camera_angles_to_process) + self.yml_gt_normalized_dir_name = 'yml_gt_normalized' + self.ds_path = Path(data_dir, phase) + if not self.ds_path.is_dir(): + raise Exception(f"Could not find a dataset in path [{self.ds_path}]") + self.sketches_path = self.ds_path.joinpath("sketches") + if not self.sketches_path.is_dir(): + raise Exception(f"Could not find a sketches in path [{self.sketches_path}]") + self.sketch_transforms = transforms.Compose([ + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)), + ]) + normalize_labels(data_dir, phase, self.yml_gt_normalized_dir_name, params_descriptors, self.train_with_visibility_label) + + obj_gt_dir = self.ds_path.joinpath('obj_gt') + self.file_names = [f.stem for f in obj_gt_dir.glob("*.obj")] + if self.phase == "real" or self.phase == "clipasso" or self.phase == "traced": + self.file_names = [f.stem for f in self.sketches_path.glob("*.png")] + + num_files = len(self.file_names) + if self.phase == "real" or self.phase == "clipasso" or self.phase == "traced": + self.size = num_files + else: + self.size = num_files * self.num_sketches_camera_angles + + def __getitem__(self, _index): + if self.phase == "real" or self.phase == "clipasso" or self.phase == "traced": + file_idx = _index + sketch_idx = 0 + else: + file_idx = _index // self.num_sketches_camera_angles + sketch_idx = _index % self.num_sketches_camera_angles + file_name = self.file_names[file_idx] + + # load target vectors, for test phase, some examples may not have a yml file attached to them + yml_path = self.ds_path.joinpath(self.yml_gt_normalized_dir_name, f"{file_name}.yml") + yml_obj = None + if yml_path.is_file(): + with open(yml_path, 'r') as f: + yml_obj = yaml.load(f, Loader=yaml.FullLoader) + else: + # for training and validation we must have a yml file for each sample, for certain phases, yml file is not mandatory + assert self.phase == "test" or self.phase == "coseg" or self.phase == "real" or self.phase == "clipasso" or self.phase == "traced" + + # assemble the vectors in the requested order of parameters + targets = assemble_targets(yml_obj, self.inputs_to_eval) + + sketch_files = sorted(self.sketches_path.glob(f"{file_name}_*.png")) + if self.phase == "real" or self.phase == "clipasso" or self.phase == "traced": + sketch_files = sorted(self.sketches_path.glob(f"{file_name}.png")) + # filter out sketches from camera angles that are excluded + if self.phase != "real" and self.phase != "clipasso" and self.phase != "traced": + sketch_files = [f for f in sketch_files if any( camera_angle in f.name for camera_angle in self.camera_angles_to_process )] + if len(sketch_files) != len(self.camera_angles_to_process): + raise Exception(f"Object [{file_name}] is missing sketch files") + sketch_file = sketch_files[sketch_idx] + sketch = Image.open(sketch_file).convert("RGB") + if sketch.size[0] != sketch.size[0]: + raise Exception(f"Images should be square, got [{sketch.size}] instead.") + if sketch.size[0] != 224: + sketch = sketch.resize((224, 224), Image.BILINEAR) + # augmentation for the sketches + if self.phase == "train": + # three augmentation options: 1) original 2) erosion 3) erosion then dilation + aug_idx = random.randint(0, 2) + if aug_idx == 1: + sketch = np.array(sketch) + sketch = erosion(sketch) + sketch = Image.fromarray(sketch) + if aug_idx == 2: + sketch = np.array(sketch) + eroded = erosion(sketch) + sketch = dilation(eroded) + sketch = Image.fromarray(sketch) + sketch = self.sketch_transforms(sketch) + if not self.pretrained_vgg: + sketch = sketch[0].unsqueeze(0) # sketch.shape = [1, 224, 224] + + curr_file_camera_angle = 'angle_na' + for camera_angle in self.camera_angles_to_process: + if camera_angle in str(sketch_file): + curr_file_camera_angle = camera_angle + break + if self.phase != "real" and self.phase != "clipasso" and self.phase != "traced": + assert curr_file_camera_angle != 'angle_na' + + # dataloaders are not allowed to return None, anything empty is converted to [] + return file_name, curr_file_camera_angle, sketch, targets, yml_obj if yml_obj else [] + + def __len__(self): + return self.size diff --git a/data/dataset_util.py b/data/dataset_util.py new file mode 100644 index 0000000..e26da24 --- /dev/null +++ b/data/dataset_util.py @@ -0,0 +1,21 @@ +import torch +import numpy as np +from typing import Dict, List + + +def assemble_targets(yml_obj: Dict, inputs_to_eval: List[str]): + targets = [] + if yml_obj: + for param_name in inputs_to_eval: + if param_name[-2:] == ' x': + targets.append(yml_obj[param_name[:-2]]['x']) + elif param_name[-2:] == ' y': + targets.append(yml_obj[param_name[:-2]]['y']) + elif param_name[-2:] == ' z': + targets.append(yml_obj[param_name[:-2]]['z']) + else: + targets.append(yml_obj[param_name]) + + # convert from list to numpy array and then to torch tensor + targets = torch.from_numpy(np.asarray(targets)) + return targets diff --git a/dataset_generator/__init__.py b/dataset_generator/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/dataset_generator/dataset_generator.py b/dataset_generator/dataset_generator.py new file mode 100644 index 0000000..702c4ee --- /dev/null +++ b/dataset_generator/dataset_generator.py @@ -0,0 +1,439 @@ +#!/usr/bin/env python3 + +import sys +import bpy +import time +import yaml +import json +import hashlib +import argparse +import traceback +import subprocess +from pathlib import Path +import importlib + +def import_parents(level=1): + global __package__ + file = Path(__file__).resolve() + parent, top = file.parent, file.parents[level] + + sys.path.append(str(top)) + try: + sys.path.remove(str(parent)) + except ValueError: + # already removed + pass + + __package__ = '.'.join(parent.parts[len(top.parts):]) + importlib.import_module(__package__) + +if __name__ == '__main__' and __package__ is None: + import_parents(level=1) + +from common.param_descriptors import ParamDescriptors +from common.file_util import get_recipe_yml_obj, hash_file_name, get_source_recipe_file_path, save_yml +from common.bpy_util import refresh_obj_in_viewport, select_objs, select_shape, get_geometric_nodes_modifier, save_obj +from common.domain import Domain +from common.input_param_map import get_input_param_map, randomize_all_params, yml_to_shape +from dataset_generator.shape_validators.common_validations import object_sanity_check +from dataset_generator.shape_validators.shape_validator_factory import ShapeValidatorFactory + + +def shape_to_yml(gnodes_mod): + shape_yml_obj = {} + for input in gnodes_mod.node_group.inputs: + param_name = str(input.name) + param_val = gnodes_mod[input.identifier] + if input.bl_label == "Vector": + shape_yml_obj[param_name] = {} + shape_yml_obj[param_name]['x'] = param_val[0] + shape_yml_obj[param_name]['y'] = param_val[1] + shape_yml_obj[param_name]['z'] = param_val[2] + else: + shape_yml_obj[param_name] = param_val + return shape_yml_obj + + +def save_obj_label(gnodes_mod, targe_yml_file_path: Path): + """ + convert the object to parameter space and save it as a yaml file + """ + shape_yml_obj = shape_to_yml(gnodes_mod) + save_yml(shape_yml_obj, targe_yml_file_path) + + +def update_base_shape_in_yml(gnodes_mod, recipe_file_path: Path): + """ + This will completely overwrite the base shape in the given yml file while keeping other + fields in the yaml untouched. + If the yml file does not exist, it will generate a new base shape as yaml in the given file. + """ + print(f'Updating the base shape in the YML file [{recipe_file_path}]') + # init an empty object in case the file does not exist + recipe_yml_obj = {} + if recipe_file_path.is_file(): + recipe_yml_obj = get_recipe_yml_obj(recipe_file_path) + base_yml_obj = shape_to_yml(gnodes_mod) + # completely overwrite 'base' in the final yml object + recipe_yml_obj['base'] = base_yml_obj + # save the object as YML file + save_yml(recipe_yml_obj, recipe_file_path) + return recipe_yml_obj + + +def json_hash(json_obj): + return hashlib.md5(json.dumps(json_obj).encode("utf-8")).hexdigest().strip() + + +def update_recipe_yml_obj_with_metadata(recipe_yml_obj, gnodes_mod): + # loops through all the inputs in the geometric node group + data_types = {} + for input in gnodes_mod.node_group.inputs: + param_name = str(input.name) + data_types[param_name] = {} + data_types[param_name]['type'] = input.bl_label + if input.bl_label != 'Boolean': + data_types[param_name]['min'] = input.min_value + data_types[param_name]['max'] = input.max_value + recipe_yml_obj['data_types'] = data_types + + +def generate_dataset(domain, dataset_dir: Path, phase, random_shapes_per_value, parallel=1, mod=None): + """ + Params: + random_shapes_per_value - number of random shapes we will generate per parameter value + mod - allows running this in parallel + """ + try: + # all other processes must wait for the folder to be created before continuing + phase_dir = dataset_dir.joinpath(phase) + yml_gt_dir = phase_dir.joinpath('yml_gt') + obj_gt_dir = phase_dir.joinpath('obj_gt') + if parallel > 1 and mod != 0: + while not (dataset_dir.is_dir() and phase_dir.is_dir() and yml_gt_dir.is_dir() and obj_gt_dir.is_dir()): + time.sleep(2) + + dataset_dir.mkdir(exist_ok=True) + phase_dir.mkdir(exist_ok=True) + + obj = select_shape() + # get the geometric nodes modifier for the object + gnodes_mod = get_geometric_nodes_modifier(obj) + recipe_file_path = get_source_recipe_file_path(domain) + if parallel <= 1: + update_base_shape_in_yml(gnodes_mod, recipe_file_path) + + # load recipe file and add some required metadata to it + recipe_yml_obj = get_recipe_yml_obj(recipe_file_path) + base_shape_yml = recipe_yml_obj['base'].copy() # used to return the viewport shape to this base shape at the end of the dataset generation + update_recipe_yml_obj_with_metadata(recipe_yml_obj, gnodes_mod) + + if parallel <= 1: + # save the recipe object as yml file in the dataset main dir (since it now also contains additional required metadata) + target_recipe_file_path = dataset_dir.joinpath('recipe.yml') + save_yml(recipe_yml_obj, target_recipe_file_path) + + yml_gt_dir.mkdir(exist_ok=True) + obj_gt_dir.mkdir(exist_ok=True) + + input_params_map = get_input_param_map(gnodes_mod, recipe_yml_obj) + inputs_to_eval = list(input_params_map.keys()) + param_descriptors = ParamDescriptors(recipe_yml_obj, inputs_to_eval) + param_descriptors_map = param_descriptors.get_param_descriptors_map() + dup_hashes_attempts = [] + + shape_validator = ShapeValidatorFactory.create_validator(domain) + + existing_samples = {} + for curr_param_name, curr_input_param in input_params_map.items(): + for value_idx, curr_param_value in enumerate(curr_input_param.possible_values): + shape_idx = 0 + while shape_idx < random_shapes_per_value: + if parallel > 1: + if hash_file_name(f'{curr_param_name}_{value_idx}_{shape_idx}') % parallel != mod: + shape_idx += 1 + continue + + curr_param_value_str_for_file = f"{curr_param_value:.4f}".replace('.', '_') + file_name = f"{domain}_{curr_input_param.get_name_for_file()}_{curr_param_value_str_for_file}_{shape_idx:04d}" + obj_file = obj_gt_dir.joinpath(f"{file_name}.obj") + yml_file = yml_gt_dir.joinpath(f"{file_name}.yml") + if obj_file.is_file() and yml_file.is_file() and object_sanity_check(obj_file): + with open(yml_file, 'r') as file: + param_values_map_from_yml = yaml.load(file, Loader=yaml.FullLoader) + param_values_map = {} + for param_name, _ in input_params_map.items(): + if param_name[-2:] in [' x', ' y', ' z']: + param_values_map[param_name] = param_values_map_from_yml[param_name[:-2]][param_name[-1:]] + else: + param_values_map[param_name] = param_values_map_from_yml[param_name] + + shape_yml = {} + for param_name, param_value in param_values_map.items(): + if not param_descriptors_map[param_name].is_visible(param_values_map): + shape_yml[param_name] = -1 + else: + shape_yml[param_name] = round(param_value, 4) + sample_hash = json_hash(shape_yml) + if sample_hash in existing_samples: + raise Exception("Found a duplicate within a single process") + existing_samples[sample_hash] = file_name + shape_idx += 1 + continue + + param_values_map = randomize_all_params(input_params_map) + param_values_map[curr_param_name] = curr_param_value + + if not param_descriptors.check_constraints(param_values_map): + with open(f'./retry_{mod}.log', 'a') as f: + f.write(f'constraints {file_name}\n') + continue + + if not param_descriptors_map[curr_param_name].is_visible(param_values_map): + with open(f'./retry_{mod}.log', 'a') as f: + f.write(f'visibility conditions {file_name} [{param_descriptors_map[curr_param_name].visibility_condition}] \n') + continue + + for param_name, param_value in param_values_map.items(): + input_params_map[param_name].assign_value(param_value) + + # sanity check to make sure we did not override what we are trying to generate + assert abs(input_params_map[curr_param_name].get_value() - curr_param_value) < 1e-6 + # must refresh the shape at this point + refresh_obj_in_viewport(obj) + + # shape-specific validation + is_valid, msg = shape_validator.validate_shape(input_params_map) + if not is_valid: + with open(f'./retry_{mod}.log', 'a') as f: + f.write(f'Shape invalid with message [{msg}] for file {file_name}\n') + continue + + # make sure this is a completely new shape + shape_yml = {} + for param_name, param_value in param_values_map.items(): + if not param_descriptors_map[param_name].is_visible(param_values_map): + shape_yml[param_name] = -1 + else: + shape_yml[param_name] = round(input_params_map[param_name].get_value(), 4) + sample_hash = json_hash(shape_yml) + if sample_hash in existing_samples: + dup_hashes_attempts.append(sample_hash) + with open(f'./retry_{mod}.log', 'a') as f: + f.write(f'already exists {sample_hash} {file_name}\n') + continue + existing_samples[sample_hash] = file_name + + targe_yml_file_path = yml_gt_dir.joinpath(f"{file_name}.yml") + save_obj_label(gnodes_mod, targe_yml_file_path) + targe_obj_file_path = obj_gt_dir.joinpath(f"{file_name}.obj") + dup_obj = save_obj(targe_obj_file_path) + # delete the duplicate object + select_objs(dup_obj) + bpy.ops.object.delete() + shape_idx += 1 + + # return the shape in Blender viewport to its original state + yml_to_shape(base_shape_yml, input_params_map, ignore_sanity_check=True) + dup_hashes_attempts_file_path = dataset_dir.joinpath(f"dup_hashes_attempts_{mod}.txt") + with open(dup_hashes_attempts_file_path, 'a') as dup_hashes_attempts_file: + dup_hashes_attempts_file.writelines([f"{h}\n" for h in dup_hashes_attempts]) + dup_hashes_attempts_file.write('---\n') + + return existing_samples + + except Exception as e: + with open(f'./err_{mod}.log', 'a') as f: + f.write(repr(e)) + f.write('\n') + f.write(traceback.format_exc()) + f.write('\n\n') + + +def main_generate_dataset_single_proc(args, blender_exe, blend_file): + assert blender_exe + assert blend_file + # show the main collections (if it is already shown, there is no effect) + bpy.context.scene.view_layers['View Layer'].layer_collection.children['Main'].hide_viewport = False + bpy.context.scene.view_layers['View Layer'].layer_collection.children['Main'].exclude = False + + try: + dataset_dir = Path(args.dataset_dir).expanduser() + existing_samples = generate_dataset(args.domain, dataset_dir, args.phase, + args.num_variations, parallel=args.parallel, mod=args.mod) + samples_hashes_file_path = dataset_dir.joinpath(f"sample_hashes_{args.mod}.json") + with open(samples_hashes_file_path, 'w') as samples_hashes_file: + json.dump(existing_samples, samples_hashes_file) + print(f"Process [{args.mod}] done") + except Exception as e: + print(repr(e)) + print(traceback.format_exc()) + + +def main_generate_dataset_parallel(args, blender_exe, blend_file): + dataset_dir = Path(args.dataset_dir).expanduser() + dataset_dir.mkdir(exist_ok=True) + + phase_dir = dataset_dir.joinpath(args.phase) + phase_dir.mkdir(exist_ok=True) + + try: + for existing_shapes_json_file_path in dataset_dir.glob("sample_hashes_*.json"): + existing_shapes_json_file_path.unlink() + + # select the procedural shape in blender + obj = select_shape() + # get the geometric nodes modifier fo the object + gnodes_mod = get_geometric_nodes_modifier(obj) + recipe_file_path = get_source_recipe_file_path(args.domain) + recipe_yml_obj = get_recipe_yml_obj(str(recipe_file_path)) + update_base_shape_in_yml(gnodes_mod, recipe_file_path) + update_recipe_yml_obj_with_metadata(recipe_yml_obj, gnodes_mod) + # save the recipe.yml file in the dataset's main dir (it now also contains required metadata) + target_recipe_file_path = dataset_dir.joinpath('recipe.yml') + save_yml(recipe_yml_obj, target_recipe_file_path) + input_params_map = get_input_param_map(gnodes_mod, recipe_yml_obj) + # loops through all the inputs in the geometric node group + data_types = {} + for input in gnodes_mod.node_group.inputs: + param_name = str(input.name) + data_types[param_name] = {} + data_types[param_name]['type'] = input.bl_label + if input.bl_label != 'Boolean': + data_types[param_name]['min'] = input.min_value + data_types[param_name]['max'] = input.max_value + recipe_yml_obj['data_types'] = data_types + inputs_to_eval = list(input_params_map.keys()) + param_descriptors = ParamDescriptors(recipe_yml_obj, inputs_to_eval) + + expected_number_of_samples = param_descriptors.get_overall_num_of_classes_without_visibility_label() * args.num_variations + print(f"Overall expected number of objects to generate is [{expected_number_of_samples}]") + + iteration_count = 0 + while True: + iteration_count += 1 + duplicates = [] + + # load any current saved state (this is to generate val, test, and train, since we cannot do this simultaneously) + # please note that all the sample_hashes files are common to all phases of the dataset (val, test, and train) so we + # will avoid creating the same sample across all of these phases + + existing_samples = {} + # add all the samples hashes from any other phase to avoid duplicates with other phases + sample_hashes_json_file_path = dataset_dir.joinpath("sample_hashes.json") + if sample_hashes_json_file_path.is_file(): + with open(sample_hashes_json_file_path, 'r') as existing_samples_file: + existing_samples = json.load(existing_samples_file) + # clear any current phase sample hashes as they are added by the processes + existing_samples[args.phase] = {} + + # every process generates a 'sample_hashes_.json' file containing hash -> file_name map + for existing_samples_json_file_path in dataset_dir.glob("sample_hashes_*.json"): + with open(existing_samples_json_file_path, 'r') as existing_samples_json_file: + existing_samples_json = json.load(existing_samples_json_file) + print(existing_samples_json_file_path) + print(existing_samples_json) + for hash, file_name in existing_samples_json.items(): + is_dup = False + for phase, sample_hashes in existing_samples.items(): + if hash in sample_hashes: + duplicates.append(file_name) + is_dup = True + break + if not is_dup: + existing_samples[args.phase][hash] = file_name + + # delete the duplicated files so they will be regenerated + if duplicates: + print(f"Found [{len(duplicates)}] duplicates that will be regenerated") + print("\n\t".join(duplicates)) + for file_name in duplicates: + obj_file = phase_dir.joinpath(file_name + ".obj") + yml_file = phase_dir.joinpath(file_name + ".yml") + obj_file.unlink() + yml_file.unlink() + + # backup the current sample_hashes + with open(f"{dataset_dir}/sample_hashes.json", 'w') as sample_hashes_file: + json.dump(existing_samples, sample_hashes_file) + + if len(existing_samples[args.phase]) > expected_number_of_samples: + raise Exception("Something went wrong, make sure you know how to count") + print(len(existing_samples[args.phase])) + print(expected_number_of_samples) + if len(existing_samples[args.phase]) == expected_number_of_samples: + print('Done creating the requested dataset') + break + + # since we need another iteration, we remove all the sample_hashes files + for existing_samples_json_file_path in dataset_dir.glob("sample_hashes_*.json"): + existing_samples_json_file_path.unlink() + + processes = [] + dataset_generator_path = Path(__file__).parent.joinpath('dataset_generator.py').resolve() + for mod in range(args.parallel): + try: + cmd = [str(blender_exe), str(blend_file), '-b', '--python', str(dataset_generator_path), '--', + 'generate-dataset-single-process', + '--dataset-dir', str(dataset_dir), + '--domain', str(args.domain), + '--phase', args.phase, + '--num-variations', str(args.num_variations), + '--parallel', str(args.parallel), + '--mod', str(mod)] + print(f'Mod {mod}:') + print(" ".join(cmd)) + process = subprocess.Popen(cmd, stdout=subprocess.DEVNULL) # DEVNULL is to avoid processes getting stuck + processes.append(process) + except Exception as e: + print(repr(e)) + print(traceback.format_exc()) + for process in processes: + process.wait() + + print(f"Dataset generation iteration [{iteration_count}] is done") + except Exception as e: + print(repr(e)) + print(traceback.format_exc()) + + +def main(): + if '--' in sys.argv: + # refer to https://b3d.interplanety.org/en/how-to-pass-command-line-arguments-to-a-blender-python-script-or-add-on/ + argv = sys.argv[sys.argv.index('--') + 1:] + else: + raise Exception("Expected \'--\' followed by arguments to the script") + + parser = argparse.ArgumentParser(prog='dataset_generator') + common_parser = argparse.ArgumentParser(add_help=False) + common_parser.add_argument('--dataset-dir', type=str, required=True, help='Path to dataset directory') + + sp = parser.add_subparsers() + sp_gen_single_proc = sp.add_parser('generate-dataset-single-process', parents=[common_parser]) + sp_gen_parallel = sp.add_parser('generate-dataset', parents=[common_parser]) + + sp_gen_single_proc.set_defaults(func=main_generate_dataset_single_proc) + sp_gen_single_proc.add_argument('--domain', type=Domain, choices=list(Domain), required=True, help='The domain name to generate the dataset for.') + sp_gen_single_proc.add_argument('--phase', type=str, required=True, help='E.g. train, val, or test') + sp_gen_single_proc.add_argument('--num-variations', type=int, default=3) + sp_gen_single_proc.add_argument('--parallel', type=int, default=1, help='Number of processes that are running the script in parallel') + sp_gen_single_proc.add_argument('--mod', type=int, default=None, help='The modulo for this process to match files\' hash') + + sp_gen_parallel.set_defaults(func=main_generate_dataset_parallel) + sp_gen_parallel.add_argument('--domain', type=Domain, choices=list(Domain), required=True, help='The domain name to generate the dataset for.') + sp_gen_parallel.add_argument('--phase', type=str, required=True, help='E.g. train, val, or test') + sp_gen_parallel.add_argument('--num-variations', type=int, default=3, help='The number of random shapes to generate for each parameter value') + sp_gen_parallel.add_argument('--parallel', type=int, default=1, help='Number of processes that will run the script') + + blender_exe_path = Path(sys.argv[0]).resolve() + blend_file_path = Path(sys.argv[1]).resolve() + try: + args = parser.parse_known_args(argv)[0] + args.func(args, blender_exe_path, blend_file_path) + except Exception as e: + print(repr(e)) + print(traceback.format_exc()) + + +if __name__ == "__main__": + main() diff --git a/dataset_generator/recipe_files/recipe_chair.yml b/dataset_generator/recipe_files/recipe_chair.yml new file mode 100644 index 0000000..df981b2 --- /dev/null +++ b/dataset_generator/recipe_files/recipe_chair.yml @@ -0,0 +1,318 @@ +base: + scale: + x: 1.0 + y: 1.0 + z: 1.0 + bevel_rails: 0.0 + pillow_state: 1 + pillow_fill_edge: 1 + seat_shape: 0.423308789730072 + seat_pos: 0.5646687746047974 + cr_count: 5 + cr_scale_y: 0.7333 + cr_scale_z: 0.98 + cr_offset_bottom: 0.4071 + cr_offset_top: 0.7786 + cr_shape_1: 0.0 + curvature: 0.25 + is_top_rail: 1 + tr_fill_edge: 1 + tr_scale_y: 0.6800000071525574 + tr_scale_z: 1.3600000143051147 + tr_shape_1: 0.5525987148284912 + is_vertical_rail: 1 + vr_count: 4 + vr_scale_x: 0.7799999713897705 + vr_scale_y: 0.3499999940395355 + vr_shape_1: 0.0 + is_back_rest: 0 + legs_shape_1: 1.0 + legs_shape_2: 1.0 + legs_bevel: 0.5 + is_monoleg: 0 + is_monoleg_tent: 0 + monoleg_tent_pct: 0.4 + monoleg_tent_count: 4 + monoleg_bezier_start_x_offset: 0.0 + monoleg_bezier_start_handle_x_offset: 0.0 + monoleg_bezier_start_handle_z_pct: 0.4 + monoleg_bezier_end_x_offset: 0.8199999332427979 + monoleg_bezier_end_handle_x_offset: 0.4 + monoleg_bezier_end_handle_z_pct: 1.0 + back_frame_top_y_offset_pct: 0.0 + back_frame_mid_y_offset_pct: 0.0 + back_leg_bottom_y_offset_pct: 0.0 + back_leg_mid_y_offset_pct: 0.0 + handles_state: 0 + is_handles_support: 1 + is_handles_cusion: 1 + handles_base_pos_z_pct: 0.15 + handles_mid_pos_x_pct: 0.0 + handles_mid_pos_y_pct: 0.42329999804496765 + handles_mid_pos_z_pct: 0.5 + handles_edge_pos_x_pct: 0.0 + handles_bottom_pos_along_seat_pct: 0.4 + handles_profile_width: 0.9 + handles_profile_height: 0.9 + handles_support_mid_x: 0.5800000429153442 + handles_support_mid_y: 0.0 + handles_support_top_pos: 0.5 + handles_support_thickness: 0.8167 + handles_cusion_cover_pct: 1.0 +dataset_generation: + scale: + x: + min: 0.5 + max: 2.0 + samples: 10 + y: + min: 0.5 + max: 2.0 + samples: 10 + z: + min: 0.5 + max: 2.0 + samples: 10 + bevel_rails: + min: 0.0 + max: 1.0 + samples: 3 + pillow_state: + min: 0 + max: 2 + pillow_fill_edge: + min: 0 + max: 1 + seat_shape: + min: 0.0 + max: 1.0 + samples: 5 + seat_pos: + min: 0.2 + max: 1.0 + samples: 9 + cr_count: + min: 3 + max: 8 + cr_scale_y: + min: 0.5 + max: 1.2 + samples: 4 + cr_scale_z: + min: 0.3 + max: 2.0 + samples: 6 + cr_offset_bottom: + min: 0.15 + max: 0.75 + samples: 8 + cr_offset_top: + min: 0.35 + max: 0.95 + samples: 8 + cr_shape_1: + min: 0.0 + max: 2.0 + samples: 7 + curvature: + min: 0.0 + max: 1.0 + samples: 5 + is_top_rail: + min: 0 + max: 1 + tr_fill_edge: + min: 0 + max: 1 + tr_scale_y: + min: 0.5 + max: 1.2 + samples: 5 + tr_scale_z: + min: 0.5 + max: 1.5 + samples: 5 + tr_shape_1: + min: 0.0 + max: 1.0 + samples: 5 + is_vertical_rail: + min: 0 + max: 1 + vr_count: + min: 3 + max: 8 + vr_scale_x: + min: 0.3 + max: 2.0 + samples: 6 + vr_scale_y: + min: 0.2 + max: 1.0 + samples: 4 + vr_shape_1: + min: 0.0 + max: 0.4 + samples: 5 + is_back_rest: + min: 0 + max: 1 + legs_shape_1: + min: 0.0 + max: 1.0 + samples: 3 + legs_shape_2: + min: 0.0 + max: 1.0 + samples: 3 + legs_bevel: + min: 0.0 + max: 1.0 + samples: 3 + is_monoleg: + min: 0 + max: 1 + is_monoleg_tent: + min: 0 + max: 1 + monoleg_tent_pct: + min: 0.2 + max: 0.8 + samples: 7 + monoleg_tent_count: + min: 3 + max: 8 + monoleg_bezier_start_x_offset: + min: 0.0 + max: 1.0 + samples: 6 + monoleg_bezier_start_handle_x_offset: + min: 0.0 + max: 1.0 + samples: 6 + monoleg_bezier_start_handle_z_pct: + min: 0.0 + max: 1.0 + samples: 6 + monoleg_bezier_end_x_offset: + min: 0.0 + max: 1.0 + samples: 6 + monoleg_bezier_end_handle_x_offset: + min: 0.0 + max: 1.0 + samples: 6 + monoleg_bezier_end_handle_z_pct: + min: 0.0 + max: 1.0 + samples: 6 + back_frame_top_y_offset_pct: + min: 0.0 + max: 1.0 + samples: 6 + back_frame_mid_y_offset_pct: + min: 0.0 + max: 1.0 + samples: 6 + back_leg_bottom_y_offset_pct: + min: 0.0 + max: 1.0 + samples: 6 + back_leg_mid_y_offset_pct: + min: 0.0 + max: 1.0 + samples: 6 + handles_state: + min: 0 + max: 2 + is_handles_support: + min: 0 + max: 1 + is_handles_cusion: + min: 0 + max: 1 + handles_profile_width: + min: 0.5 + max: 1.0 + samples: 6 + handles_profile_height: + min: 0.5 + max: 1.0 + samples: 6 + handles_base_pos_z_pct: + min: 0.15 + max: 0.8 + samples: 8 + handles_mid_pos_x_pct: + min: 0.0 + max: 1.0 + samples: 7 + handles_mid_pos_y_pct: + min: 0.1 + max: 0.9 + samples: 7 + handles_mid_pos_z_pct: + min: 0.0 + max: 1.0 + samples: 7 + handles_edge_pos_x_pct: + min: 0.0 + max: 1.0 + samples: 6 + handles_bottom_pos_along_seat_pct: + min: 0.3 + max: 0.9 + samples: 7 + handles_support_mid_x: + min: 0.0 + max: 1.0 + samples: 6 + handles_support_mid_y: + min: 0.0 + max: 1.0 + samples: 6 + handles_support_top_pos: + min: 0.3 + max: 0.9 + samples: 7 + handles_support_thickness: + min: 0.4 + max: 0.9 + samples: 7 + handles_cusion_cover_pct: + min: 0.0 + max: 1.0 + samples: 6 +constraints: + rule1: cr_offset_top - cr_offset_bottom >= 0.1 + rule2: monoleg_bezier_end_x_offset >= monoleg_bezier_start_x_offset +visibility_conditions: + bevel_rails: ( not is_back_rest ) + cr_: ( not is_back_rest ) and ( not is_top_rail or not is_vertical_rail ) + vr_: ( not is_back_rest ) and is_top_rail and is_vertical_rail + tr_: is_top_rail + is_vertical_rail: ( not is_back_rest ) and ( is_top_rail ) + legs_shape_: ( is_monoleg and is_monoleg_tent ) or ( not is_monoleg ) + is_monoleg_tent: is_monoleg + monoleg_tent_pct: is_monoleg and is_monoleg_tent + monoleg_tent_count: is_monoleg and is_monoleg_tent + monoleg_bezier: is_monoleg + pillow_fill_edge: pillow_state > 0 + back_leg_: ( not is_monoleg ) + is_handles_support: handles_state == 1 + handles_support_: ( handles_state == 1 ) and ( is_handles_support ) + is_handles_cusion: handles_state > 0 + handles_cusion_cover_pct: ( handles_state > 0 ) and ( is_handles_cusion ) + handles_profile_: handles_state > 0 + handles_base_: handles_state > 0 + handles_mid_: handles_state > 0 + handles_edge_: handles_state == 1 + handles_bottom_pos_along_seat_pct: ( handles_state == 2 ) or ( handles_state == 1 and is_handles_support == 1 ) +camera_angles_train: +- - -30.0 + - 35.0 +- - -30.0 + - 55.0 +camera_angles_test: +- - -30.0 + - 15.0 diff --git a/dataset_generator/recipe_files/recipe_table.yml b/dataset_generator/recipe_files/recipe_table.yml new file mode 100644 index 0000000..9d11255 --- /dev/null +++ b/dataset_generator/recipe_files/recipe_table.yml @@ -0,0 +1,198 @@ +base: + table_top_scale_x: 2.5999999046325684 + table_top_scale_y: 2.5999999046325684 + table_top_height: 0.0 + table_top_shape: 1.0 + table_top_thickness: 0.0 + table_top_profile_state: 2 + table_top_profile_strength: 0.0 + legs_shape_1: 1.0 + legs_shape_2: 1.0 + legs_bevel: 0.0 + std_legs_bottom_offset_y: 1.0 + std_legs_mid_offset_y: 0.0 + std_legs_top_offset_x: 0.26999998092651367 + std_legs_top_offset_y: 0.0 + std_legs_rotation: 0.0 + is_std_legs_support_x: 1 + std_legs_support_x_height: 0.7200000286102295 + std_legs_support_x_curvature: 0.0 + std_legs_support_x_profile_width: 1.0 + std_legs_support_x_profile_height: 0.27000004053115845 + is_std_legs_support_y: 1 + std_legs_support_y_height: 0.36000001430511475 + std_legs_support_y_curvature: 0.0 + std_legs_support_y_profile_width: 1.0 + std_legs_support_y_profile_height: 1.0 + is_monoleg: 0 + is_monoleg_tent: 1 + monoleg_tent_pct: 0.6649518609046936 + monoleg_tent_base_radius: 0.0 + monoleg_tent_count: 5 + monoleg_bezier_start_x_offset: 0.5273312330245972 + monoleg_bezier_start_handle_x_offset: 1.0 + monoleg_bezier_start_handle_z_pct: 0.46000000834465027 + monoleg_bezier_end_x_offset: 0.2499999850988388 + monoleg_bezier_end_handle_x_offset: 0.10999999940395355 + monoleg_bezier_end_handle_z_pct: 0.20999999344348907 +dataset_generation: + table_top_scale_x: + min: 0.6 + max: 2.6 + samples: 12 + table_top_scale_y: + min: 0.6 + max: 2.6 + samples: 12 + table_top_height: + min: 0.0 + max: 1.0 + samples: 8 + table_top_shape: + min: 0.0 + max: 1.0 + samples: 11 + table_top_thickness: + min: 0.0 + max: 1.0 + samples: 6 + table_top_profile_state: + min: 0 + max: 3 + table_top_profile_strength: + min: 0.0 + max: 1.0 + samples: 6 + legs_shape_1: + min: 0.0 + max: 1.0 + samples: 3 + legs_shape_2: + min: 0.0 + max: 1.0 + samples: 3 + legs_bevel: + min: 0.0 + max: 1.0 + samples: 3 + std_legs_bottom_offset_y: + min: 0.0 + max: 1.0 + samples: 6 + std_legs_mid_offset_y: + min: 0.0 + max: 1.0 + samples: 6 + std_legs_top_offset_x: + min: 0.0 + max: 1.0 + samples: 6 + std_legs_top_offset_y: + min: 0.0 + max: 1.0 + samples: 6 + std_legs_rotation: + min: 0.0 + max: 1.0 + samples: 6 + is_std_legs_support_x: + min: 0 + max: 1 + std_legs_support_x_height: + min: 0.0 + max: 1.0 + samples: 6 + std_legs_support_x_curvature: + min: 0.0 + max: 1.0 + samples: 6 + std_legs_support_x_profile_width: + min: 0.0 + max: 1.0 + samples: 5 + std_legs_support_x_profile_height: + min: 0.0 + max: 1.0 + samples: 5 + is_std_legs_support_y: + min: 0 + max: 1 + std_legs_support_y_height: + min: 0.0 + max: 1.0 + samples: 6 + std_legs_support_y_curvature: + min: 0.0 + max: 1.0 + samples: 6 + std_legs_support_y_profile_width: + min: 0.0 + max: 1.0 + samples: 5 + std_legs_support_y_profile_height: + min: 0.0 + max: 1.0 + samples: 5 + is_monoleg: + min: 0 + max: 1 + is_monoleg_tent: + min: 0 + max: 1 + monoleg_tent_pct: + min: 0.2 + max: 0.8 + samples: 7 + monoleg_tent_base_radius: + min: 0.0 + max: 1.0 + samples: 11 + monoleg_tent_count: + min: 3 + max: 8 + monoleg_bezier_start_x_offset: + min: 0.0 + max: 1.0 + samples: 6 + monoleg_bezier_start_handle_x_offset: + min: 0.0 + max: 1.0 + samples: 6 + monoleg_bezier_start_handle_z_pct: + min: 0.0 + max: 1.0 + samples: 6 + monoleg_bezier_end_x_offset: + min: 0.2 + max: 1.0 + samples: 5 + monoleg_bezier_end_handle_x_offset: + min: 0.0 + max: 1.0 + samples: 6 + monoleg_bezier_end_handle_z_pct: + min: 0.0 + max: 1.0 + samples: 6 +constraints: + rule1: monoleg_bezier_end_x_offset >= monoleg_bezier_start_x_offset +visibility_conditions: + legs_shape_: ( is_monoleg and is_monoleg_tent ) or ( not is_monoleg ) + legs_bevel: ( is_monoleg and is_monoleg_tent ) or ( not is_monoleg ) + is_monoleg_tent: is_monoleg + monoleg_tent_pct: is_monoleg and is_monoleg_tent + monoleg_tent_count: is_monoleg and is_monoleg_tent + monoleg_tent_base_radius: is_monoleg and is_monoleg_tent + monoleg_bezier: is_monoleg and ( not is_monoleg_tent or monoleg_tent_pct < 0.4 ) + std_legs_support_x_: ( not is_monoleg ) and is_std_legs_support_x + std_legs_support_y_: ( not is_monoleg ) and is_std_legs_support_y + std_legs: ( not is_monoleg ) + table_top_profile_strength: table_top_profile_state > 0 +camera_angles_train: +- - -30.0 + - 35.0 +- - -30.0 + - 55.0 +camera_angles_test: +- - -30.0 + - 15.0 diff --git a/dataset_generator/recipe_files/recipe_vase.yml b/dataset_generator/recipe_files/recipe_vase.yml new file mode 100644 index 0000000..22475cc --- /dev/null +++ b/dataset_generator/recipe_files/recipe_vase.yml @@ -0,0 +1,208 @@ +base: + body_height: 1.2143 + body_width: 0.2944 + body_bottom_curve_width: 0.0 + body_bottom_curve_height: 0.6 + body_mouth_width: 0.6 + body_top_curve_width: 0.2 + body_top_curve_height: 0.9 + body_profile_blend: 0.8 + has_body_thickness: 0 + body_thickness_val: 0.05999999865889549 + handle_count: 5 + hndl_type: 2 + hndl_profile_width: 1.0 + hndl_profile_height: 0.2 + hndl_profile_blend: 0.8 + hndl_base_z: 0.4699999988079071 + hndl_base_bezier_handle_angle: 0.4 + hndl_base_bezier_handle_length: 0.2 + hndl_radius_along_path: 0.0 + hndl1_top_z: 0.5 + hndl1_end_bezier_handle_angle: 0.2 + hndl1_end_bezier_handle_length: 0.6 + hndl2_end_x: 0.0 + hndl2_end_z: 0.4099999964237213 + hndl2_end_bezier_handle_x: 0.0 + hndl2_end_bezier_handle_z: 0.2 + has_neck: 1 + neck_end_x: 0.8 + neck_end_z: 0.8 + neck_end_bezier_handle_x: 0.8 + neck_end_bezier_handle_z: 0.0 + has_base: 1 + base_start_x: 0.7 + base_start_z: 0.6 + base_mid_x: 0.2 + base_mid_z: 0.8 + has_lid: 0 + has_lid_handle: 1 + lid_handle_radius: 0.02 +dataset_generation: + body_height: + min: 0.5 + max: 1.5 + samples: 15 + body_width: + min: 0.1 + max: 0.45 + samples: 10 + body_bottom_curve_width: + min: 0.0 + max: 1.0 + samples: 11 + body_bottom_curve_height: + min: 0.1 + max: 0.9 + samples: 9 + body_mouth_width: + min: 0.0 + max: 0.7 + samples: 8 + body_top_curve_width: + min: 0.0 + max: 0.9 + samples: 10 + body_top_curve_height: + min: 0.1 + max: 0.9 + samples: 9 + body_profile_blend: + min: 0.0 + max: 1.0 + samples: 11 + has_body_thickness: + min: 0 + max: 1 + body_thickness_val: + min: 0.01 + max: 0.07 + samples: 7 + handle_count: + min: 0 + max: 6 + hndl_type: + min: 1 + max: 2 + hndl_profile_width: + min: 0.0 + max: 1.0 + samples: 6 + hndl_profile_height: + min: 0.0 + max: 1.0 + samples: 6 + hndl_profile_blend: + min: 0.0 + max: 1.0 + samples: 6 + hndl_base_z: + min: 0.1 + max: 0.6 + samples: 6 + hndl_base_bezier_handle_angle: + min: 0.0 + max: 1.0 + samples: 11 + hndl_base_bezier_handle_length: + min: 0.0 + max: 1.0 + samples: 6 + hndl_radius_along_path: + min: 0.0 + max: 1.0 + samples: 11 + hndl1_top_z: + min: 0.2 + max: 0.8 + samples: 7 + hndl1_end_bezier_handle_angle: + min: 0.0 + max: 1.0 + samples: 11 + hndl1_end_bezier_handle_length: + min: 0.0 + max: 1.0 + samples: 6 + hndl2_end_x: + min: 0.0 + max: 1.0 + samples: 11 + hndl2_end_z: + min: 0.0 + max: 1.0 + samples: 11 + hndl2_end_bezier_handle_x: + min: 0.0 + max: 1.0 + samples: 11 + hndl2_end_bezier_handle_z: + min: 0.1 + max: 1.0 + samples: 10 + has_neck: + min: 0 + max: 1 + neck_end_x: + min: 0.1 + max: 1.0 + samples: 10 + neck_end_z: + min: 0.0 + max: 1.0 + samples: 11 + neck_end_bezier_handle_x: + min: 0.0 + max: 1.0 + samples: 11 + neck_end_bezier_handle_z: + min: 0.0 + max: 1.0 + samples: 11 + has_base: + min: 0 + max: 1 + base_start_x: + min: 0.0 + max: 1.0 + samples: 11 + base_start_z: + min: 0.0 + max: 1.0 + samples: 6 + base_mid_x: + min: 0.0 + max: 1.0 + samples: 11 + base_mid_z: + min: 0.0 + max: 1.0 + samples: 6 + has_lid: + min: 0 + max: 1 + has_lid_handle: + min: 0 + max: 1 + lid_handle_radius: + min: 0.02 + max: 0.07 + samples: 6 +visibility_conditions: + body_thickness_val: has_body_thickness and not has_lid + hndl_: handle_count > 0 + hndl1_: handle_count > 0 and hndl_type == 1 + hndl2_: handle_count > 0 and hndl_type == 2 + neck_: has_neck + base_start_: has_base + base_mid_: has_base + has_lid_handle: has_lid + lid_handle_radius: has_lid and has_lid_handle +camera_angles_train: +- - -30.0 + - 35.0 +- - -30.0 + - 55.0 +camera_angles_test: +- - -30.0 + - 15.0 diff --git a/dataset_generator/shape_validators/__init__.py b/dataset_generator/shape_validators/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/dataset_generator/shape_validators/chair_validator.py b/dataset_generator/shape_validators/chair_validator.py new file mode 100644 index 0000000..501e85c --- /dev/null +++ b/dataset_generator/shape_validators/chair_validator.py @@ -0,0 +1,41 @@ +from dataset_generator.shape_validators.shape_validator_interface import ShapeValidatorInterface +from dataset_generator.shape_validators.common_validations import validate_monoleg +from common.intersection_util import find_self_intersections, find_cross_intersections + + +class ChairValidator(ShapeValidatorInterface): + def validate_shape(self, input_params_map) -> (bool, str): + if not self.is_valid_chair(input_params_map): + return False, "Collision" + return True, "Valid" + + def is_valid_chair(self, input_params_map): + if input_params_map['is_back_rest'].get_value() == 0 \ + and input_params_map['is_top_rail'].get_value() == 1 \ + and input_params_map['is_vertical_rail'].get_value() == 1: + # reaching here means the vertical rails are visible + # also note the assumption that the min vertical rails count is 3 + if find_self_intersections('vertical_rails_out') > 0: + # try again, as we have vertical rails intersecting each other + # print("vr, found collisions") + return False + if input_params_map['is_back_rest'].get_value() == 0 \ + and (input_params_map['is_top_rail'].get_value() == 0 + or input_params_map['is_vertical_rail'].get_value() == 0): + # reaching here means the cross rails are visible + # also note the assumption that the min cross rails count is 3 + if find_self_intersections('cross_rails_and_top_rail_out') > 0: + # try again, as we have cross rails intersecting each other or the top rail + # print("cr, found collisions...") + return False + if input_params_map['handles_state'].get_value() == 1 and input_params_map['is_handles_support'].get_value(): + if find_self_intersections('handles_support_and_back_frame') > 0: + return False + if input_params_map['handles_state'].get_value() > 0: + if find_cross_intersections('handles_left_side', 'handles_right_side') > 0: + # the handles in both sides of the chair should never intersect + return False + if input_params_map['is_monoleg'].get_value() > 0 and input_params_map['is_monoleg_tent'].get_value() == 0: + if not validate_monoleg('monoleg'): + return False + return True diff --git a/dataset_generator/shape_validators/common_validations.py b/dataset_generator/shape_validators/common_validations.py new file mode 100644 index 0000000..8168e70 --- /dev/null +++ b/dataset_generator/shape_validators/common_validations.py @@ -0,0 +1,58 @@ +import bpy +import numpy as np +from common.file_util import load_obj +from common.bpy_util import select_shape +from common.intersection_util import isolate_node_as_final_geometry + + +def triangle_area(x): + a = x[:, 0, :] - x[:, 1, :] + b = x[:, 0, :] - x[:, 2, :] + cross = np.cross(a, b) + area = 0.5 * np.norm(cross, dim=1) + return area + + +def object_sanity_check(obj_file): + try: + vertices, faces = load_obj(obj_file) + vertices = vertices.reshape(1, vertices.shape[0], vertices.shape[1]) + faces = vertices.squeeze()[faces] + triangle_area(faces) + except Exception: + print('Invalid sample') + return False + return True + + +def validate_monoleg(node_label, factor=0.08): + chair = select_shape() + revert_isolation = isolate_node_as_final_geometry(chair, node_label) + + dup_obj = chair.copy() + dup_obj.data = chair.data.copy() + dup_obj.animation_data_clear() + bpy.context.collection.objects.link(dup_obj) + # move for clarity + dup_obj.location.x += 2.0 + # set active + bpy.ops.object.select_all(action='DESELECT') + dup_obj.select_set(True) + bpy.context.view_layer.objects.active = dup_obj + # apply the modifier to turn the geometry node to a mesh + bpy.ops.object.modifier_apply(modifier="GeometryNodes") + # export the object + assert dup_obj.type == 'MESH' + + revert_isolation() + + bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY', center='BOUNDS') + center_of_volume = dup_obj.location[2] + # another option is to use (type='ORIGIN_CENTER_OF_MASS', center='MEDIAN') as the center of mass + bpy.ops.object.origin_set(type='ORIGIN_CENTER_OF_VOLUME', center='MEDIAN') + center_of_mass = dup_obj.location[2] + height = dup_obj.dimensions[2] + + if center_of_volume - center_of_mass > factor * height: + return True + return False diff --git a/dataset_generator/shape_validators/shape_validator_factory.py b/dataset_generator/shape_validators/shape_validator_factory.py new file mode 100644 index 0000000..4204e69 --- /dev/null +++ b/dataset_generator/shape_validators/shape_validator_factory.py @@ -0,0 +1,18 @@ +from common.domain import Domain +from dataset_generator.shape_validators.shape_validator_interface import ShapeValidatorInterface +from dataset_generator.shape_validators.chair_validator import ChairValidator +from dataset_generator.shape_validators.vase_validator import VaseValidator +from dataset_generator.shape_validators.table_validator import TableValidator + + +class ShapeValidatorFactory: + @staticmethod + def create_validator(domain) -> ShapeValidatorInterface: + if domain == Domain.chair: + return ChairValidator() + elif domain == Domain.vase: + return VaseValidator() + elif domain == Domain.table: + return TableValidator() + else: + raise Exception(f"Domain [{domain}] is not recognized.") diff --git a/dataset_generator/shape_validators/shape_validator_interface.py b/dataset_generator/shape_validators/shape_validator_interface.py new file mode 100644 index 0000000..e99f54f --- /dev/null +++ b/dataset_generator/shape_validators/shape_validator_interface.py @@ -0,0 +1,4 @@ +class ShapeValidatorInterface: + def validate_shape(self, input_params_map) -> (bool, str): + """validate the shape and return True if valid""" + pass diff --git a/dataset_generator/shape_validators/table_validator.py b/dataset_generator/shape_validators/table_validator.py new file mode 100644 index 0000000..e6659f0 --- /dev/null +++ b/dataset_generator/shape_validators/table_validator.py @@ -0,0 +1,18 @@ +from dataset_generator.shape_validators.shape_validator_interface import ShapeValidatorInterface +from dataset_generator.shape_validators.common_validations import validate_monoleg +from common.intersection_util import find_self_intersections + + +class TableValidator(ShapeValidatorInterface): + def validate_shape(self, input_params_map) -> (bool, str): + table_top_and_legs_support_intersections = find_self_intersections('table_top_and_legs_support') + if table_top_and_legs_support_intersections > 0: + return False, "Table top intersects with the legs supports" + floor_and_legs_support_intersections = find_self_intersections('floor_and_legs_support') + if floor_and_legs_support_intersections > 0: + return False, "Legs supports intersect with the floor" + if input_params_map['is_monoleg'].get_value() > 0 and input_params_map['is_monoleg_tent'].get_value() == 0: + if not validate_monoleg('monoleg', factor=0.16): + # the factor is more restricting since the tables can be much wider than chairs + return False, "Invalid monoleg" + return True, "Valid" diff --git a/dataset_generator/shape_validators/vase_validator.py b/dataset_generator/shape_validators/vase_validator.py new file mode 100644 index 0000000..6828d69 --- /dev/null +++ b/dataset_generator/shape_validators/vase_validator.py @@ -0,0 +1,20 @@ +from dataset_generator.shape_validators.shape_validator_interface import ShapeValidatorInterface +from common.intersection_util import find_self_intersections, find_cross_intersections + + +class VaseValidator(ShapeValidatorInterface): + def validate_shape(self, input_params_map) -> (bool, str): + body_self_intersections = find_self_intersections('Body Self Intersections') + if body_self_intersections > 0: + return False, "Self intersection in the body" + if input_params_map['handle_count'].get_value() > 0: + handle_self_intersections = find_self_intersections('Handle Self Intersections') + if handle_self_intersections > 0: + return False, "self intersection in the handle" + base_handle_intersections = find_self_intersections('Base and Handle Intersections') + if base_handle_intersections > 0: + return False, "Base intersects with handles" + floor_handle_intersections = find_self_intersections('Floor and Handle Intersections') + if floor_handle_intersections > 0: + return False, "Floor intersects with handles" + return True, "Valid" diff --git a/dataset_generator/sketch_generator.py b/dataset_generator/sketch_generator.py new file mode 100644 index 0000000..754af49 --- /dev/null +++ b/dataset_generator/sketch_generator.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 + +import sys +import traceback +import bpy +import time +from mathutils import Vector +import math +import mathutils +import random +import argparse +from tqdm import tqdm +from pathlib import Path +import importlib + +def import_parents(level=1): + global __package__ + file = Path(__file__).resolve() + parent, top = file.parent, file.parents[level] + + sys.path.append(str(top)) + try: + sys.path.remove(str(parent)) + except ValueError: + # already removed + pass + + __package__ = '.'.join(parent.parts[len(top.parts):]) + importlib.import_module(__package__) + +if __name__ == '__main__' and __package__ is None: + import_parents(level=1) + +from common.bpy_util import normalize_scale, look_at, del_obj, clean_scene +from common.file_util import get_recipe_yml_obj, hash_file_name + + +""" +Shader references: + pencil shader - https://www.youtube.com/watch?v=71KGlu_Yxtg + white background (compositing) - https://www.youtube.com/watch?v=aegiN7XeLow + creating transparent object - https://www.katsbits.com/codex/transparency-cycles/ +""" + + +def main(dataset_dir: Path, phase, parallel, mod): + try: + clean_scene() + + # setup to avoid rendering surfaces and only render the freestyle curves + bpy.context.scene.view_layers["View Layer"].use_pass_z = False + bpy.context.scene.view_layers["View Layer"].use_pass_combined = False + bpy.context.scene.view_layers["View Layer"].use_sky = False + bpy.context.scene.view_layers["View Layer"].use_solid = False + bpy.context.scene.view_layers["View Layer"].use_volumes = False + bpy.context.scene.view_layers["View Layer"].use_strand = True # freestyle curves + + recipe_file_path = dataset_dir.joinpath('recipe.yml') + recipe_yml_obj = get_recipe_yml_obj(recipe_file_path) + camera_angles = recipe_yml_obj['camera_angles_train'] + recipe_yml_obj['camera_angles_test'] + # euler setting + radius = 2 + eulers = [mathutils.Euler((math.radians(camera_angle[0]), 0.0, math.radians(camera_angle[1])), 'XYZ') for camera_angle in camera_angles] + + obj_gt_dir = dataset_dir.joinpath(phase, 'obj_gt') + path_to_sketches = dataset_dir.joinpath(phase, 'sketches') # output folder + if (parallel == 1 or mod == 0) and not path_to_sketches.is_dir(): + path_to_sketches.mkdir() + + if parallel == 1 and mod != 0: + while not path_to_sketches.is_dir(): + time.sleep(2) + + obj_files = sorted(obj_gt_dir.glob('*.obj')) + # filter out files that were already processed + obj_files = [file for file in obj_files if + not all( + list(path_to_sketches.glob(f'{file.stem}_{camera_angle[0]}_{camera_angle[1]}.png')) + for camera_angle in camera_angles)] + # remove any file that is not handled in this job + if parallel > 1: + obj_files = [file for file in obj_files if hash_file_name(file.name) % parallel == mod] + + for obj_file in tqdm(obj_files): + file_name = obj_file.name + + filepath = obj_gt_dir.joinpath(file_name) + bpy.ops.import_scene.obj(filepath=str(filepath), axis_forward='-Z', axis_up='Y', filter_glob="*.obj;*.mtl") + obj = bpy.context.selected_objects[0] + + # normalize the object + normalize_scale(obj) + + for i, eul in enumerate(eulers): + filename_no_ext = obj_file.stem + target_file_name = f"{filename_no_ext}_{camera_angles[i][0]:.1f}_{camera_angles[i][1]:.1f}.png" + target_file = path_to_sketches.joinpath(target_file_name) + if target_file.is_file(): + continue + + # camera setting + cam_pos = mathutils.Vector((0.0, -radius, 0.0)) + cam_pos.rotate(eul) + if i < 4: + # camera position perturbation + rand_x = random.uniform(-2.0, 2.0) + rand_z = random.uniform(-3.0, 3.0) + eul_perturb = mathutils.Euler((math.radians(rand_x), 0.0, math.radians(rand_z)), 'XYZ') + cam_pos.rotate(eul_perturb) + + scene = bpy.context.scene + bpy.ops.object.camera_add(enter_editmode=False, location=cam_pos) + new_camera = bpy.context.active_object + new_camera.name = "camera_tmp" + new_camera.data.name = "camera_tmp" + new_camera.data.lens_unit = 'FOV' + new_camera.data.angle = math.radians(60) + look_at(new_camera, Vector((0.0, 0.0, 0.0))) + + # render + scene.camera = new_camera + scene.render.filepath = str(target_file) + scene.render.resolution_x = 224 + scene.render.resolution_y = 224 + bpy.context.scene.cycles.samples = 10 + bpy.ops.render.render(write_still=True) + + # prepare for the next camera + del_obj(new_camera) + + # delete the obj to prepare for the next one + del_obj(obj) + + # clean the scene + clean_scene() + except Exception as e: + print(repr(e)) + print(traceback.format_exc()) + + +if __name__ == "__main__": + argv = sys.argv + if '--' in sys.argv: + # refer to https://b3d.interplanety.org/en/how-to-pass-command-line-arguments-to-a-blender-python-script-or-add-on/ + argv = sys.argv[sys.argv.index('--') + 1:] + else: + raise Exception("Expected \'--\' followed by arguments to the script") + + parser = argparse.ArgumentParser() + parser.add_argument('--dataset-dir', type=str, required=True, help='Path to dataset directory') + parser.add_argument('--parallel', type=int, default=1, help='Number of processes that will run the script') + parser.add_argument('--mod', type=int, default=0, help='The modulo for this process to match files\' hash') + parser.add_argument('--phases', type=str, required=True, nargs='+', help='List of phases to generate the sketches for') + + args = parser.parse_args(argv) + + # hide the main collections (if it is already hidden, there is no effect) + bpy.context.scene.view_layers['View Layer'].layer_collection.children['Main'].hide_viewport = True + bpy.context.scene.view_layers['View Layer'].layer_collection.children['Main'].exclude = True + + dataset_dir = Path(args.dataset_dir).expanduser() + phases = args.phases + for phase in phases: + main(dataset_dir, phase, args.parallel, args.mod) diff --git a/dataset_processing/__init__.py b/dataset_processing/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/dataset_processing/prepare_coseg.py b/dataset_processing/prepare_coseg.py new file mode 100644 index 0000000..fd9426c --- /dev/null +++ b/dataset_processing/prepare_coseg.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 + +import bpy +import sys +from tqdm import tqdm +from subprocess import Popen, PIPE +import argparse +from pathlib import Path +import importlib +import io +from contextlib import redirect_stdout + +def import_parents(level=1): + global __package__ + file = Path(__file__).resolve() + parent, top = file.parent, file.parents[level] + + sys.path.append(str(top)) + try: + sys.path.remove(str(parent)) + except ValueError: + # already removed + pass + + __package__ = '.'.join(parent.parts[len(top.parts):]) + importlib.import_module(__package__) + +if __name__ == '__main__' and __package__ is None: + import_parents(level=1) + +from common.bpy_util import normalize_scale + + +def main(): + if '--' in sys.argv: + # refer to https://b3d.interplanety.org/en/how-to-pass-command-line-arguments-to-a-blender-python-script-or-add-on/ + argv = sys.argv[sys.argv.index('--') + 1:] + else: + raise Exception("Expected \'--\' followed by arguments to the script") + + parser = argparse.ArgumentParser("prepare_coseg") + parser.add_argument('--shapes-dir', type=str, required=True, help='Path to COSEG raw shapes directory which contains the .off files') + parser.add_argument('--target-dataset-dir', type=str, required=True, help='Path to dataset directory where the normalized COSEG. obj files will be stored') + parser.add_argument('--target-phase', type=str, required=True, help='The name of the phase will hold the .obj files, e.g. \"coseg\"') + args = parser.parse_args(argv) + + shapes_dir = Path(args.shapes_dir) + target_dataset_dir = Path(args.target_dataset_dir).expanduser() + target_phase_dir = target_dataset_dir.joinpath(args.target_phase) + target_phase_dir.mkdir(exist_ok=True) + target_obj_gt_dir = target_phase_dir.joinpath('obj_gt') + target_obj_gt_dir.mkdir(exist_ok=True) + + print("Converting .off files to obj files...") + for off_file in tqdm(list(shapes_dir.glob('*.off'))): + obj_file = target_obj_gt_dir.joinpath(f'{off_file.stem}.obj') + if obj_file.is_file(): + continue + path_to_converter = Path(__file__).parent.joinpath('model-converter-python', 'convert.py').resolve() + cmd = [str(path_to_converter), '-i', str(off_file), '-o', str(obj_file)] + print(" ".join(cmd)) + process = Popen(cmd, stdout=PIPE) + process.wait() + + print("Normalizing obj files...") + for obj_file in tqdm(list(target_obj_gt_dir.glob("*.obj"))): + with redirect_stdout(io.StringIO()): + bpy.ops.import_scene.obj(filepath=str(obj_file)) + imported_object = bpy.context.selected_objects[0] + imported_object.data.materials.clear() + normalize_scale(imported_object) + bpy.ops.object.select_all(action='DESELECT') + bpy.context.view_layer.objects.active = imported_object + imported_object.select_set(True) + bpy.ops.export_scene.obj(filepath=str(obj_file), use_selection=True, use_materials=False, use_triangles=True) + bpy.ops.object.delete() + + +if __name__ == "__main__": + main() diff --git a/dataset_processing/save_obj.py b/dataset_processing/save_obj.py new file mode 100644 index 0000000..c55bfa1 --- /dev/null +++ b/dataset_processing/save_obj.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 + +import sys +import bpy +import argparse +import traceback +from pathlib import Path +import importlib + +def import_parents(level=1): + global __package__ + file = Path(__file__).resolve() + parent, top = file.parent, file.parents[level] + sys.path.append(str(top)) + try: + sys.path.remove(str(parent)) + except ValueError: + # already removed + pass + + __package__ = '.'.join(parent.parts[len(top.parts):]) + importlib.import_module(__package__) + +if __name__ == '__main__' and __package__ is None: + import_parents(level=1) + +from common.file_util import get_recipe_yml_obj +from common.input_param_map import get_input_param_map, load_shape_from_yml +from common.bpy_util import clean_scene, select_shape, select_objs, get_geometric_nodes_modifier, save_obj + + +def save_obj_from_yml(args): + if args.simplification_ratio: + assert 0.0 <= args.simplification_ratio <= 1.0 + target_obj_file_path = Path(args.target_obj_file_path) + assert target_obj_file_path.suffix == ".obj" + clean_scene(start_with_strings=["Camera", "Light"]) + bpy.context.scene.view_layers['View Layer'].layer_collection.children['Main'].hide_viewport = False + bpy.context.scene.view_layers['View Layer'].layer_collection.children['Main'].exclude = False + obj = select_shape() + gnodes_mod = get_geometric_nodes_modifier(obj) + recipe_yml = get_recipe_yml_obj(args.recipe_file_path) + input_params_map = get_input_param_map(gnodes_mod, recipe_yml) + load_shape_from_yml(args.yml_file_path, input_params_map, ignore_sanity_check=args.ignore_sanity_check) + dup_obj = save_obj(target_obj_file_path, simplification_ratio=args.simplification_ratio) + bpy.data.collections["Main"].hide_render = False + chair_obj = select_shape() + dup_obj.hide_render = False + chair_obj.hide_render = True + dup_obj.data.materials.clear() + select_objs(dup_obj) + bpy.ops.object.delete() + + +def main(): + if '--' in sys.argv: + # refer to https://b3d.interplanety.org/en/how-to-pass-command-line-arguments-to-a-blender-python-script-or-add-on/ + argv = sys.argv[sys.argv.index('--') + 1:] + else: + raise Exception("Expected \'--\' followed by arguments to the script") + + parser = argparse.ArgumentParser(prog='save_obj') + parser.add_argument('--recipe-file-path', type=str, required=True, help='Path to recipe.yml file') + parser.add_argument('--yml-file-path', type=str, required=True, help='Path to yaml file to convert to object') + parser.add_argument('--target-obj-file-path', type=str, required=True, help='Path the obj file that will be created') + parser.add_argument('--simplification-ratio', type=float, default=None, help='Simplification ratio to decimate the mesh') + parser.add_argument('--ignore-sanity-check', action='store_true', default=False, help='Do not check the shape\'s parameters') + + try: + args = parser.parse_known_args(argv)[0] + save_obj_from_yml(args) + except Exception as e: + print(repr(e)) + print(traceback.format_exc()) + +if __name__ == '__main__': + main() diff --git a/dataset_processing/simplified_mesh_dataset.py b/dataset_processing/simplified_mesh_dataset.py new file mode 100644 index 0000000..ac7f510 --- /dev/null +++ b/dataset_processing/simplified_mesh_dataset.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 + +import shutil +import argparse +import multiprocessing +from pathlib import Path +from functools import partial +from subprocess import Popen, PIPE + + +def simplify_and_save_obj_process(yml_file_path: Path, dst_phase_dir: Path, simplification_ratio: float, recipe_file_path, + blender_exe: Path, blend_file: Path): + assert 0.0 <= simplification_ratio <= 1.0 + simplification_ratio_str = f'{simplification_ratio:.3f}'.replace('.', '_') + out_file_name_no_ext = f'{yml_file_path.stem}_simplification_ratio_{simplification_ratio_str}' + print(f"Converting [{yml_file_path}] to obj file [{out_file_name_no_ext}]") + new_yml_path = dst_phase_dir.joinpath('yml_gt', f'{out_file_name_no_ext}.yml') + new_obj_path = dst_phase_dir.joinpath('obj_gt', f'{out_file_name_no_ext}.obj') + shutil.copy(yml_file_path, new_yml_path) + save_obj_script_path = Path(__file__).parent.joinpath('save_obj.py').resolve() + cmd = [str(blender_exe.expanduser()), str(blend_file.expanduser()), '-b', '--python', + str(save_obj_script_path), '--', + '--recipe-file-path', str(recipe_file_path), + '--yml-file-path', str(yml_file_path), + '--target-obj-file-path', str(new_obj_path), + '--simplification-ratio', str(simplification_ratio)] + print(" ".join(cmd)) + process = Popen(cmd, stdout=PIPE) + process.wait() + + +def main(): + parser = argparse.ArgumentParser("simplified_mesh_dataset") + parser.add_argument('--dataset-dir', type=str, required=True, help='Path to dataset directory') + parser.add_argument('--src-phase', type=str, required=True, help='Directory name of the source dataset phase') + parser.add_argument('--dst-phase', type=str, default='simplified', help='Directory name of the destination dataset phase') + parser.add_argument('--blender-exe', type=str, required=True, help='Path to blender executable') + parser.add_argument('--blend-file', type=str, required=True, help='Path to blend file') + parser.add_argument('--simplification-ratios', type=float, required=True, nargs='+', help='List of simplification ratios to generate datasets for,' + 'note that 1.0 means the original shape is used with no simplification') + args = parser.parse_args() + + assert all([0.0 <= simplification_ratio <= 1.0 for simplification_ratio in args.simplification_ratios]) + + ds_dir = Path(args.dataset_dir).expanduser() + src_phase_dir = ds_dir.joinpath(args.src_phase) + dst_phase_dir = ds_dir.joinpath(args.dst_phase) + if dst_phase_dir.is_dir(): + raise Exception(f'Simplified mesh dataset target directory [{dst_phase_dir}] already exists') + dst_phase_dir.mkdir() + + dest_yml_gt_path = dst_phase_dir.joinpath('yml_gt') + dest_yml_gt_path.mkdir(exist_ok=True) + dest_obj_gt_path = dst_phase_dir.joinpath('obj_gt') + dest_obj_gt_path.mkdir(exist_ok=True) + + blender_exe = Path(args.blender_exe) + blend_file = Path(args.blend_file) + + # create all the obj from the prediction yaml files + # note that for point cloud we have one yml and for sketch we have multiple yml files (one for each camera angle) + cpu_count = multiprocessing.cpu_count() + print(f"Generating simplified-mesh dataset in [{src_phase_dir}] with [{cpu_count}] processes") + recipe_file_path = Path(args.dataset_dir, 'recipe.yml') + src_yml_gt_path = src_phase_dir.joinpath('yml_gt') + yml_file_paths = [yml_file_path for yml_file_path in src_yml_gt_path.glob("*.yml")] + # note that 1.0 means no simplification takes place + for simplification_ratio in args.simplification_ratios: + print(f"Started generating simplified-mesh dataset with ratio [{simplification_ratio}]") + simplify_and_save_obj_process_partial = partial(simplify_and_save_obj_process, + dst_phase_dir=dst_phase_dir, + simplification_ratio=simplification_ratio, + recipe_file_path=recipe_file_path, + blender_exe=blender_exe, + blend_file=blend_file) + p = multiprocessing.Pool(cpu_count) + p.map(simplify_and_save_obj_process_partial, yml_file_paths) + p.close() + p.join() + + +if __name__ == "__main__": + main() diff --git a/environment.yml b/environment.yml new file mode 100644 index 0000000..293dced --- /dev/null +++ b/environment.yml @@ -0,0 +1,29 @@ +name: geocode +channels: + - pytorch + - defaults + - dglteam + - bottler + - conda-forge + - fvcore + - iopath + - pytorch3d +dependencies: + - python=3.8 + - pytorch=1.10.2 + - torchvision + - numpy=1.21.2 + - matplotlib=3.5.1 + - pytorch-lightning=1.5.10 + - neptune-client=0.16.4 + - nvidiacub=1.10.0 + - dgl=0.9.1 + - scikit-image=0.19.2 + - iopath=0.1.9 + - fvcore=0.1.5.post20210915 + - tqdm>=4.64.0 + - gdown>=4.5.4 + - pytorch3d + - pip + - pip: + - git+https://github.com/otaheri/chamfer_distance diff --git a/geocode/__init__.py b/geocode/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/geocode/barplot_util.py b/geocode/barplot_util.py new file mode 100644 index 0000000..c455d63 --- /dev/null +++ b/geocode/barplot_util.py @@ -0,0 +1,68 @@ +import json +import numpy as np + + +def gen_and_save_barplot(barplot_json_path, title, barplot_target_image_path=None): + from matplotlib import pyplot as plt + with open(barplot_json_path, 'r') as barplot_json_file: + data = json.load(barplot_json_file) + + inputs_to_eval = data['inputs_to_eval'] + correct_arr_pc = data['correct_arr_pc'] + correct_arr_sketch = data['correct_arr_sketch'] + total_pc = data['total_pc'] + total_sketch = data['total_sketch'] + + correct = [a + b for a, b in zip(correct_arr_pc, correct_arr_sketch)] + accuracy_avg = [a / (total_pc + total_sketch) for a in correct] + accuracy_pc = [a / total_pc for a in correct_arr_pc] + accuracy_sketch = [a / total_sketch for a in correct_arr_sketch] + + overall_acc_avg = (sum(correct_arr_pc) + sum(correct_arr_sketch)) / ( + len(inputs_to_eval) * (total_pc + total_sketch)) + overall_acc_pc = sum(correct_arr_pc) / (len(inputs_to_eval) * total_pc) + overall_acc_sketch = sum(correct_arr_sketch) / (len(inputs_to_eval) * total_sketch) + + is_only_sketches = False + is_only_pcs = False + if all([param_acc == 0 for param_acc in accuracy_pc]): + # only sketches + overall_acc_avg = overall_acc_sketch + is_only_sketches = True + if all([param_acc == 0 for param_acc in accuracy_sketch]): + # only pcs + overall_acc_avg = overall_acc_pc + is_only_pcs = True + + # sort by average accuracy + inputs_to_eval, accuracy_avg, accuracy_pc, accuracy_sketch = zip( + *sorted(zip(inputs_to_eval, accuracy_avg, accuracy_pc, accuracy_sketch), key=lambda x: x[1])) + + inputs_to_eval += ("Overall",) + accuracy_avg += (overall_acc_avg,) + accuracy_pc += (overall_acc_pc,) + accuracy_sketch += (overall_acc_sketch,) + + fig, ax = plt.subplots(figsize=(16, 14)) + X_axis = np.arange(len(inputs_to_eval)) * 2.6 + if not is_only_pcs and not is_only_sketches: + pps = ax.barh(X_axis + 0.7, accuracy_avg, 0.7, color='steelblue') + ax.barh(X_axis - 0.0, accuracy_pc, 0.7, color='lightsteelblue') + ax.barh(X_axis - 0.7, accuracy_sketch, 0.7, color='wheat') + ax.legend(labels=['Average', 'Point Clouds', 'Sketches']) + elif is_only_pcs: + pps = ax.barh(X_axis - 0.0, accuracy_pc, 0.7, color='lightsteelblue') + ax.legend(labels=['Point Clouds']) + elif is_only_sketches: + pps = ax.barh(X_axis - 0.7, accuracy_sketch, 0.7, color='wheat') + ax.legend(labels=['Sketches']) + else: + raise Exception("Either point cloud or sketch input should be processed") + + ax.bar_label(pps, fmt='%.2f', label_type='center', fontsize=8) + ax.set_yticks(X_axis, inputs_to_eval) + ax.set_title(title) + + if barplot_target_image_path: + plt.savefig(barplot_target_image_path) + return fig diff --git a/geocode/calculator_accuracy.py b/geocode/calculator_accuracy.py new file mode 100644 index 0000000..a04e28a --- /dev/null +++ b/geocode/calculator_accuracy.py @@ -0,0 +1,52 @@ +import torch +from calculator_util import eval_metadata + + +class AccuracyCalculator(): + def __init__(self, inputs_to_eval, param_descriptors): + self.inputs_to_eval = inputs_to_eval + self.normalized_classes_all, self.num_classes_all_shifted_cumulated, self.num_classes_all, self.regression_params_indices \ + = eval_metadata(inputs_to_eval, param_descriptors) + self.param_descriptors = param_descriptors + + def eval(self, pred, targets, top_k_acc): + batch_size = pred.shape[0] + device = targets.device + normalized_classes_all = self.normalized_classes_all.to(device) + num_classes_all_shifted_cumulated = self.num_classes_all_shifted_cumulated.to(device) + num_classes_all = self.num_classes_all.to(device) + correct = [[0] * len(self.inputs_to_eval) for _ in range(top_k_acc)] + targets_interleaved = torch.repeat_interleave(targets, num_classes_all.view(-1), dim=1) + normalized_classes_all_repeated = normalized_classes_all.repeat(batch_size, 1).to(device) + target_class = torch.abs(normalized_classes_all_repeated - targets_interleaved) + target_class = torch.where(target_class < 1e-3)[1].view(batch_size, -1) # take the indices along dim=1 since target is of size [1, param_count] + if len(self.regression_params_indices) > 0: + regression_params_indices_repeated = self.regression_params_indices.repeat(batch_size, 1).to(device) + target_class = torch.cat((target_class, regression_params_indices_repeated), dim=1) + target_class, _ = torch.sort(target_class, dim=1) + assert target_class.shape[1] == len(self.inputs_to_eval) + target_class = target_class - num_classes_all_shifted_cumulated + pred_split = torch.split(pred, list(num_classes_all), dim=1) + class_indices_diff = [(torch.argmax(p, dim=1) - t if p.shape[1] > 1 else None) for p, t in zip( pred_split, target_class.T )] + + l1_distance = [None] * targets.shape[1] + if len(self.regression_params_indices) > 0: + for param_idx, (p, t) in enumerate(zip(pred_split, targets.T)): + if self.param_descriptors[self.inputs_to_eval[param_idx]].is_regression: + adjusted_pred = p[:, 0].clone() + adjusted_pred[p[:, 1] >= 0.5] = -1.0 + l1_distance[param_idx] = torch.abs(adjusted_pred.squeeze() - t) + + for i, param_name in enumerate(self.inputs_to_eval): + if self.param_descriptors[param_name].is_regression: + # regression parameter + normalized_acc_threshold = self.param_descriptors[param_name].normalized_acc_threshold + for j in range(top_k_acc): + assert len(l1_distance[i]) == batch_size + correct[j][i] += torch.sum((l1_distance[i] < normalized_acc_threshold * (j + 1)).int()).item() + else: + cid = class_indices_diff[i] + assert len(cid) == batch_size + for j in range(top_k_acc): + correct[j][i] += len(cid[(cid <= j) & (cid >= -j)]) + return correct diff --git a/geocode/calculator_loss.py b/geocode/calculator_loss.py new file mode 100644 index 0000000..2800e12 --- /dev/null +++ b/geocode/calculator_loss.py @@ -0,0 +1,58 @@ +import torch +from calculator_util import eval_metadata + + +MSE = torch.nn.MSELoss() +CElossSum = torch.nn.CrossEntropyLoss(reduction='sum') + + +class LossCalculator(): + def __init__(self, inputs_to_eval, param_descriptors): + self.inputs_to_eval = inputs_to_eval + self.normalized_classes_all, self.num_classes_all_shifted_cumulated, self.num_classes_all, self.regression_params_indices \ + = eval_metadata(inputs_to_eval, param_descriptors) + self.param_descriptors = param_descriptors + + def loss(self, pred, targets): + """ + _pred: (B, TARGET_VEC_LEN) + """ + batch_size = pred.shape[0] + device = targets.device + normalized_classes_all = self.normalized_classes_all.to(device) + num_classes_all_shifted_cumulated = self.num_classes_all_shifted_cumulated.to(device) + num_classes_all = self.num_classes_all.to(device) + targets_interleaved = torch.repeat_interleave(targets, num_classes_all.view(-1), dim=1) + normalized_classes_all_repeated = normalized_classes_all.repeat(batch_size, 1).to(device) + target_class = torch.abs(normalized_classes_all_repeated - targets_interleaved) + target_class = torch.where(target_class < 1e-3)[1].view(batch_size, -1) # take the indices along dim=1 + if len(self.regression_params_indices) > 0: + regression_params_indices_repeated = self.regression_params_indices.repeat(batch_size, 1).to(device) + target_class = torch.cat((target_class, regression_params_indices_repeated), dim=1) + target_class, _ = torch.sort(target_class, dim=1) + assert target_class.shape[1] == len(self.inputs_to_eval) + target_class = target_class - num_classes_all_shifted_cumulated + # target_class = target_class.to(_pred.get_device()) + pred_split = torch.split(pred, list(num_classes_all), dim=1) + detailed_ce_loss = [(CElossSum(p, t) if p.shape[1] > 1 else None) for p, t in zip( pred_split, target_class.T )] + + detailed_mse_loss = [None] * targets.shape[1] + if len(self.regression_params_indices) > 0: + for param_idx, (p, t) in enumerate(zip(pred_split, targets.T)): + if self.param_descriptors[self.inputs_to_eval[param_idx]].is_regression: + t_visibility = torch.zeros(t.shape[0]) + t_visibility[t >= 0.0] = 0.0 + t_visibility[t == -1.0] = 1.0 + t_visibility = t_visibility.to(device) + t_clone = t.clone() + t_clone = t_clone.float() + t_clone[t_clone == -1] = p[t_clone == -1,0] + t_adjusted = torch.concat((t_clone.unsqueeze(1), t_visibility.unsqueeze(1)), dim=1) + detailed_mse_loss[param_idx] = MSE(p, t_adjusted) + detailed_mse_loss_no_none = [e for e in detailed_mse_loss if e] + detailed_ce_loss_no_none = [e for e in detailed_ce_loss if e] + mse_loss_range = 1.0 if not detailed_mse_loss_no_none else (max(detailed_mse_loss_no_none).item() - min(detailed_mse_loss_no_none).item()) + ce_loss_range = max(detailed_ce_loss_no_none).item() - min(detailed_ce_loss_no_none).item() + detailed_loss = [(ce_loss / ce_loss_range) if not mse_loss else (mse_loss / mse_loss_range) for ce_loss, mse_loss in zip(detailed_ce_loss, detailed_mse_loss)] + + return sum(detailed_loss), detailed_loss diff --git a/geocode/calculator_util.py b/geocode/calculator_util.py new file mode 100644 index 0000000..3f50427 --- /dev/null +++ b/geocode/calculator_util.py @@ -0,0 +1,27 @@ +import torch +from typing import Dict, List +from common.param_descriptors import ParamDescriptor + + +def eval_metadata(inputs_to_eval: List[str], param_descriptors_map: Dict[str, ParamDescriptor]): + num_classes_all = torch.empty(0) + normalized_classes_all = torch.empty(0) + for i, param_name in enumerate(inputs_to_eval): + param_descriptor = param_descriptors_map[param_name] + num_classes = param_descriptor.num_classes # Including the visibility label. If using regression then num_classes=2. + num_classes_all = torch.cat((num_classes_all, torch. tensor([num_classes]))).long() + if param_descriptor.normalized_classes is not None: + normalized_classes = torch.from_numpy(param_descriptor.normalized_classes) + else: + # high values so that eval and loss methods will work when using regression + normalized_classes = torch.tensor([100000.0, 100000.0]) + normalized_classes_all = torch.cat((normalized_classes_all, normalized_classes.view(-1))) + num_classes_all_shifted = torch.cat((torch.tensor([0]), num_classes_all))[0:-1] # shift right + drop right-most element + num_classes_all_shifted_cumulated = torch.cumsum(num_classes_all_shifted, dim=0).view(1, -1) + + # get the indices of all the regression params, then shift them to match the expanded vector + regression_params = torch.tensor([param_descriptors_map[param_name].is_regression for param_name in inputs_to_eval], dtype=torch.int) + regression_params_indices = torch.where(regression_params)[0] + regression_params_indices = torch.tensor([num_classes_all_shifted_cumulated[0, idx] for idx in regression_params_indices]) + + return normalized_classes_all, num_classes_all_shifted_cumulated, num_classes_all, regression_params_indices diff --git a/geocode/geocode.py b/geocode/geocode.py new file mode 100644 index 0000000..a58f827 --- /dev/null +++ b/geocode/geocode.py @@ -0,0 +1,62 @@ +#!/usr/bin/env python3 + +import argparse +from geocode_util import InputType +from geocode_train import train +from geocode_test import test + + +def str2bool(v): + if isinstance(v, bool): + return v + if v.lower() in ('yes', 'true', 't', 'y', '1'): + return True + elif v.lower() in ('no', 'false', 'f', 'n', '0'): + return False + else: + raise argparse.ArgumentTypeError('Boolean value expected but got [{}].'.format(v)) + + +def main(): + parser = argparse.ArgumentParser(prog='ShapeEditing') + + common_parser = argparse.ArgumentParser(add_help=False) + common_parser.add_argument('--dataset-dir', type=str, required=True, help='Path to dataset directory') + common_parser.add_argument('--models-dir', type=str, required=True, help='Directory where experiments will be saved') + common_parser.add_argument('--exp-name', type=str, required=True, help='Experiment directory within the models directory, where checkpoints will be saved') + common_parser.add_argument('--input-type', type=InputType, nargs='+', default='pc sketch', help='Either \"pc\", \"sketch\" or \"pc sketch\"') + common_parser.add_argument('--increase-network-size', action='store_true', default=False, help='Use larger encoders networks sizes') + common_parser.add_argument('--normalize-embeddings', action='store_true', default=False, help='Normalize embeddings before using the decoders') + common_parser.add_argument('--pretrained-vgg', action='store_true', default=False, help='Use a pretrained VGG network') + common_parser.add_argument('--use-regression', action='store_true', default=False, help='Use regression instead of classification for continuous parameters') + + sp = parser.add_subparsers() + sp_train = sp.add_parser('train', parents=[common_parser]) + sp_test = sp.add_parser('test', parents=[common_parser]) + + sp_train.set_defaults(func=train) + sp_test.set_defaults(func=test) + + sp_train.add_argument('--batch_size', type=int, required=True, help='Batch size') + sp_train.add_argument('--nepoch', type=int, required=True, help='Number of epochs to train') + + sp_test.add_argument('--phase', type=str, default='test') + sp_test.add_argument('--blender-exe', type=str, required=True, help='Path to blender executable') + sp_test.add_argument('--blend-file', type=str, required=True, help='Path to blend file') + sp_test.add_argument('--random-pc', type=int, default=None, help='Use only random point cloud sampling with specified number of points') + sp_test.add_argument('--gaussian', type=float, default=0.0, help='Add Gaussian noise to the point cloud with the specified STD') + sp_test.add_argument('--normalize-pc', action='store_true', default='False', help='Automatically normalize the input point clouds') + sp_test.add_argument('--scanobjectnn', action='store_true', default='False', help='ScanObjectNN dataset which has only point clouds input') + # we augment in phases "train", "val", "test" and experiments "coseg", "simplify_mesh", and "gaussian" + # use `--augment-with-random-points false` to disable + sp_test.add_argument('--augment-with-random-points', type=str2bool, default='True', help='Augment FPS point cloud with randomly sampled points') + + args = parser.parse_args() + args.func(args) + + # either pc or sketch, or both must be trained + assert InputType.pc in args.input_type or InputType.sketch in args.input_type + + +if __name__ == "__main__": + main() diff --git a/geocode/geocode_model.py b/geocode/geocode_model.py new file mode 100644 index 0000000..9e339bd --- /dev/null +++ b/geocode/geocode_model.py @@ -0,0 +1,339 @@ +import yaml +import json +import shutil +import torch +import torch.optim as optim +from torch.optim.lr_scheduler import StepLR +import pytorch_lightning as pl +from barplot_util import gen_and_save_barplot +from neptune.new.types import File +from models.dgcnn import DGCNN +from models.vgg import vgg11_bn +from models.decoder import DecodersNet +from calculator_accuracy import AccuracyCalculator +from calculator_loss import LossCalculator +from common.param_descriptors import ParamDescriptors +from pathlib import Path +from geocode_util import InputType + + +class Model(pl.LightningModule): + def __init__(self, top_k_acc, batch_size, detailed_vec_size, increase_network_size, normalize_embeddings, + pretrained_vgg, input_type, inputs_to_eval, lr, sched_step_size, sched_gamma, + exp_name=None, trainer=None, param_descriptors: ParamDescriptors = None, models_dir: Path = None, + results_dir: Path = None, test_dir: Path = None, test_dataloaders_types=None, test_input_type=None, + use_regression=False): + super().__init__() + # saved hyper parameters + self.input_type = input_type + self.inputs_to_eval = inputs_to_eval + self.batch_size = batch_size + self.lr = lr + self.sched_step_size = sched_step_size + self.sched_gamma = sched_gamma + self.top_k_acc = top_k_acc + + # non-saved parameters + self.trainer = trainer + self.param_descriptors = param_descriptors + self.param_descriptors_map = self.param_descriptors.get_param_descriptors_map() + self.results_dir = results_dir + self.test_dir = test_dir + self.exp_name = exp_name + self.models_dir = models_dir + self.test_dataloaders_types = test_dataloaders_types + self.test_type = test_input_type + self.use_regression = use_regression + + regression_params = None + if self.use_regression: + regression_params = [param_descriptor.input_type == "Float" or param_descriptor.input_type == "Vector" for + param_name, param_descriptor in self.param_descriptors_map.items()] + self.acc_calc = AccuracyCalculator(self.inputs_to_eval, self.param_descriptors_map) + self.loss_calc = LossCalculator(self.inputs_to_eval, self.param_descriptors_map) + self.decoders_net = DecodersNet(output_channels=detailed_vec_size, increase_network_size=increase_network_size, + regression_params=regression_params) + self.dgcnn = None + self.vgg = None + if InputType.pc in self.input_type: + self.dgcnn = DGCNN(increase_network_size=increase_network_size, normalize_embeddings=normalize_embeddings) + if InputType.sketch in self.input_type: + self.vgg = vgg11_bn(pretrained=pretrained_vgg, progress=True, encoder_only=True, + increase_network_size=increase_network_size, normalize_embeddings=normalize_embeddings) + + self.save_hyperparameters( + ignore=["trainer", "param_descriptors", "models_dir", "results_dir", "test_dir", "test_input_type", + "use_regression", "exp_name", "test_dataloaders_types"]) + + def configure_optimizers(self): + params = list(self.decoders_net.parameters()) + if InputType.pc in self.input_type: + params += list(self.dgcnn.parameters()) + if InputType.sketch in self.input_type: + params += list(self.vgg.parameters()) + optimizer = optim.Adam(params, lr=self.lr) + lr_scheduler = StepLR(optimizer, step_size=self.sched_step_size, gamma=self.sched_gamma) + return [optimizer], [lr_scheduler] + + def log_accuracy(self, phase, correct_arr, metric_type): + """ + :param phase: e.g. "train" or "val" + :param correct_arr: two dim array, rows are k in top-k accuracy, cols are the parameters + the value is the number of correct predictions for a parameter when considering + top-k accuracy (top-k with k=1 is equivalent to argmax) + :param metric_type: either "pc" or "sketch" + """ + for i, param_name in enumerate(self.inputs_to_eval): + for j in range(self.top_k_acc): + acc_metric = correct_arr[j][i] / self.batch_size + self.log(f"{phase}/acc_top{j + 1}/{metric_type}/{param_name}", acc_metric, on_step=False, on_epoch=True, + logger=True, batch_size=self.batch_size) + for j in range(self.top_k_acc): + acc_avg_metric = sum(correct_arr[j]) / (self.batch_size * len(self.inputs_to_eval)) + self.log(f"{phase}/acc_top{j + 1}/{metric_type}/avg", acc_avg_metric, on_step=False, on_epoch=True, + logger=True, batch_size=self.batch_size) + + def get_decoder_loss(self, pc_emb, sketch_emb, targets_pcs, targets_sketches): + if InputType.pc in self.input_type: + batch_size = pc_emb.shape[0] + else: + batch_size = sketch_emb.shape[0] + pred_pc = None + pred_sketch = None + if InputType.pc in self.input_type and InputType.sketch in self.input_type: + targets_both = torch.cat((targets_pcs, targets_sketches), dim=0) + pred_both = self.decoders_net.decode(torch.cat((pc_emb, sketch_emb), dim=0)) + pred_pc = pred_both[:batch_size, :] + pred_sketch = pred_both[batch_size:, :] + decoders_loss, detailed_decoder_loss = self.loss_calc.loss(pred_both, targets_both) + elif InputType.pc in self.input_type: + pred_pc = self.decoders_net.decode(pc_emb) + decoders_loss, detailed_decoder_loss = self.loss_calc.loss(pred_pc, targets_pcs) + elif InputType.sketch in self.input_type: + pred_sketch = self.decoders_net.decode(sketch_emb) + decoders_loss, detailed_decoder_loss = self.loss_calc.loss(pred_sketch, targets_sketches) + else: + raise Exception("Illegal input type") + return decoders_loss, detailed_decoder_loss, pred_pc, pred_sketch + + def training_step(self, train_batch, batch_idx): + targets_pcs = None + targets_sketches = None + pc_emb = None + sketch_emb = None + if InputType.pc in self.input_type: + _, pcs, targets_pcs, _ = train_batch["pc"] + pcs = pcs.transpose(2, 1) + pc_emb = self.dgcnn(pcs) + if InputType.sketch in self.input_type: + _, _, sketches, targets_sketches, _ = train_batch["sketch"] + sketch_emb = self.vgg(sketches) + + decoders_loss, detailed_decoder_loss, pred_pc, pred_sketch = self.get_decoder_loss(pc_emb, sketch_emb, + targets_pcs, + targets_sketches) + # log detailed decoding loss + for i, param_name in enumerate(self.inputs_to_eval): + self.log(f"train/loss/{param_name}", detailed_decoder_loss[i], on_step=False, on_epoch=True, logger=True, + batch_size=self.batch_size) + + train_loss = decoders_loss + self.log("train/loss/total", train_loss, on_step=False, on_epoch=True, logger=True, batch_size=self.batch_size) + + # compute and log accuracy for point cloud and sketch + if InputType.pc in self.input_type: + correct_arr_pc = self.acc_calc.eval(pred_pc, targets_pcs, self.top_k_acc) + self.log_accuracy("train", correct_arr_pc, "pc") + if InputType.sketch in self.input_type: + correct_arr_sketch = self.acc_calc.eval(pred_sketch, targets_sketches, self.top_k_acc) + self.log_accuracy("train", correct_arr_sketch, "sketch") + + return train_loss + + def validation_step(self, val_batch, batch_idx): + targets_pcs = None + targets_sketches = None + pc_emb = None + sketch_emb = None + if InputType.pc in self.input_type: + _, pcs, targets_pcs, _ = val_batch["pc"] + pcs = pcs.transpose(2, 1) + pc_emb = self.dgcnn(pcs) + if InputType.sketch in self.input_type: + _, _, sketches, targets_sketches, _ = val_batch["sketch"] + sketch_emb = self.vgg(sketches) + + decoders_loss, detailed_decoder_loss, pred_pc, pred_sketch = self.get_decoder_loss(pc_emb, sketch_emb, + targets_pcs, + targets_sketches) + # log detailed decoding loss + for i, param_name in enumerate(self.inputs_to_eval): + self.log(f"val/loss/{param_name}", detailed_decoder_loss[i], on_step=False, on_epoch=True, logger=True, + batch_size=self.batch_size) + val_loss = decoders_loss + self.log("val/loss/total", val_loss, on_step=False, on_epoch=True, logger=True, batch_size=self.batch_size) + + # compute and log accuracy for point cloud and sketch + correct_arr_pc = None + if InputType.pc in self.input_type: + correct_arr_pc = self.acc_calc.eval(pred_pc, targets_pcs, self.top_k_acc) + self.log_accuracy("val", correct_arr_pc, "pc") + correct_arr_sketch = None + if InputType.sketch in self.input_type: + correct_arr_sketch = self.acc_calc.eval(pred_sketch, targets_sketches, self.top_k_acc) + self.log_accuracy("val", correct_arr_sketch, "sketch") + + # log average validation accuracy + pc_acc_avg = [] + if InputType.pc in self.input_type: + for j in range(self.top_k_acc): + curr_avg = sum(correct_arr_pc[j]) / (self.batch_size * len(self.inputs_to_eval)) + pc_acc_avg.append(curr_avg) + sketch_acc_avg = [] + if InputType.sketch in self.input_type: + for j in range(self.top_k_acc): + curr_avg = sum(correct_arr_sketch[j]) / (self.batch_size * len(self.inputs_to_eval)) + sketch_acc_avg.append(curr_avg) + avg_acc = [] + for j in range(self.top_k_acc): + if InputType.pc in self.input_type and InputType.sketch in self.input_type: + curr_avg = (pc_acc_avg[j] + sketch_acc_avg[j]) / 2 + elif InputType.pc in self.input_type: + curr_avg = pc_acc_avg[j] + elif InputType.sketch in self.input_type: + curr_avg = sketch_acc_avg[j] + else: + raise Exception("Illegal input type") + avg_acc.append(curr_avg) + self.log(f"val/acc_top{j + 1}/avg", curr_avg, on_step=False, on_epoch=True, logger=True, + batch_size=self.batch_size) + return avg_acc, correct_arr_pc, correct_arr_sketch + + def validation_epoch_end(self, validation_step_outputs): + num_batches = len(validation_step_outputs) + num_samples = num_batches * self.batch_size + for j in range(self.top_k_acc): + avg_acc = 0 + correct_arr_pc = [0] * len(self.inputs_to_eval) + correct_arr_sketch = [0] * len(self.inputs_to_eval) + for avg_acc_batch, correct_arr_pc_batch, correct_arr_sketch_batch in validation_step_outputs: + avg_acc += avg_acc_batch[j] + for i in range(len(self.inputs_to_eval)): + if correct_arr_pc_batch: + correct_arr_pc[i] += correct_arr_pc_batch[j][i] + if correct_arr_sketch_batch: + correct_arr_sketch[i] += correct_arr_sketch_batch[j][i] + avg_acc /= len(validation_step_outputs) + + if f'val/acc_top{j + 1}/best_avg' in self.trainer.callback_metrics: + best_avg_val_acc = max(self.trainer.callback_metrics[f'val/acc_top{j + 1}/best_avg'], avg_acc) + else: + best_avg_val_acc = avg_acc + self.log(f"val/acc_top{j + 1}/best_avg", best_avg_val_acc, on_step=False, on_epoch=True, logger=True, + batch_size=self.batch_size) + if avg_acc == best_avg_val_acc: + barplot_data = { + "inputs_to_eval": self.inputs_to_eval, + "correct_arr_pc": correct_arr_pc, + "total_pc": num_samples, + "correct_arr_sketch": correct_arr_sketch, + "total_sketch": num_samples, + "accuracy_top_k": j + 1, + } + barplot_data_file_path = f'{self.models_dir}/{self.exp_name}/val_barplot_top_{j + 1}.json' + with open(barplot_data_file_path, 'w') as barplot_data_file: + json.dump(barplot_data, barplot_data_file) + # log the bar plot as image + fig = gen_and_save_barplot(barplot_data_file_path, + title=f"Validation Accuracy Top {j + 1} @ Epoch {self.trainer.current_epoch}") + if self.logger: + self.logger.run["barplot"].log(File.as_image(fig)) + + def test_step(self, test_batch, batch_idx, dataloader_idx=0): + assert self.batch_size == 1 + + pc, sketch = None, None + if self.test_dataloaders_types[dataloader_idx] == 'pc': + file_name, pc, targets, shape = test_batch + elif self.test_dataloaders_types[dataloader_idx] == 'sketch': + file_name, sketch_camera_angle, sketch, targets, shape = test_batch + else: + raise Exception(f"Unrecognized dataloader type [{self.test_dataloaders_types[dataloader_idx]}]") + + if pc is not None: + pcs = pc.transpose(2, 1) + pred_pc = self.decoders_net.decode(self.dgcnn(pcs)) + pred_map_pc = self.param_descriptors.convert_prediction_vector_to_map(pred_pc.cpu(), use_regression=self.use_regression) + pc_pred_yaml_file_path = self.results_dir.joinpath('yml_predictions_pc', f'{file_name[0]}_pred_pc.yml') + with open(pc_pred_yaml_file_path, 'w') as yaml_file: + yaml.dump(pred_map_pc, yaml_file) + elif sketch is not None: + pred_sketch = self.decoders_net.decode(self.vgg(sketch)) + pred_map_sketch = self.param_descriptors.convert_prediction_vector_to_map(pred_sketch.cpu(), use_regression=self.use_regression) + sketch_pred_yaml_file_path = self.results_dir.joinpath('yml_predictions_sketch', + f'{file_name[0]}_{sketch_camera_angle[0]}_pred_sketch.yml') + with open(sketch_pred_yaml_file_path, 'w') as yaml_file: + yaml.dump(pred_map_sketch, yaml_file) + else: + raise Exception("No pc and no sketch input") + + # compute accuracy for point cloud and sketch + correct_arr_pc = None + correct_arr_sketch = None + if shape: + # this means we had a target vector to compare against + if pc is not None: + correct_arr_pc = self.acc_calc.eval(pred_pc, targets, self.top_k_acc) + if sketch is not None: + correct_arr_sketch = self.acc_calc.eval(pred_sketch, targets, self.top_k_acc) + # save ground truth yaml + expanded_targets_vector = self.param_descriptors.expand_target_vector(targets.cpu()) + gt_map = self.param_descriptors.convert_prediction_vector_to_map(expanded_targets_vector, use_regression=self.use_regression) + gt_yaml_file_path = self.results_dir.joinpath('yml_gt', f'{file_name[0]}_gt.yml') + with open(gt_yaml_file_path, 'w') as yaml_file: + yaml.dump(gt_map, yaml_file) + + # save the gt sketches, note that obj gt are saved during the visualization step + sketches_dir = self.test_dir.joinpath("sketches") + if sketches_dir.is_dir(): + sketch_files = sketches_dir.glob(f'{file_name[0]}*') + for sketch_file in sketch_files: + shutil.copy(sketch_file, self.results_dir.joinpath('sketch_gt', sketch_file.name)) + + return correct_arr_pc, correct_arr_sketch + + def test_epoch_end(self, test_step_outputs): + assert self.batch_size == 1 + assert self.test_type + if InputType.pc in self.test_type and InputType.sketch in self.test_type: + test_step_outputs_pc_and_sketch = test_step_outputs[0] + test_step_outputs[1] + else: + test_step_outputs_pc_and_sketch = test_step_outputs + for top_k in range(self.top_k_acc): + total_pcs = 0 + total_sketches = 0 + correct_arr_pc = [0] * len(self.inputs_to_eval) + correct_arr_sketch = [0] * len(self.inputs_to_eval) + for correct_arr_pc_batch, correct_arr_sketch_batch in test_step_outputs_pc_and_sketch: + if correct_arr_pc_batch is not None: + total_pcs += 1 + for i in range(len(self.inputs_to_eval)): + correct_arr_pc[i] += correct_arr_pc_batch[top_k][i] + if correct_arr_sketch_batch is not None: + total_sketches += 1 + for i in range(len(self.inputs_to_eval)): + correct_arr_sketch[i] += correct_arr_sketch_batch[top_k][i] + barplot_data = { + "inputs_to_eval": self.inputs_to_eval, + "correct_arr_pc": correct_arr_pc, + "total_pc": total_pcs, + "correct_arr_sketch": correct_arr_sketch, + "total_sketch": total_sketches, + "accuracy_top_k": top_k + 1, + } + if total_pcs > 0 and total_sketches > 0: + print( + f'Saving test barplot for [{total_pcs}] point cloud sampels and [{total_sketches}] sketch samples') + barplot_data_file_path = f'{self.models_dir}/{self.exp_name}/test_barplot_top_{top_k + 1}.json' + with open(barplot_data_file_path, 'w') as barplot_data_file: + json.dump(barplot_data, barplot_data_file) diff --git a/geocode/geocode_test.py b/geocode/geocode_test.py new file mode 100644 index 0000000..0977ec3 --- /dev/null +++ b/geocode/geocode_test.py @@ -0,0 +1,262 @@ +import json +import torch +import shutil +import traceback +import numpy as np +import multiprocessing +from pathlib import Path +from functools import partial +import pytorch_lightning as pl +from subprocess import Popen, PIPE +from data.dataset_pc import DatasetPC +from data.dataset_sketch import DatasetSketch +from barplot_util import gen_and_save_barplot +from common.param_descriptors import ParamDescriptors +from geocode_util import InputType, get_inputs_to_eval, calc_prediction_vector_size +from geocode_model import Model +from torch.utils.data import DataLoader +from chamfer_distance import ChamferDistance as chamfer_dist +from common.sampling_util import sample_surface +from common.file_util import load_obj, get_recipe_yml_obj +from common.point_cloud_util import normalize_point_cloud + + +def sample_pc_random(obj_path, num_points=10_000, apply_point_cloud_normalization=False): + """ + Chamfer evaluation + """ + vertices, faces = load_obj(obj_path) + vertices = vertices.reshape(1, vertices.shape[0], vertices.shape[1]) + vertices = torch.from_numpy(vertices) + faces = torch.from_numpy(faces) + point_cloud = sample_surface(faces, vertices, num_points=num_points) + if apply_point_cloud_normalization: + point_cloud = normalize_point_cloud(point_cloud) + # assert that the point cloud is normalized + max_dist_in_pc = torch.max(torch.sqrt(torch.sum((point_cloud ** 2), dim=1))) + threshold = 0.1 + assert abs(1.0 - max_dist_in_pc) <= threshold, f"PC of obj [{obj_path}] is not normalized, max distance in PC was [{abs(1.0 - max_dist_in_pc)}] but required to be <= [{threshold}]." + return point_cloud + + +def get_chamfer_distance(target_pc, gt_pc, device, num_points_in_pc, check_rot=False): + """ + num_points_in_pc - for sanity check + check_rot - was done for the tables since they are symmetric, and sketches are randomly flipped + it is ok to leave it on for the vase and chair, just makes it slower + """ + gt_pc = gt_pc.reshape(1, gt_pc.shape[0], gt_pc.shape[1]) # add batch dim + target_pc = target_pc.reshape(1, target_pc.shape[0], target_pc.shape[1]) # add batch dim + assert gt_pc.shape[1] == target_pc.shape[1] == num_points_in_pc + dist1, dist2, idx1, idx2 = chamfer_dist(target_pc.float().to(device), gt_pc.float().to(device)) + chamfer_distance = (torch.sum(dist1) + torch.sum(dist2)) / (gt_pc.shape[1] * 2) + if check_rot: + rot_mat = torch.tensor([[0.0, 0.0, -1.0], [0.0, 1.0, 0.0], [1.0, 0.0, 0.0]], dtype=torch.float64) + target_pc_rot = torch.matmul(rot_mat, target_pc.squeeze().t()).t().unsqueeze(0) + dist1, dist2, idx1, idx2 = chamfer_dist(target_pc_rot.float().to(device), gt_pc.float().to(device)) + chamfer_distance_rot = (torch.sum(dist1) + torch.sum(dist2)) / (gt_pc.shape[1] * 2) + return torch.min(chamfer_distance, chamfer_distance_rot) + return chamfer_distance + + +def save_as_obj_proc(pred_yml_file_path: Path, recipe_file_path: Path, results_dir: Path, out_dir: str, blender_exe: str, blend_file: str): + target_obj_file_path = results_dir.joinpath(out_dir, f'{pred_yml_file_path.stem}.obj') + print(f"Converting [{pred_yml_file_path}] to obj file [{target_obj_file_path}]") + save_obj_script_path = Path(__file__).parent.joinpath('..', 'dataset_processing', 'save_obj.py').resolve() + cmd = [f'{str(Path(blender_exe).expanduser())}', f'{str(Path(blend_file).expanduser())}', '-b', '--python', + f"{str(save_obj_script_path)}", '--', + '--recipe-file-path', str(recipe_file_path), + '--yml-file-path', str(pred_yml_file_path), + '--target-obj-file-path', str(target_obj_file_path), + '--ignore-sanity-check'] + print(" ".join(cmd)) + process = Popen(cmd, stdout=PIPE) + process.wait() + + +def test(opt): + recipe_file_path = Path(opt.dataset_dir, 'recipe.yml') + print(recipe_file_path) + if not recipe_file_path.is_file(): + raise Exception(f'No \'recipe.yml\' file found in path [{recipe_file_path}]') + recipe_yml_obj = get_recipe_yml_obj(str(recipe_file_path)) + + inputs_to_eval = get_inputs_to_eval(recipe_yml_obj) + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + camera_angles_to_process = recipe_yml_obj['camera_angles_train'] + recipe_yml_obj['camera_angles_test'] + camera_angles_to_process = [f'{a}_{b}' for a, b in camera_angles_to_process] + + param_descriptors = ParamDescriptors(recipe_yml_obj, inputs_to_eval, opt.use_regression) + param_descriptors_map = param_descriptors.get_param_descriptors_map() + detailed_vec_size = calc_prediction_vector_size(param_descriptors_map) + print(f"Prediction vector length is set to [{sum(detailed_vec_size)}]") + + # setup required dirs + required_dirs = ['barplot', + 'yml_gt', 'yml_predictions_pc', 'yml_predictions_sketch', + 'obj_gt', 'obj_predictions_pc', 'obj_predictions_sketch', + 'render_gt', 'render_predictions_pc', 'render_predictions_sketch', + 'sketch_gt'] + test_dir = Path(opt.dataset_dir, opt.phase) + test_dir_obj_gt = test_dir.joinpath('obj_gt') + results_dir = test_dir.joinpath(f'results_{opt.exp_name}') + results_dir.mkdir(exist_ok=True) + for dir in required_dirs: + results_dir.joinpath(dir).mkdir(exist_ok=True) + + # save the recipe to the results directory + shutil.copy(recipe_file_path, results_dir.joinpath('recipe.yml')) + + # find the best checkpoint (the one with the highest epoch number out of the saved checkpoints) + exp_dir = Path(opt.models_dir, opt.exp_name) + best_model_and_highest_epoch = None + highest_epoch = 0 + for ckpt_file in exp_dir.glob("*.ckpt"): + file_name = ckpt_file.name + if 'epoch' not in file_name: + continue + epoch_start_idx = file_name.find('epoch') + len('epoch') + epoch = int(file_name[epoch_start_idx:epoch_start_idx + 3]) + if epoch > highest_epoch: + best_model_and_highest_epoch = ckpt_file + highest_epoch = epoch + print(f'Best model with highest epoch is [{best_model_and_highest_epoch}]') + + batch_size = 1 + test_dataloaders = [] + test_dataloaders_types = [] + # pc + if InputType.pc in opt.input_type: + test_dataset_pc = DatasetPC(inputs_to_eval, device, param_descriptors_map, + opt.dataset_dir, opt.phase, random_pc=opt.random_pc, + gaussian=opt.gaussian, apply_point_cloud_normalization=opt.normalize_pc, + scanobjectnn=opt.scanobjectnn, augment_with_random_points=opt.augment_with_random_points) + test_dataloader_pc = DataLoader(test_dataset_pc, batch_size=batch_size, shuffle=False, + num_workers=2, prefetch_factor=2) + test_dataloaders.append(test_dataloader_pc) + test_dataloaders_types.append('pc') + # sketch + if InputType.sketch in opt.input_type: + test_dataset_sketch = DatasetSketch(inputs_to_eval, param_descriptors_map, + camera_angles_to_process, opt.pretrained_vgg, + opt.dataset_dir, opt.phase) + test_dataloader_sketch = DataLoader(test_dataset_sketch, batch_size=batch_size, shuffle=False, + num_workers=2, prefetch_factor=2) + test_dataloaders.append(test_dataloader_sketch) + test_dataloaders_types.append('sketch') + + pl_model = Model.load_from_checkpoint(str(best_model_and_highest_epoch), batch_size=1, + param_descriptors=param_descriptors, results_dir=results_dir, + test_dir=test_dir, models_dir=opt.models_dir, + test_dataloaders_types=test_dataloaders_types, test_input_type=opt.input_type, + exp_name=opt.exp_name) + + if True: + trainer = pl.Trainer(gpus=1) + trainer.test(model=pl_model, dataloaders=test_dataloaders, ckpt_path=best_model_and_highest_epoch) + + # save the validation and test bar-plots as image + barplot_target_dir = results_dir.joinpath('barplot') + for barplot_type in ['val', 'test']: + barplot_json_path = Path(opt.models_dir, opt.exp_name, f'{barplot_type}_barplot_top_1.json') + if not barplot_json_path.is_file(): + print(f"Could not find barplot [{barplot_json_path}] skipping copy") + continue + barplot_target_image_path = barplot_target_dir.joinpath(f'{barplot_type}_barplot.png') + title = "Validation Accuracy" if barplot_type == 'val' else "Test Accuracy" + gen_and_save_barplot(barplot_json_path, title, barplot_target_image_path=barplot_target_image_path) + shutil.copy(barplot_json_path, barplot_target_dir.joinpath(barplot_json_path.name)) + + gt_dir = results_dir.joinpath('yml_gt') + model_predictions_pc_dir = results_dir.joinpath('yml_predictions_pc') + model_predictions_sketch_dir = results_dir.joinpath('yml_predictions_sketch') + file_names = sorted([f.stem for f in test_dir_obj_gt.glob("*.obj")]) + + random_pc_dir = test_dir.joinpath("point_cloud_random") + if opt.scanobjectnn: + # [:-2] removed the _0 suffix + file_names = [str(f.stem)[:-2] for f in random_pc_dir.glob("*.npy")] + + # create all the obj from the prediction yaml files + # note that for pc we have one yml and for sketch we have multiple yml files (one for each camera angle) + if True: + cpu_count = multiprocessing.cpu_count() + print(f"Converting yml files to obj files using [{cpu_count}] processes") + for yml_dir, out_dir in [(gt_dir, 'obj_gt'), (model_predictions_pc_dir, 'obj_predictions_pc'), (model_predictions_sketch_dir, 'obj_predictions_sketch')]: + try: + # for each gt obj file we might have multiple yml files as predictions, like for the sketches + yml_files = sorted(yml_dir.glob("*.yml")) + # filter out existing + yml_files_filtered = [yml_file for yml_file in yml_files if not results_dir.joinpath(out_dir, f'{yml_file.stem}.obj').is_file()] + if out_dir == 'obj_gt' and not yml_files: + # COSEG (or any external ds for which we do not have ground truth yml files) + for obj_file in test_dir_obj_gt.glob("*.obj"): + print(f"shutil [{obj_file}]") + shutil.copy(obj_file, str(Path(results_dir, out_dir, f"{obj_file.stem}_gt.obj"))) + continue + save_as_obj_proc_partial = partial(save_as_obj_proc, + recipe_file_path=recipe_file_path, + results_dir=results_dir, + out_dir=out_dir, + blender_exe=opt.blender_exe, + blend_file=opt.blend_file) + p = multiprocessing.Pool(cpu_count) + p.map(save_as_obj_proc_partial, yml_files_filtered) + p.close() + p.join() + except Exception as e: + print(traceback.format_exc()) + print(repr(e)) + print("Done converting yml files to obj files") + + if True: + num_points_in_pc_for_chamfer = 10000 + chamfer_json = {'pc': {}, 'sketch': {}} + chamfer_summary_json = {'pc': {'chamfer_sum': 0.0, 'num_samples': 0}, 'sketch': {'chamfer_sum': 0.0, 'num_samples': 0}} + + for file_idx, file_name in enumerate(file_names): # for each unique test object + # get ground truth point cloud (uniform) + gt_file_name = file_name + if "_decimate_ratio_0" in file_name: + # edge case for comparing on the decimated ds + gt_file_name = gt_file_name.replace("_decimate_ratio_0_100", "_decimate_ratio_1_000") + gt_file_name = gt_file_name.replace("_decimate_ratio_0_010", "_decimate_ratio_1_000") + gt_file_name = gt_file_name.replace("_decimate_ratio_0_005", "_decimate_ratio_1_000") + assert "_decimate_ratio_0_100" not in gt_file_name + assert "_decimate_ratio_0_010" not in gt_file_name + assert "_decimate_ratio_0_005" not in gt_file_name + + if opt.scanobjectnn: + random_pc_path = random_pc_dir.joinpath(f"{file_name}_0.npy") + gt_pc = np.load(str(random_pc_path)) + gt_pc = torch.from_numpy(gt_pc).float() + num_points_in_pc_for_chamfer = 2048 + else: + gt_pc = sample_pc_random(results_dir.joinpath('obj_gt',f'{file_name}_gt.obj'), + num_points=num_points_in_pc_for_chamfer, + apply_point_cloud_normalization=opt.normalize_pc) + + for input_type, model_prediction_dir in [('pc', model_predictions_pc_dir), ('sketch', model_predictions_sketch_dir)]: + yml_files = sorted(model_prediction_dir.glob(f"{file_name}_*.yml")) + for yml_file in yml_files: + yml_file_base_name_no_ext = yml_file.stem + target_pc = sample_pc_random(results_dir.joinpath(f'obj_predictions_{input_type}', f'{yml_file_base_name_no_ext}.obj'), + num_points=num_points_in_pc_for_chamfer, + apply_point_cloud_normalization=opt.normalize_pc) + chamf_distance = get_chamfer_distance(target_pc, gt_pc, device, num_points_in_pc_for_chamfer, check_rot=(input_type == 'sketch')) + chamfer_summary_json[input_type]['chamfer_sum'] += chamf_distance.item() + chamfer_summary_json[input_type]['num_samples'] += 1 + chamfer_json[input_type][yml_file_base_name_no_ext] = chamf_distance.item() + + # compute overall average + if chamfer_summary_json['pc']['num_samples'] > 0: + chamfer_json['pc']['avg'] = chamfer_summary_json['pc']['chamfer_sum'] / chamfer_summary_json['pc']['num_samples'] + if chamfer_summary_json['sketch']['num_samples'] > 0: + chamfer_json['sketch']['avg'] = chamfer_summary_json['sketch']['chamfer_sum'] / chamfer_summary_json['sketch']['num_samples'] + + # save chamfer json to the results dir + with open(results_dir.joinpath("chamfer.json"), "w") as outfile: + json.dump(chamfer_json, outfile) + + print("Done calculating Chamfer distances") diff --git a/geocode/geocode_train.py b/geocode/geocode_train.py new file mode 100644 index 0000000..762206a --- /dev/null +++ b/geocode/geocode_train.py @@ -0,0 +1,150 @@ +import yaml +import json +from pathlib import Path +import torch +from torch.utils.data import DataLoader +import pytorch_lightning as pl +from pytorch_lightning.callbacks import ModelCheckpoint +from pytorch_lightning.loggers import NeptuneLogger +import neptune.new as neptune +from data.dataset_pc import DatasetPC +from data.dataset_sketch import DatasetSketch +from common.param_descriptors import ParamDescriptors +from common.file_util import get_recipe_yml_obj +from geocode_model import Model +from pytorch_lightning.trainer.supporters import CombinedLoader +from geocode_util import InputType, get_inputs_to_eval, calc_prediction_vector_size + + +def train(opt): + torch.set_printoptions(precision=4) + torch.multiprocessing.set_sharing_strategy('file_system') # to prevent "received 0 items of data" errors + recipe_file_path = Path(opt.dataset_dir, 'recipe.yml') + if not recipe_file_path.is_file(): + raise Exception(f'No \'recipe.yml\' file found in path [{recipe_file_path}]') + recipe_yml_obj = get_recipe_yml_obj(str(recipe_file_path)) + + inputs_to_eval = get_inputs_to_eval(recipe_yml_obj) + + top_k_acc = 2 + camera_angles_to_process = [f'{a}_{b}' for a, b in recipe_yml_obj['camera_angles_train']] + param_descriptors = ParamDescriptors(recipe_yml_obj, inputs_to_eval, use_regression=opt.use_regression) + param_descriptors_map = param_descriptors.get_param_descriptors_map() + detailed_vec_size = calc_prediction_vector_size(param_descriptors_map) + print(f"Prediction vector length is set to [{sum(detailed_vec_size)}]") + + # create datasets + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + train_loaders_map = {} + val_loaders_map = {} + + # pc + if InputType.pc in opt.input_type: + train_dataset_pc = DatasetPC(inputs_to_eval, device, param_descriptors_map, + opt.dataset_dir, "train", augment_with_random_points=True) + train_dataloader_pc = DataLoader(train_dataset_pc, batch_size=opt.batch_size, shuffle=True, num_workers=5, prefetch_factor=5) + val_dataset_pc = DatasetPC(inputs_to_eval, device, param_descriptors_map, + opt.dataset_dir, "val", augment_with_random_points=True) + val_dataloader_pc = DataLoader(val_dataset_pc, batch_size=opt.batch_size, shuffle=False, num_workers=5, prefetch_factor=5) + train_loaders_map['pc'] = train_dataloader_pc + val_loaders_map['pc'] = val_dataloader_pc + print(f"Point cloud train dataset size [{len(train_dataset_pc)}] val dataset size [{len(val_dataset_pc)}]") + + # sketch + if InputType.sketch in opt.input_type: + train_dataset_sketch = DatasetSketch(inputs_to_eval, param_descriptors_map, camera_angles_to_process, opt.pretrained_vgg, + opt.dataset_dir, "train") + train_dataloader_sketch = DataLoader(train_dataset_sketch, batch_size=opt.batch_size, shuffle=True, num_workers=5, prefetch_factor=5) + val_dataset_sketch = DatasetSketch(inputs_to_eval, param_descriptors_map, camera_angles_to_process, opt.pretrained_vgg, + opt.dataset_dir, "val") + val_dataloader_sketch = DataLoader(val_dataset_sketch, batch_size=opt.batch_size, shuffle=False, num_workers=5, prefetch_factor=5) + train_loaders_map['sketch'] = train_dataloader_sketch + val_loaders_map['sketch'] = val_dataloader_sketch + print(f"Sketch train dataset size [{len(train_dataset_sketch)}] val dataset size [{len(val_dataset_sketch)}]") + + combined_train_dataloader = CombinedLoader(train_loaders_map, mode="max_size_cycle") + combined_val_dataloader = CombinedLoader(val_loaders_map, mode="max_size_cycle") + + if InputType.pc in opt.input_type and InputType.sketch in opt.input_type: + assert ( len(camera_angles_to_process) * len(train_dataset_pc) ) == len(train_dataset_sketch) + assert ( len(camera_angles_to_process) * len(val_dataset_pc) ) == len(val_dataset_sketch) + + print(f"Experiment name [{opt.exp_name}]") + + exp_dir = Path(opt.models_dir, opt.exp_name) + exp_dir.mkdir(exist_ok=True) + + neptune_short_id = None + neptune_short_id_file_path = exp_dir.joinpath('neptune_session.json') + if neptune_short_id_file_path.is_file(): + with open(neptune_short_id_file_path, 'r') as neptune_short_id_file: + try: + neptune_session_json = json.load(neptune_short_id_file) + if 'short_id' in neptune_session_json: + neptune_short_id = neptune_session_json['short_id'] + print(f'Continuing Neptune run [{neptune_short_id}]') + except: + print("Could not resume neptune session") + + # create/load NeptuneLogger + neptune_logger = None + neptune_config_file_path = Path(__file__).parent.joinpath('..', 'config', 'neptune_config.yml').resolve() + if neptune_config_file_path.is_file(): + print(f"Found neptune config file [{neptune_config_file_path}]") + with open(neptune_config_file_path) as neptune_config_file: + config = yaml.safe_load(neptune_config_file) + api_token = config['neptune']['api_token'] + project = config['neptune']['project'] + tags = ["train"] + if neptune_short_id: + neptune_logger = NeptuneLogger( run=neptune.init(run=neptune_short_id, project=project, api_token=api_token, tags=tags), log_model_checkpoints=False ) + else: + # log_model_checkpoints=False avoids saving the models to Neptune + neptune_logger = NeptuneLogger(api_key=api_token, project=project, tags=tags, log_model_checkpoints=False) + if neptune_short_id is None: + # new experiment + neptune_short_id = neptune_logger.run.fetch()['sys']['id'] # e.g. IN-105 (-) + with open(neptune_short_id_file_path, 'w') as neptune_short_id_file: + json.dump({'short_id': neptune_short_id}, neptune_short_id_file) + print(f'Started a new Neptune.ai run with id [{neptune_short_id}]') + + # log parameters to Neptune + params = { + "exp_name": opt.exp_name, + "lr": 1e-2, + "bs": opt.batch_size, + "n_parameters": len(inputs_to_eval), + "sched_step_size": 20, + "sched_gamma": 0.9, + "normalize_embeddings": opt.normalize_embeddings, + "increase_net_size": opt.increase_network_size, + "pretrained_vgg": opt.pretrained_vgg, + "use_regression": opt.use_regression, + } + if neptune_logger: + neptune_logger.run['parameters'] = params + + checkpoint_callback = ModelCheckpoint( + dirpath=exp_dir, + filename='ise-epoch{epoch:03d}-val_loss{val/loss/total:.2f}-val_acc{val/acc_top1/avg:.2f}', + auto_insert_metric_name=False, + save_last=True, + monitor="val/acc_top1/avg", + mode="max", + save_top_k=3) + + trainer = pl.Trainer(gpus=1, max_epochs=opt.nepoch, logger=neptune_logger, callbacks=[checkpoint_callback]) + last_ckpt_file_name = f"{checkpoint_callback.CHECKPOINT_NAME_LAST}{checkpoint_callback.FILE_EXTENSION}" # "last.ckpt" by default + last_checkpoint_file_path = exp_dir.joinpath(last_ckpt_file_name) + if last_checkpoint_file_path.is_file(): + print(f"Loading checkpoint file [{last_checkpoint_file_path}]...") + pl_model = Model.load_from_checkpoint(str(last_checkpoint_file_path), + param_descriptors=param_descriptors, + trainer=trainer, + models_dir=opt.models_dir, + exp_name=opt.exp_name) + else: + pl_model = Model(top_k_acc, opt.batch_size, detailed_vec_size, opt.increase_network_size, opt.normalize_embeddings, opt.pretrained_vgg, + opt.input_type, inputs_to_eval, params['lr'], params['sched_step_size'], params['sched_gamma'], opt.exp_name, + trainer=trainer, param_descriptors=param_descriptors, models_dir=opt.models_dir, use_regression=opt.use_regression) + trainer.fit(pl_model, train_dataloaders=combined_train_dataloader, val_dataloaders=combined_val_dataloader, ckpt_path=None) diff --git a/geocode/geocode_util.py b/geocode/geocode_util.py new file mode 100644 index 0000000..9309002 --- /dev/null +++ b/geocode/geocode_util.py @@ -0,0 +1,35 @@ +from enum import Enum +from typing import Dict +from common.param_descriptors import ParamDescriptor + + +class InputType(Enum): + sketch = 'sketch' + pc = 'pc' + + def __str__(self): + return self.value + + def __eq__(self, other): + return self.value == other.value + + +def get_inputs_to_eval(recipe_yml_obj): + inputs_to_eval = [] + for param_name, param_dict in recipe_yml_obj['dataset_generation'].items(): + is_vector = False + for axis in ['x', 'y', 'z']: + if axis in param_dict: + inputs_to_eval.append(f'{param_name} {axis}') + is_vector = True + if not is_vector: + inputs_to_eval.append(param_name) + print("Inputs that will be evaluated:") + print("\t" + "\n\t".join(inputs_to_eval)) + return inputs_to_eval + + +def calc_prediction_vector_size(param_descriptors_map: Dict[str, ParamDescriptor]): + detailed_vec_size = [param_descriptor.num_classes for param_name, param_descriptor in param_descriptors_map.items()] + print(f"Found [{len(detailed_vec_size)}] parameters with a combined number of classes of [{sum(detailed_vec_size)}]") + return detailed_vec_size diff --git a/resources/teaser.png b/resources/teaser.png new file mode 100644 index 0000000..2651a6a Binary files /dev/null and b/resources/teaser.png differ diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/scripts/download_ds.py b/scripts/download_ds.py new file mode 100644 index 0000000..314b48b --- /dev/null +++ b/scripts/download_ds.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 + +import gdown +import zipfile +import hashlib +import requests +import traceback +from pathlib import Path +from common.domain import Domain +from argparse import ArgumentParser + + +def download_file(url, target_file_path): + print(f"Downloading file from [{url}] as [{target_file_path}]") + req = requests.get(url, allow_redirects=True) + with open(target_file_path, 'wb') as target_file: + target_file.write(req.content) + + +def download_ds(args): + datasets_dir = None + if args.datasets_dir: + datasets_dir = Path(args.datasets_dir) + if not datasets_dir.is_dir(): + raise Exception(f'Given datasets path [{datasets_dir}] is not an existing directory.') + models_dir = None + if args.models_dir: + models_dir = Path(args.models_dir) + if not models_dir.is_dir(): + raise Exception(f'Given models path [{models_dir}] is not an existing directory.') + blends_dir = None + if args.blends_dir: + blends_dir = Path(args.blends_dir) + if not blends_dir.is_dir(): + raise Exception(f'Given blends path [{blends_dir}] is not an existing directory.') + + if args.domain == Domain.chair: + md5 = "27c283fa6893b23400a9bba6aca92854" + ds_url = "https://drive.google.com/uc?id=1Mmq1dVuMjoTZxtcC2wm4qMnIIeI_0OI-" + ds_zip_file_name = "ChairDataset.zip" + best_epoch = 585 + elif args.domain == Domain.vase: + md5 = "1200bfb9552513ea6c9a3b9050af470e" + ds_url = "https://drive.google.com/uc?id=1nCK4Jx2uIKcK5hubhKN29b8-OxdASEfl" + ds_zip_file_name = "VaseDataset.zip" + best_epoch = 573 + elif args.domain == Domain.table: + md5 = "c7a0fc73c2b3f39dcd02f8cd3380d9dd" + ds_url = "https://drive.google.com/uc?id=1W8s5MvNKPY56wKB0QJA-7VIqFBMd-U2k" + ds_zip_file_name = "TableDataset.zip" + best_epoch = 537 + else: + raise Exception(f'Domain [{args.domain}] is not recognized.') + + if args.datasets_dir: + target_ds_zip_file_path = datasets_dir.joinpath(ds_zip_file_name) + # download requested dataset zip file from Google Drive + if not target_ds_zip_file_path.is_file(): + gdown.download(ds_url, target_ds_zip_file_path, quiet=False) + else: + print(f"Skipping downloading dataset from Google Drive, file [{target_ds_zip_file_path}] already exists.") + + unzipped_dataset_dir = datasets_dir.joinpath(f"{str(args.domain).title()}Dataset") + + if not unzipped_dataset_dir.is_dir(): + # verify md5 + print("Verifying MD5 hash...") + assert hashlib.md5(open(target_ds_zip_file_path, 'rb').read()).hexdigest() == md5 + + print("Unzipping dataset...") + with zipfile.ZipFile(target_ds_zip_file_path, 'r') as target_ds_zip_file: + target_ds_zip_file.extractall(datasets_dir) + else: + print(f"Skipping dataset unzip, directory [{unzipped_dataset_dir}] already exists.") + + release_url = "https://github.com/threedle/GeoCode/releases/latest/download" + + if args.models_dir: + best_ckpt_file_name = f"procedural_{args.domain}_last_ckpt.zip" + latest_ckpt_file_name = f"procedural_{args.domain}_epoch{best_epoch:03d}_ckpt.zip" + exp_target_dir = models_dir.joinpath(f"exp_geocode_{args.domain}") + exp_target_dir.mkdir(exist_ok=True) + + best_ckpt_url = f"{release_url}/{best_ckpt_file_name}" + best_ckpt_file_path = exp_target_dir.joinpath(best_ckpt_file_name) + download_file(best_ckpt_url, best_ckpt_file_path) + + print(f"Unzipping checkpoint file [{best_ckpt_file_path}]...") + with zipfile.ZipFile(best_ckpt_file_path, 'r') as best_ckpt_file: + best_ckpt_file.extractall(exp_target_dir) + + latest_ckpt_url = f"{release_url}/{latest_ckpt_file_name}" + latest_ckpt_file_path = exp_target_dir.joinpath(latest_ckpt_file_name) + download_file(latest_ckpt_url, latest_ckpt_file_path) + + print(f"Unzipping checkpoint file [{latest_ckpt_file_path}]...") + with zipfile.ZipFile(latest_ckpt_file_path, 'r') as latest_ckpt_file: + latest_ckpt_file.extractall(exp_target_dir) + + if args.blends_dir: + blend_file_name = f"procedural_{args.domain}.blend" + blend_file_path = blends_dir.joinpath(blend_file_name) + blend_url = f"{release_url}/{blend_file_name}" + download_file(blend_url, blend_file_path) + + print("Done") + + +def main(): + parser = ArgumentParser(prog='download_ds') + parser.add_argument('--domain', type=Domain, choices=list(Domain), required=True, help='The domain name to download the dataset for.') + parser.add_argument('--datasets-dir', type=str, default=None, help='The directory to download the dataset to.') + parser.add_argument('--models-dir', type=str, default=None, help='The directory to download checkpoint file to.') + parser.add_argument('--blends-dir', type=str, default=None, help='The directory to download blend file to.') + + try: + args = parser.parse_args() + download_ds(args) + except Exception as e: + print(repr(e)) + print(traceback.format_exc()) + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..64d340d --- /dev/null +++ b/setup.py @@ -0,0 +1,2 @@ +from setuptools import setup, find_packages +setup(name='geocode', packages=find_packages()) diff --git a/stability_metric/__init__.py b/stability_metric/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/stability_metric/stability.py b/stability_metric/stability.py new file mode 100644 index 0000000..5826da7 --- /dev/null +++ b/stability_metric/stability.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python3 + +import bpy +import sys +import argparse +import numpy as np +import traceback +from mathutils import Vector + +from pathlib import Path +import importlib + +def import_parents(level=1): + global __package__ + file = Path(__file__).resolve() + parent, top = file.parent, file.parents[level] + + sys.path.append(str(top)) + try: + sys.path.remove(str(parent)) + except ValueError: + # already removed + pass + + __package__ = '.'.join(parent.parts[len(top.parts):]) + importlib.import_module(__package__) + +if __name__ == '__main__' and __package__ is None: + import_parents(level=1) + +from common.bpy_util import normalize_scale, select_obj +from common.intersection_util import detect_cross_intersection + + +class DropSimulator: + def __init__(self, args): + self.drop_height = 0.1 + self.duration_sec = 5 + self.skip_components_check = args.skip_components_check + self.apply_normalization = args.apply_normalization + self.obj_file_path = Path(args.obj_path).expanduser() + + def simulate(self): + print(f"Importing object file [{self.obj_file_path}]") + bpy.ops.import_scene.obj(filepath=str(self.obj_file_path), use_split_objects=False) + obj = bpy.context.selected_objects[0] + obj.data.materials.clear() + select_obj(obj) + if self.apply_normalization: + print("Normalizing object...") + normalize_scale(obj) + # set origin to center of mass + bpy.ops.object.origin_set(type='ORIGIN_CENTER_OF_MASS', center='MEDIAN') + vertices = np.array([(obj.matrix_world @ v.co) for v in obj.data.vertices]) + # verify that the object is normalized + max_diff = 0.05 + max_dist_from_center = abs(1.0 - np.max(np.sqrt(np.sum((vertices ** 2), axis=1)))) + assert max_dist_from_center < max_diff, f"Point cloud is not normalized [{max_dist_from_center} > {max_diff}] for sample [{self.obj_file_path.name}]. If this is an external dataset, please consider adding --apply-normalization flag." + # position the object at drop height + obj.location = Vector((0, 0, obj.location.z - min(vertices[:, 2]) + self.drop_height)) + height_before_drop = max(vertices[:, 2]) - min(vertices[:, 2]) + + # apply rigid body simulation + bpy.ops.rigidbody.object_add() + frame_end = self.duration_sec * 25 + area = [a for a in bpy.context.screen.areas if a.type == "VIEW_3D"][0] + with bpy.context.temp_override(area=area): + select_obj(obj) + bpy.ops.rigidbody.bake_to_keyframes(frame_start=1, frame_end=frame_end, step=1) + bpy.context.scene.frame_current = frame_end + obj.data.update() + bpy.context.view_layer.update() + print("Simulation completed") + self.eval(obj, height_before_drop) + + def eval(self, obj, height_before_drop): + vertices = np.array([(obj.matrix_world @ v.co) for v in obj.data.vertices]) + height_after_drop = max(vertices[:, 2]) - min(vertices[:, 2]) + score = min(height_after_drop, height_before_drop) / max(height_after_drop, height_before_drop) + score = 1.0 if score > 1.0 else score + print(f"Height before simulation [{height_before_drop:.5f}]") + print(f"Height after simulation [{height_after_drop:.5f}]") + print(f"Score [{score:.5f}]") + print(f"is_stable (score > 0.98) [{score > 0.98}]") + self.reset_simulation(obj) + # structural evaluation + if self.skip_components_check: + print("is_structurally_valid (shape is connected) [True] (check skipped)") + else: + print("Checking structural validity...") + obj_is_valid = self.is_structurally_connected() + print(f"is_structurally_valid (shape is connected) [{obj_is_valid}]") + + def is_structurally_connected(self): + """ + return True if all the parts that make the shape are reachable from any other part + two parts are connected if they are intersecting or there is a path from one part + to the other that passes only through intersecting parts + """ + bpy.ops.import_scene.obj(filepath=str(self.obj_file_path), use_split_objects=False) + obj = bpy.context.selected_objects[0] + select_obj(obj) + bpy.ops.object.mode_set(mode='EDIT') + bpy.ops.mesh.separate(type='LOOSE') + bpy.ops.object.mode_set(mode='OBJECT') + parts = bpy.context.selected_objects + # in the beginning, each part is put into a separate set + components = [] + for part in parts: + components.append({part}) + + idx_a = 0 + while idx_a + 1 < len(components): + component_a = components[idx_a] + found_intersection = False + for idx_b in range(idx_a + 1, len(components)): + component_b = components[idx_b] + for part_a in component_a: + for part_b in component_b: + assert part_a.name != part_b.name + if len(detect_cross_intersection(part_a, part_b)) > 0: + components.remove(component_a) + components.remove(component_b) + components.append(component_a.union(component_b)) + found_intersection = True + break + if found_intersection: + break + if found_intersection: + break + if not found_intersection: + idx_a += 1 + # note that we can 'break' here and return False if we are only looking to have a single connected component + bpy.ops.object.delete() + return len(components) <= 1 + + @staticmethod + def reset_simulation(obj): + bpy.context.scene.frame_current = 0 + obj.data.update() + bpy.context.view_layer.update() + bpy.ops.object.delete() + + +def main(): + if '--' in sys.argv: + # refer to https://b3d.interplanety.org/en/how-to-pass-command-line-arguments-to-a-blender-python-script-or-add-on/ + argv = sys.argv[sys.argv.index('--') + 1:] + else: + raise Exception("Expected \'--\' followed by arguments to the script") + + parser = argparse.ArgumentParser(prog='stability') + parser.add_argument('--obj-path', type=str, required=True, help='Path to the object to test') + parser.add_argument('--apply-normalization', action='store_true', default=False, help='Apply normalization on the object upon importing') + parser.add_argument('--skip-components-check', action='store_true', default=False, help='Do not check that the shape is structurally valid') + + try: + args = parser.parse_known_args(argv)[0] + drop_simulator = DropSimulator(args) + drop_simulator.simulate() + except Exception as e: + print(repr(e)) + print(traceback.format_exc()) + + +if __name__ == '__main__': + main() diff --git a/stability_metric/stability_parallel.py b/stability_metric/stability_parallel.py new file mode 100644 index 0000000..203bd7e --- /dev/null +++ b/stability_metric/stability_parallel.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 + +import json +import random +import argparse +import traceback +import multiprocessing +import subprocess +from subprocess import Popen +from functools import partial +from pathlib import Path + + +def calculate_stability_proc(pred_obj_file_path, apply_normalization, skip_components_check, blender_exe: Path): + print(f"Calculating stability for object [{pred_obj_file_path}]") + simulation_blend_file_path = Path(__file__).parent.joinpath('stability_simulation.blend').resolve() + stability_script_path = Path(__file__).parent.joinpath('stability.py').resolve() + cmd = [str(blender_exe.expanduser()), + str(simulation_blend_file_path), + '-b', '--python', str(stability_script_path), '--', + 'sim-obj', '--obj-path', str(pred_obj_file_path)] + if apply_normalization: + cmd.append('--apply-normalization') + if skip_components_check: + cmd.append('--skip-components-check') + print(" ".join(cmd)) + process = Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True) + out, err = process.communicate() + result = out.splitlines() + score_str_list = [line for line in result if 'Score [' in line] + structurally_valid_str_list = [line for line in result if 'is_structurally_valid' in line] + assert score_str_list, out + score = float(score_str_list[0][-8:-1]) + assert structurally_valid_str_list, out + is_structurally_valid = True if 'True' in structurally_valid_str_list[0] else False + assert score > 0.1 + return score, is_structurally_valid + + +def sim_dir_parallel(args): + cpu_count = multiprocessing.cpu_count() + stability_json = {} + count_stable = 0 + count_structurally_valid = 0 + count_good = 0 + dir_path = Path(args.dir_path).resolve() + print(f"Calculating stability for dir [{dir_path}] with [{cpu_count}] processes") + try: + obj_files = sorted(dir_path.glob("*.obj")) + print(len(obj_files)) + if args.limit and args.limit < len(obj_files): + obj_files = random.sample(obj_files, args.limit) + blender_exe = Path(args.blender_exe).resolve() + calculate_stability_proc_partial = partial(calculate_stability_proc, + apply_normalization=args.apply_normalization, + skip_components_check=args.skip_components_check, + blender_exe=blender_exe) + p = multiprocessing.Pool(cpu_count) + stability_results = p.map(calculate_stability_proc_partial, obj_files) + p.close() + p.join() + for obj_file_idx, obj_file in enumerate(obj_files): + stability_json[str(obj_file)] = stability_results[obj_file_idx] + score = stability_results[obj_file_idx][0] + is_structurally_valid = stability_results[obj_file_idx][1] + count_stable += 1 if score > 0.98 else 0 + count_structurally_valid += 1 if is_structurally_valid else 0 + count_good += 1 if (score > 0.98 and is_structurally_valid) else 0 + except Exception as e: + print(traceback.format_exc()) + print(repr(e)) + sample_count = len(stability_json) + print(f"# stable samples [{count_stable}] out of total [{sample_count}]") + print(f"# structurally valid samples [{count_structurally_valid}] out of total [{sample_count}]") + print(f"# good samples [{count_good}] out of total [{sample_count}] = [{(count_good/sample_count) * 100}%]") + # save the detailed results to a json file + stability_json['execution-details'] = {} + stability_json['execution-details']['dir-path'] = str(dir_path) + json_result_file_path = Path(__file__).parent.joinpath('stability_results.json').resolve() + with open(json_result_file_path, 'w') as json_result_file: + json.dump(stability_json, json_result_file) + print(f"Results per .obj file were saved to [{json_result_file_path}]") + + +def main(): + parser = argparse.ArgumentParser(prog='stability_parallel') + parser.add_argument('--dir-path', type=str, required=True, help='Path to the dir to test with the \'stability metric\'') + parser.add_argument('--blender-exe', type=str, required=True, help='Path to Blender executable') + parser.add_argument('--skip-components-check', action='store_true', default=False, help='Skip checking if the shape is structurally valid') + parser.add_argument('--apply-normalization', action='store_true', default=False, help='Apply normalization on the imported objects') + parser.add_argument('--limit', type=int, default=None, help='Limit the number of shapes that will be evaluated, randomly selected shapes will be tested') + + try: + args = parser.parse_args() + sim_dir_parallel(args) + except Exception as e: + print(repr(e)) + print(traceback.format_exc()) + + +if __name__ == '__main__': + main() diff --git a/stability_metric/stability_simulation.blend b/stability_metric/stability_simulation.blend new file mode 100644 index 0000000..42b32d8 Binary files /dev/null and b/stability_metric/stability_simulation.blend differ diff --git a/visualize_results/visualize.py b/visualize_results/visualize.py new file mode 100644 index 0000000..a46c8c0 --- /dev/null +++ b/visualize_results/visualize.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 + +import sys +import bpy +import math +import random +import argparse +import traceback +import mathutils +from pathlib import Path +from mathutils import Vector +import importlib + +def import_parents(level=1): + global __package__ + file = Path(__file__).resolve() + parent, top = file.parent, file.parents[level] + sys.path.append(str(top)) + try: + sys.path.remove(str(parent)) + except ValueError: + # already removed + pass + + __package__ = '.'.join(parent.parts[len(top.parts):]) + importlib.import_module(__package__) + +if __name__ == '__main__' and __package__ is None: + import_parents(level=1) + +from common.file_util import hash_file_name +from common.bpy_util import clean_scene, setup_lights, select_objs, normalize_scale, del_obj, look_at + + +def add_3d_text(obj_to_align_with, text): + """ + Adds 3D text object in front of the normalized object + """ + bpy.ops.object.select_all(action='DESELECT') + font_curve = bpy.data.curves.new(type="FONT", name="Font Curve") + font_curve.body = text + font_obj = bpy.data.objects.new(name="Font Object", object_data=font_curve) + bpy.context.scene.collection.objects.link(font_obj) + font_obj.select_set(True) + bpy.context.view_layer.objects.active = font_obj + bpy.ops.object.transform_apply(location=False, rotation=False, scale=True) + bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY', center='MEDIAN') + font_obj.location.x = obj_to_align_with.location.x + font_obj.location.y = obj_to_align_with.location.y + font_obj.location.z = 0 + font_obj.scale.x *= 0.2 + font_obj.scale.y *= 0.2 + font_obj.location.y -= 1 + return font_obj + + +def visualize_results(args): + """ + Before using this method, a test dataset should be evaluated using the model + """ + test_ds_dir = Path(args.dataset_dir, args.phase).expanduser() + if not test_ds_dir.is_dir(): + raise Exception(f"Expected a \'{args.phase}\' dataset directory with 3D object to evaluate") + + results_dir = test_ds_dir.joinpath(f'results_{args.exp_name}') + model_predictions_pc_dir = results_dir.joinpath('yml_predictions_pc') + if not model_predictions_pc_dir.is_dir(): + raise Exception(f"Expected a \'results_{args.exp_name}/yml_predictions_pc\' directory with predictions from point clouds") + + model_predictions_sketch_dir = results_dir.joinpath('yml_predictions_sketch') + if not model_predictions_sketch_dir.is_dir(): + raise Exception(f"Expected a \'results_{args.exp_name}/yml_predictions_sketch\' directory with predictions from skeches") + + obj_gt_dir = results_dir.joinpath('obj_gt') + render_gt_dir = results_dir.joinpath('render_gt') + obj_predictions_pc_dir = results_dir.joinpath('obj_predictions_pc') + render_predictions_pc_dir = results_dir.joinpath('render_predictions_pc') + obj_predictions_sketch_dir = results_dir.joinpath('obj_predictions_sketch') + render_predictions_sketch_dir = results_dir.joinpath('render_predictions_sketch') + + work = [ + (obj_gt_dir, render_gt_dir, "GT"), # render original 3D objs + (obj_predictions_pc_dir, render_predictions_pc_dir, "PRED FROM PC"), # render predictions from point cloud input + (obj_predictions_sketch_dir, render_predictions_sketch_dir, "PRED FROM SKETCH") # render predictions from sketch input + ] + + try: + clean_scene(start_with_strings=["Camera", "Light"]) + setup_lights() + # hide the main collections + bpy.context.scene.view_layers['View Layer'].layer_collection.children['Main'].hide_viewport = True + bpy.context.scene.view_layers['View Layer'].layer_collection.children['Main'].exclude = True + for obj_dir, render_dir, title in work: + file_names = sorted([f.stem for f in obj_dir.glob("*.obj")]) + if args.parallel > 1: + file_names = [file for file in file_names if hash_file_name(file) % args.parallel == args.mod] + for file_name in file_names: + original_obj_file_path = obj_dir.joinpath(f'{file_name}.obj') + bpy.ops.import_scene.obj(filepath=str(original_obj_file_path)) + imported_object = bpy.context.selected_objects[0] + imported_object.hide_render = False + imported_object.data.materials.clear() + normalize_scale(imported_object) + title_obj = add_3d_text(imported_object, title) + render_images(render_dir, file_name) + select_objs(title_obj, imported_object) + bpy.ops.object.delete() + except Exception as e: + print(repr(e)) + print(traceback.format_exc()) + + +def render_images(target_dir: Path, file_name, suffix=None): + # euler setting + camera_angles = [ + [-30.0, -35.0] + ] + radius = 2 + eulers = [mathutils.Euler((math.radians(camera_angle[0]), 0.0, math.radians(camera_angle[1])), 'XYZ') for + camera_angle in camera_angles] + + for i, eul in enumerate(eulers): + target_file_name = f"{file_name}{(f'_{suffix}' if suffix else '')}_at_{camera_angles[i][0]:.1f}_{camera_angles[i][1]:.1f}.png" + target_file = target_dir.joinpath(target_file_name) + + # camera setting + cam_pos = mathutils.Vector((0.0, -radius, 0.0)) + cam_pos.rotate(eul) + if i < 4: + rand_x = random.uniform(-2.0, 2.0) + rand_z = random.uniform(-5.0, 5.0) + eul_perturb = mathutils.Euler((math.radians(rand_x), 0.0, math.radians(rand_z)), 'XYZ') + cam_pos.rotate(eul_perturb) + + scene = bpy.context.scene + bpy.ops.object.camera_add(enter_editmode=False, location=cam_pos) + new_camera = bpy.context.active_object + new_camera.name = "camera_tmp" + new_camera.data.name = "camera_tmp" + new_camera.data.lens_unit = 'FOV' + new_camera.data.angle = math.radians(60) + look_at(new_camera, Vector((0.0, 0.0, 0.0))) + + # render + # scene.render.engine = 'BLENDER_EEVEE' # don't need anything fancy here + scene.camera = new_camera + scene.render.filepath = str(target_file) + scene.render.resolution_x = 224 + scene.render.resolution_y = 224 + bpy.context.scene.cycles.samples = 5 + # bpy.context.scene.eevee.taa_render_samples = 32 + # disable the sketch shader + bpy.context.scene.render.use_freestyle = False + bpy.ops.render.render(write_still=True) + + # prepare for the next camera + del_obj(new_camera) + + +def main(): + if '--' in sys.argv: + # refer to https://b3d.interplanety.org/en/how-to-pass-command-line-arguments-to-a-blender-python-script-or-add-on/ + argv = sys.argv[sys.argv.index('--') + 1:] + else: + raise Exception("Expected \'--\' followed by arguments to the script") + + parser = argparse.ArgumentParser(prog='dataset_generator') + parser.add_argument('--dataset-dir', type=str, required=True, help='Path to dataset directory') + parser.add_argument('--phase', type=str, required=True, help='E.g. train, test or val') + parser.add_argument('--exp-name', type=str, required=True) + parser.add_argument('--parallel', type=int, default=1) + parser.add_argument('--mod', type=int, default=None) + + try: + args = parser.parse_known_args(argv)[0] + visualize_results(args) + except Exception as e: + print(repr(e)) + print(traceback.format_exc()) + + +if __name__ == '__main__': + main()