User/evliu/height calib

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,25 +108,34 @@ 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
 
     label_path = output_path / f"{model_name}.json"
     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")
         fastener = bpy.data.objects[model_name]
 
         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)
 

sim/scripts/sim_runner.py

@@ -2,6 +2,7 @@
 import json
 import logging
 import os
+import shutil
 import subprocess
 
 from pathlib import Path
@@ -11,26 +12,45 @@
 @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)
     assert input_folder_path.exists(), f"{input_folder_path} does not exist"
     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()