diff --git a/sim/scripts/run_sim.py b/sim/scripts/run_sim.py index b5f4c4d..3bffc36 100644 --- a/sim/scripts/run_sim.py +++ b/sim/scripts/run_sim.py @@ -56,8 +56,10 @@ def init_and_sim_fastener(C, fastener, scenes, height=60): fastener.select_set(True) apply_texture(C, fastener) - x = random.randint(-2000, 2000) / 1000.0 - y = random.randint(-2000, 2000) / 1000.0 + #x = random.randint(-2000, 2000) / 1000.0 + #y = random.randint(-2000, 2000) / 1000.0 + x = 0 + y = 0 x_angle = math.radians(random.randint(0, 360)) @@ -67,8 +69,10 @@ def init_and_sim_fastener(C, fastener, scenes, height=60): y_angle = math.radians(random.randint(150, 210)) z_angle = math.radians(random.randint(0, 359)) - fastener.location = (x, y, 60) - fastener.rotation_euler = (x_angle, y_angle, z_angle) + #fastener.location = (x, y, 60) + fastener.location = (x, y, 10) + fastener.rotation_euler = (0, 0, 0) + #fastener.rotation_euler = (x_angle, y_angle, z_angle) bpy.ops.rigidbody.objects_add(type='ACTIVE') bpy.ops.object.origin_set(type='ORIGIN_CENTER_OF_VOLUME') @@ -78,7 +82,7 @@ def init_and_sim_fastener(C, fastener, scenes, height=60): post_rotation = fastener.matrix_world.to_euler() return post_rotation[-1] -def play_scene(C, until_frame=250): +def play_scene(C, until_frame=230): frame = 0 for _ in range(until_frame): C.scene.frame_set(frame) @@ -104,6 +108,7 @@ def rotate_and_take_image(C, fastener, output_model_path, model_name, z_angle=No take_image(C, rotate_image_path) def run_sim(model_path, output_path, copies, label): + copies = 1 output_path.mkdir(exist_ok=True) model_name = model_path.stem @@ -111,7 +116,12 @@ def run_sim(model_path, output_path, copies, label): with open(label_path, 'w') as f: json.dump(label, f) - for copy in range(copies): + camera = bpy.data.objects.get('top-down-cam') + increments = 75 + increment = 2 + start_z = 260 + + for i in range(increments): C = bpy.context scenes = {scene.name: scene for scene in bpy.data.scenes} bpy.ops.import_mesh.stl(filepath=str(model_path), axis_up='X', axis_forward="-Y") @@ -119,10 +129,13 @@ def run_sim(model_path, output_path, copies, label): z_angle_of_fastener = init_and_sim_fastener(C, fastener, scenes) - output_model_path = output_path / f"{model_name}_{copy}" + + z_height = start_z + increment*i + camera.location = (0, 0, z_height) + output_model_path = output_path / f"{model_name}" output_model_path.mkdir(exist_ok=True) - top_down_image_path = output_model_path / f"{model_name}_top.jpg" + top_down_image_path = output_model_path / f"{model_name}_{z_height:.1f}.jpg" C.window.scene = scenes[TOPDOWN_SCENE] play_scene(C) @@ -132,8 +145,8 @@ def run_sim(model_path, output_path, copies, label): C.window.scene = scenes[SIDEON_SCENE] play_scene(C) - z_angle_of_fastener += math.radians(random.randint(-15, 15)) - rotate_and_take_image(C, fastener, output_model_path, model_name, z_angle=z_angle_of_fastener) + #z_angle_of_fastener += math.radians(random.randint(-15, 15)) + #rotate_and_take_image(C, fastener, output_model_path, model_name, z_angle=z_angle_of_fastener) bpy.data.objects.remove(fastener, do_unlink=True) diff --git a/sim/scripts/sim_runner.py b/sim/scripts/sim_runner.py index 995918e..5e7324b 100644 --- a/sim/scripts/sim_runner.py +++ b/sim/scripts/sim_runner.py @@ -2,6 +2,7 @@ import json import logging import os +import shutil import subprocess from pathlib import Path @@ -11,7 +12,8 @@ @click.option("--output_folder_path", "output_folder_path", required=True, help="Output folder to put generated images") @click.option("--copies", "copies", required=True, help="Number of copies for each fastener") @click.option("--batch_size", "batch_size", required=True, default=10, help="Number of models to handle at a time") -def main(input_folder_path, output_folder_path, copies, batch_size): +@click.option("--curr_state_path", "curr_state_path", required=False, default=None, help="Path to curr_state.json from previous run") +def main(input_folder_path, output_folder_path, copies, batch_size, curr_state_path): logging.basicConfig(level=logging.INFO) input_folder_path = Path(input_folder_path) @@ -19,18 +21,36 @@ def main(input_folder_path, output_folder_path, copies, batch_size): output_folder_path = Path(output_folder_path) assert output_folder_path.exists(), f"{output_folder_path} does not exist" - cad_models = set([d for d in os.listdir(input_folder_path) if os.path.isdir(input_folder_path / d)]) - existing = set([d for d in os.listdir(output_folder_path) if os.path.isdir(output_folder_path / d)]) + remaining = [] + if curr_state_path: + curr_state_path = Path(curr_state_path) + assert curr_state_path.exists(), f"{curr_state_path} does not exist" + with open(curr_state_path) as f: + curr_state = json.load(f) - logging.info(f"{len(cad_models)} total") - logging.info(f"{len(existing)} existing") - logging.info(f"{len(cad_models - existing)} to convert") + to_delete = [output_folder_path / model_id for model_id in curr_state["to_convert"] + curr_state["failed"]] + for delete_path in to_delete: + if delete_path.exists(): + shutil.rmtree(delete_path) + + remaining = curr_state["remaining"] + curr_state["to_convert"] + curr_state["failed"] + else: + + cad_models = set([d for d in os.listdir(input_folder_path) if os.path.isdir(input_folder_path / d)]) + existing = set([d for d in os.listdir(output_folder_path) if os.path.isdir(output_folder_path / d)]) + + logging.info(f"{len(cad_models)} total") + logging.info(f"{len(existing)} existing") + logging.info(f"{len(cad_models - existing)} to convert") + remaining = list(cad_models - existing) curr_state = { - "remaining": list(set(cad_models) - set(existing)), + "remaining": remaining, "to_convert": [], "failed": [], } + logging.info(f"{len(curr_state['remaining'])} remaining") + curr_state_json_path = output_folder_path / "curr_state.json" while curr_state["remaining"]: @@ -56,10 +76,13 @@ def main(input_folder_path, output_folder_path, copies, batch_size): if res.returncode: logging.error(f"Errored out, adding to failed field in {str(curr_state_json_path)}") + logging.error(res.stderr) curr_state["failed"] = curr_state["to_convert"] logging.info(f"Finished with status code: {res.returncode}") curr_state["to_convert"] = [] + with open(curr_state_json_path, 'w') as f: + json.dump(curr_state, f) if __name__ == "__main__": main()