add main()

app.py

@@ -60,27 +60,27 @@ sam_segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).to(
     DEVICE, DTYPE
 )
 
-############## 3D-Fixer model
-model_dir = 'HorizonRobotics/3D-Fixer'
-pipeline = ThreeDFixerPipeline.from_pretrained(
-    model_dir, compile=False
-)
-pipeline.to(device=DEVICE)
-############## 3D-Fixer model
-
-rot = np.array([
-    [-1.0,  0.0,  0.0, 0.0],
-    [ 0.0,  0.0,  1.0, 0.0],
-    [ 0.0,  1.0,  0.0, 0.0],
-    [ 0.0,  0.0,  0.0, 1.0],
-], dtype=np.float32)
-
-c2w = torch.tensor([
-    [1.0, 0.0, 0.0, 0.0],
-    [0.0, 0.0, -1.0, 0.0],
-    [0.0, 1.0, 0.0, 0.0],
-    [0.0, 0.0, 0.0, 1.0],
-], dtype=torch.float32, device=DEVICE)
+# ############## 3D-Fixer model
+# model_dir = 'HorizonRobotics/3D-Fixer'
+# pipeline = ThreeDFixerPipeline.from_pretrained(
+#     model_dir, compile=False
+# )
+# pipeline.to(device=DEVICE)
+# ############## 3D-Fixer model
+
+# rot = np.array([
+#     [-1.0,  0.0,  0.0, 0.0],
+#     [ 0.0,  0.0,  1.0, 0.0],
+#     [ 0.0,  1.0,  0.0, 0.0],
+#     [ 0.0,  0.0,  0.0, 1.0],
+# ], dtype=np.float32)
+
+# c2w = torch.tensor([
+#     [1.0, 0.0, 0.0, 0.0],
+#     [0.0, 0.0, -1.0, 0.0],
+#     [0.0, 1.0, 0.0, 0.0],
+#     [0.0, 0.0, 0.0, 1.0],
+# ], dtype=torch.float32, device=DEVICE)
 
 save_projected_colored_pcd = lambda pts, pts_color, fpath: trimesh.PointCloud(pts.reshape(-1, 3), pts_color.reshape(-1, 3)).export(fpath)
 
@@ -881,5 +881,42 @@ with gr.Blocks() as demo:
         outputs=[download_single_glb],
     )
 
+if __name__ == '__main__':
+
+    work_space = None
+    dpt_pack = None
+    generated_object_map = {}
+
+    DEVICE = "cuda"
 
-demo.launch(allowed_paths=[TMP_DIR, EXAMPLE_DIR])
+    # Prepare models
+    ## Grounding SAM
+    segmenter_id = "facebook/sam-vit-base"
+    sam_processor = AutoProcessor.from_pretrained(segmenter_id)
+    sam_segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).to(
+        DEVICE, DTYPE
+    )
+
+    ############## 3D-Fixer model
+    model_dir = 'HorizonRobotics/3D-Fixer'
+    pipeline = ThreeDFixerPipeline.from_pretrained(
+        model_dir, compile=False
+    )
+    pipeline.to(device=DEVICE)
+    ############## 3D-Fixer model
+
+    rot = np.array([
+        [-1.0,  0.0,  0.0, 0.0],
+        [ 0.0,  0.0,  1.0, 0.0],
+        [ 0.0,  1.0,  0.0, 0.0],
+        [ 0.0,  0.0,  0.0, 1.0],
+    ], dtype=np.float32)
+
+    c2w = torch.tensor([
+        [1.0, 0.0, 0.0, 0.0],
+        [0.0, 0.0, -1.0, 0.0],
+        [0.0, 1.0, 0.0, 0.0],
+        [0.0, 0.0, 0.0, 1.0],
+    ], dtype=torch.float32, device=DEVICE)
+
+    demo.launch(allowed_paths=[TMP_DIR, EXAMPLE_DIR])