init project

app.py

@@ -39,8 +39,7 @@ import torchvision.transforms as tvf
 
 
 silent = False
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-pe3r = Models(device)
+pe3r = Models('cpu')
 
 
 def _convert_scene_output_to_glb(outdir, imgs, pts3d, mask, focals, cams2world, cam_size=0.05,
@@ -86,7 +85,7 @@ def _convert_scene_output_to_glb(outdir, imgs, pts3d, mask, focals, cams2world,
     scene.export(file_obj=outfile)
     return outfile
 
-@spaces.GPU(duration=180)
+# @spaces.GPU(duration=180)
 def get_3D_model_from_scene(outdir, scene, min_conf_thr=3, as_pointcloud=False, mask_sky=False,
                             clean_depth=False, transparent_cams=False, cam_size=0.05):
     """
@@ -246,6 +245,8 @@ def slerp_multiple(vectors, t_values):
 
 @torch.no_grad
 def get_mask_from_img_sam1(mobilesamv2, yolov8, sam1_image, yolov8_image, original_size, input_size, transform):
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
     sam_mask=[]
     img_area = original_size[0] * original_size[1]
 
@@ -299,6 +300,7 @@ def get_mask_from_img_sam1(mobilesamv2, yolov8, sam1_image, yolov8_image, origin
 
 @torch.no_grad
 def get_cog_feats(images):
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
     cog_seg_maps = []
     rev_cog_seg_maps = []
     inference_state = pe3r.sam2.init_state(images=images.sam2_images, video_height=images.sam2_video_size[0], video_width=images.sam2_video_size[1])
@@ -443,6 +445,8 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
     """
     if len(filelist) < 2:
         raise gradio.Error("Please input at least 2 images.")
+    
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     images = Images(filelist=filelist, device=device)
     
@@ -523,9 +527,11 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
 def get_3D_object_from_scene(outdir, text, threshold, scene, min_conf_thr, as_pointcloud, 
                  mask_sky, clean_depth, transparent_cams, cam_size):
     
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
     texts = [text]
     inputs = pe3r.siglip_tokenizer(text=texts, padding="max_length", return_tensors="pt")
-    inputs = {key: value.to('cuda') for key, value in inputs.items()}
+    inputs = {key: value.to(device) for key, value in inputs.items()}
     with torch.no_grad():
         text_feats =pe3r.siglip.get_text_features(**inputs)
         text_feats = text_feats / text_feats.norm(dim=-1, keepdim=True)
@@ -559,7 +565,7 @@ def set_scenegraph_options(inputfiles, winsize, refid, scenegraph_type):
 
 with tempfile.TemporaryDirectory(suffix='pe3r_gradio_demo') as tmpdirname:
     recon_fun = functools.partial(get_reconstructed_scene, tmpdirname)
-    model_from_scene_fun = functools.partial(get_3D_model_from_scene, tmpdirname)
+    # model_from_scene_fun = functools.partial(get_3D_model_from_scene, tmpdirname)
     get_3D_object_from_scene_fun = functools.partial(get_3D_object_from_scene, tmpdirname)
 
     with gradio.Blocks(css=""".gradio-container {margin: 0 !important; min-width: 100%};""", title="PE3R Demo") as demo:
@@ -594,11 +600,11 @@ with tempfile.TemporaryDirectory(suffix='pe3r_gradio_demo') as tmpdirname:
                 # adjust the camera size in the output pointcloud
                 cam_size = gradio.Slider(label="cam_size", value=0.05, minimum=0.001, maximum=0.1, step=0.001, visible=False)
             with gradio.Row():
-                as_pointcloud = gradio.Checkbox(value=True, label="As pointcloud")
+                as_pointcloud = gradio.Checkbox(value=True, label="As pointcloud", visible=False)
                 # two post process implemented
                 mask_sky = gradio.Checkbox(value=False, label="Mask sky", visible=False)
                 clean_depth = gradio.Checkbox(value=True, label="Clean-up depthmaps", visible=False)
-                transparent_cams = gradio.Checkbox(value=True, label="Transparent cameras")
+                transparent_cams = gradio.Checkbox(value=True, label="Transparent cameras", visible=False)
 
             with gradio.Row():
                 text_input = gradio.Textbox(label="Query Text")
@@ -622,30 +628,30 @@ with tempfile.TemporaryDirectory(suffix='pe3r_gradio_demo') as tmpdirname:
                                     mask_sky, clean_depth, transparent_cams, cam_size,
                                     scenegraph_type, winsize, refid],
                             outputs=[scene, outmodel]) # , outgallery
-            min_conf_thr.release(fn=model_from_scene_fun,
-                                    inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
-                                            clean_depth, transparent_cams, cam_size],
-                                    outputs=outmodel)
-            cam_size.change(fn=model_from_scene_fun,
-                            inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
-                                    clean_depth, transparent_cams, cam_size],
-                            outputs=outmodel)
-            as_pointcloud.change(fn=model_from_scene_fun,
-                                    inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
-                                            clean_depth, transparent_cams, cam_size],
-                                    outputs=outmodel)
-            mask_sky.change(fn=model_from_scene_fun,
-                            inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
-                                    clean_depth, transparent_cams, cam_size],
-                            outputs=outmodel)
-            clean_depth.change(fn=model_from_scene_fun,
-                                inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
-                                        clean_depth, transparent_cams, cam_size],
-                                outputs=outmodel)
-            transparent_cams.change(model_from_scene_fun,
-                                    inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
-                                            clean_depth, transparent_cams, cam_size],
-                                    outputs=outmodel)
+            # min_conf_thr.release(fn=model_from_scene_fun,
+            #                         inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+            #                                 clean_depth, transparent_cams, cam_size],
+            #                         outputs=outmodel)
+            # cam_size.change(fn=model_from_scene_fun,
+            #                 inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+            #                         clean_depth, transparent_cams, cam_size],
+            #                 outputs=outmodel)
+            # as_pointcloud.change(fn=model_from_scene_fun,
+            #                         inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+            #                                 clean_depth, transparent_cams, cam_size],
+            #                         outputs=outmodel)
+            # mask_sky.change(fn=model_from_scene_fun,
+            #                 inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+            #                         clean_depth, transparent_cams, cam_size],
+            #                 outputs=outmodel)
+            # clean_depth.change(fn=model_from_scene_fun,
+            #                     inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+            #                             clean_depth, transparent_cams, cam_size],
+            #                     outputs=outmodel)
+            # transparent_cams.change(model_from_scene_fun,
+            #                         inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+            #                                 clean_depth, transparent_cams, cam_size],
+            #                         outputs=outmodel)
             find_btn.click(fn=get_3D_object_from_scene_fun,
                                 inputs=[text_input, threshold, scene, min_conf_thr, as_pointcloud, mask_sky,
                                         clean_depth, transparent_cams, cam_size],