init project

app.py

@@ -43,6 +43,7 @@ from modules.mobilesamv2 import sam_model_registry
 from sam2.sam2_video_predictor import SAM2VideoPredictor
 from modules.mast3r.model import AsymmetricMASt3R
 
+from torch.nn.functional import cosine_similarity
 
 silent = False
 
@@ -448,6 +449,44 @@ def get_cog_feats(images, sam2, siglip, siglip_processor, yolov8, mobilesamv2):
     return cog_seg_maps, rev_cog_seg_maps, multi_view_clip_feats
 
 
+class Scene_cpu:
+    def __init__(self, fix_imgs, cogs, focals, cams2world, pts3d, min_conf_thr, msk):
+        self.fix_imgs = fix_imgs
+        self.cogs = cogs
+        self.focals = focals
+        self.cams2world = cams2world
+        self.pts3d = pts3d
+        self.min_conf_thr = min_conf_thr
+        self.msk = msk
+
+    def render_image(self, text_feats, threshold=0.85):
+        self.rendered_imgs = []
+        # Collect all cosine similarities to compute min-max normalization
+        all_similarities = []
+        for each_cog in self.cogs:
+            similarity_map = cosine_similarity(each_cog, text_feats.unsqueeze(1), dim=-1)
+            all_similarities.append(similarity_map.squeeze().numpy())
+        # Flatten and normalize all similarities
+        total_similarities = np.concatenate(all_similarities)
+        min_sim, max_sim = total_similarities.min(), total_similarities.max()
+        normalized_similarities = [(sim - min_sim) / (max_sim - min_sim) for sim in all_similarities]
+        # Process each image with normalized similarities
+        for i, (each_cog, heatmap) in enumerate(zip(self.cogs, normalized_similarities)):
+            mask = heatmap > threshold
+            # Scale heatmap for visualization
+            heatmap = np.uint8(255 * heatmap)
+            heatmap_color = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+            # Prepare image
+            image = self.fix_imgs[i]
+            image = image * 255.0
+            image = np.clip(image, 0, 255).astype(np.uint8)
+            # Apply mask and overlay heatmap with red RGB for masked areas
+            mask_indices = np.where(mask)  # Get indices where mask is True
+            heatmap_color[mask_indices[0], mask_indices[1]] = [0, 0, 255]  # Red color for masked regions
+            superimposed_img = np.where(np.expand_dims(mask, axis=-1), heatmap_color, image) / 255.0
+            self.rendered_imgs.append(superimposed_img)
+
+
 @spaces.GPU(duration=120)
 def get_reconstructed_scene(outdir, filelist, schedule='linear', niter=300, min_conf_thr=3.0,
                             as_pointcloud=True, mask_sky=False, clean_depth=True, transparent_cams=True, cam_size=0.05,
@@ -540,37 +579,46 @@ def get_reconstructed_scene(outdir, filelist, schedule='linear', niter=300, min_
     
     torch.cuda.empty_cache()
 
+    fix_imgs = []
+    for img in scene.fix_imgs:
+        fix_imgs.append(img)
+    cogs = []
+    for cog in scene.cogs:
+        cog_cpu = cog.cpu()
+        cogs.append(cog_cpu)
+    focals = scene.get_focals().cpu()
+    cams2world = scene.get_im_poses().cpu()
+    pts3d = to_numpy(scene.get_pts3d())
+    min_conf_thr = float(scene.conf_trf(torch.tensor(3.0)))
+    msk = to_numpy(scene.get_masks())
+    scene_cpu = Scene_cpu(fix_imgs, cogs, focals, cams2world, pts3d, min_conf_thr, msk)
 
-    return outfile 
-
-# def get_3D_object_from_scene(outdir, text, threshold, scene, min_conf_thr=3.0, as_pointcloud=True, 
-#                  mask_sky=False, clean_depth=True, transparent_cams=True, cam_size=0.05):
-    
-#     device = 'cpu'
-#     siglip_tokenizer = AutoTokenizer.from_pretrained("google/siglip-large-patch16-256")
-#     siglip = AutoModel.from_pretrained("google/siglip-large-patch16-256", device_map=device)
-
-#     texts = [text]
-#     inputs = siglip_tokenizer(text=texts, padding="max_length", return_tensors="pt")
-#     inputs = {key: value.to(device) for key, value in inputs.items()}
-#     with torch.no_grad():
-#         text_feats =siglip.get_text_features(**inputs)
-#         text_feats = text_feats / text_feats.norm(dim=-1, keepdim=True)
-#     scene.render_image(text_feats, threshold)
-#     scene.ori_imgs = scene.rendered_imgs
-
+    return scene_cpu, outfile 
 
-#     rgbimg = scene.ori_imgs
-#     focals = scene.get_focals().cpu()
-#     cams2world = scene.get_im_poses().cpu()
-#     # 3D pointcloud from depthmap, poses and intrinsics
-#     pts3d = to_numpy(scene.get_pts3d())
-#     scene.min_conf_thr = float(scene.conf_trf(torch.tensor(min_conf_thr)))
-#     msk = to_numpy(scene.get_masks())
-#     return _convert_scene_output_to_glb(outdir, rgbimg, pts3d, msk, focals, cams2world, as_pointcloud=as_pointcloud,
-#                                         transparent_cams=transparent_cams, cam_size=cam_size)
 
+def get_3D_object_from_scene(outdir, text, threshold, scene, min_conf_thr=3.0, as_pointcloud=True, 
+                 mask_sky=False, clean_depth=True, transparent_cams=True, cam_size=0.05):
+    
+    device = 'cpu'
+    siglip_tokenizer = AutoTokenizer.from_pretrained("google/siglip-large-patch16-256")
+    siglip = AutoModel.from_pretrained("google/siglip-large-patch16-256", device_map=device)
 
+    texts = [text]
+    inputs = siglip_tokenizer(text=texts, padding="max_length", return_tensors="pt")
+    inputs = {key: value.to(device) for key, value in inputs.items()}
+    with torch.no_grad():
+        text_feats =siglip.get_text_features(**inputs)
+        text_feats = text_feats / text_feats.norm(dim=-1, keepdim=True)
+    scene.render_image(text_feats, threshold)
+    scene.ori_imgs = scene.rendered_imgs
+    rgbimg = scene.ori_imgs
+    focals = scene.focals
+    cams2world = scene.cams2world
+    # 3D pointcloud from depthmap, poses and intrinsics
+    pts3d = scene.pts3d
+    msk = scene.msk
+    return _convert_scene_output_to_glb(outdir, rgbimg, pts3d, msk, focals, cams2world, as_pointcloud=as_pointcloud,
+                                        transparent_cams=transparent_cams, cam_size=cam_size)
 
 
 
@@ -579,11 +627,11 @@ tmpdirname = tempfile.mkdtemp(suffix='pe3r_gradio_demo')
 
 recon_fun = functools.partial(get_reconstructed_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)
+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:
     # scene state is save so that you can change conf_thr, cam_size... without rerunning the inference
-    # scene = gradio.State(None)
+    scene = gradio.State(None)
 
     gradio.HTML('<h2 style="text-align: center;">PE3R Demo</h2>')
     with gradio.Column():
@@ -602,9 +650,9 @@ with gradio.Blocks(css=""".gradio-container {margin: 0 !important; min-width: 10
 
         run_btn.click(fn=recon_fun,
                         inputs=[inputfiles],
-                        outputs=[outmodel]) # , outgallery, scene, 
+                        outputs=[scene, outmodel]) # , outgallery, , 
         
-        # find_btn.click(fn=get_3D_object_from_scene_fun,
-        #                     inputs=[text_input, threshold, scene],
-        #                 outputs=outmodel)
+        find_btn.click(fn=get_3D_object_from_scene_fun,
+                            inputs=[text_input, threshold, scene],
+                        outputs=outmodel)
 demo.launch(show_error=True, share=None, server_name=None, server_port=None)