init project

app.py

@@ -83,7 +83,7 @@ def _convert_scene_output_to_glb(outdir, imgs, pts3d, mask, focals, cams2world,
     if not silent:
         print('(exporting 3D scene to', outfile, ')')
     # scene.export(file_obj=outfile)
-    print('ttttt')
+
     return outfile
 
 # @spaces.GPU(duration=180)
@@ -244,197 +244,197 @@ def slerp_multiple(vectors, t_values):
     
     return interpolated_vector
 
-@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'
+# @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]
+#     sam_mask=[]
+#     img_area = original_size[0] * original_size[1]
 
-    obj_results = yolov8(yolov8_image,device=device,retina_masks=False,imgsz=1024,conf=0.25,iou=0.95,verbose=False)
-    input_boxes1 = obj_results[0].boxes.xyxy
-    input_boxes1 = input_boxes1.cpu().numpy()
-    input_boxes1 = transform.apply_boxes(input_boxes1, original_size)
-    input_boxes = torch.from_numpy(input_boxes1).to(device)
+#     obj_results = yolov8(yolov8_image,device=device,retina_masks=False,imgsz=1024,conf=0.25,iou=0.95,verbose=False)
+#     input_boxes1 = obj_results[0].boxes.xyxy
+#     input_boxes1 = input_boxes1.cpu().numpy()
+#     input_boxes1 = transform.apply_boxes(input_boxes1, original_size)
+#     input_boxes = torch.from_numpy(input_boxes1).to(device)
     
-    # obj_results = yolov8(yolov8_image,device=device,retina_masks=False,imgsz=512,conf=0.25,iou=0.9,verbose=False)
-    # input_boxes2 = obj_results[0].boxes.xyxy
-    # input_boxes2 = input_boxes2.cpu().numpy()
-    # input_boxes2 = transform.apply_boxes(input_boxes2, original_size)
-    # input_boxes2 = torch.from_numpy(input_boxes2).to(device)
-
-    # input_boxes = torch.cat((input_boxes1, input_boxes2), dim=0)
-
-    input_image = mobilesamv2.preprocess(sam1_image)
-    image_embedding = mobilesamv2.image_encoder(input_image)['last_hidden_state']
-
-    image_embedding=torch.repeat_interleave(image_embedding, 320, dim=0)
-    prompt_embedding=mobilesamv2.prompt_encoder.get_dense_pe()
-    prompt_embedding=torch.repeat_interleave(prompt_embedding, 320, dim=0)
-    for (boxes,) in batch_iterator(320, input_boxes):
-        with torch.no_grad():
-            image_embedding=image_embedding[0:boxes.shape[0],:,:,:]
-            prompt_embedding=prompt_embedding[0:boxes.shape[0],:,:,:]
-            sparse_embeddings, dense_embeddings = mobilesamv2.prompt_encoder(
-                points=None,
-                boxes=boxes,
-                masks=None,)
-            low_res_masks, _ = mobilesamv2.mask_decoder(
-                image_embeddings=image_embedding,
-                image_pe=prompt_embedding,
-                sparse_prompt_embeddings=sparse_embeddings,
-                dense_prompt_embeddings=dense_embeddings,
-                multimask_output=False,
-                simple_type=True,
-            )
-            low_res_masks=mobilesamv2.postprocess_masks(low_res_masks, input_size, original_size)
-            sam_mask_pre = (low_res_masks > mobilesamv2.mask_threshold)
-            for mask in sam_mask_pre:
-                if mask.sum() / img_area > 0.002:
-                    sam_mask.append(mask.squeeze(1))
-    sam_mask=torch.cat(sam_mask)
-    sorted_sam_mask = sorted(sam_mask, key=(lambda x: x.sum()), reverse=True)
-    keep = mask_nms(sorted_sam_mask)
-    ret_mask = filter(sorted_sam_mask, keep)
-
-    return ret_mask
-
-@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])
-    mask_num = 0
-
-    sam1_images = images.sam1_images
-    sam1_images_size = images.sam1_images_size
-    np_images = images.np_images
-    np_images_size = images.np_images_size
+#     # obj_results = yolov8(yolov8_image,device=device,retina_masks=False,imgsz=512,conf=0.25,iou=0.9,verbose=False)
+#     # input_boxes2 = obj_results[0].boxes.xyxy
+#     # input_boxes2 = input_boxes2.cpu().numpy()
+#     # input_boxes2 = transform.apply_boxes(input_boxes2, original_size)
+#     # input_boxes2 = torch.from_numpy(input_boxes2).to(device)
+
+#     # input_boxes = torch.cat((input_boxes1, input_boxes2), dim=0)
+
+#     input_image = mobilesamv2.preprocess(sam1_image)
+#     image_embedding = mobilesamv2.image_encoder(input_image)['last_hidden_state']
+
+#     image_embedding=torch.repeat_interleave(image_embedding, 320, dim=0)
+#     prompt_embedding=mobilesamv2.prompt_encoder.get_dense_pe()
+#     prompt_embedding=torch.repeat_interleave(prompt_embedding, 320, dim=0)
+#     for (boxes,) in batch_iterator(320, input_boxes):
+#         with torch.no_grad():
+#             image_embedding=image_embedding[0:boxes.shape[0],:,:,:]
+#             prompt_embedding=prompt_embedding[0:boxes.shape[0],:,:,:]
+#             sparse_embeddings, dense_embeddings = mobilesamv2.prompt_encoder(
+#                 points=None,
+#                 boxes=boxes,
+#                 masks=None,)
+#             low_res_masks, _ = mobilesamv2.mask_decoder(
+#                 image_embeddings=image_embedding,
+#                 image_pe=prompt_embedding,
+#                 sparse_prompt_embeddings=sparse_embeddings,
+#                 dense_prompt_embeddings=dense_embeddings,
+#                 multimask_output=False,
+#                 simple_type=True,
+#             )
+#             low_res_masks=mobilesamv2.postprocess_masks(low_res_masks, input_size, original_size)
+#             sam_mask_pre = (low_res_masks > mobilesamv2.mask_threshold)
+#             for mask in sam_mask_pre:
+#                 if mask.sum() / img_area > 0.002:
+#                     sam_mask.append(mask.squeeze(1))
+#     sam_mask=torch.cat(sam_mask)
+#     sorted_sam_mask = sorted(sam_mask, key=(lambda x: x.sum()), reverse=True)
+#     keep = mask_nms(sorted_sam_mask)
+#     ret_mask = filter(sorted_sam_mask, keep)
+
+#     return ret_mask
+
+# @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])
+#     mask_num = 0
+
+#     sam1_images = images.sam1_images
+#     sam1_images_size = images.sam1_images_size
+#     np_images = images.np_images
+#     np_images_size = images.np_images_size
     
-    sam1_masks = get_mask_from_img_sam1(pe3r.mobilesamv2, pe3r.yolov8, sam1_images[0], np_images[0], np_images_size[0], sam1_images_size[0], images.sam1_transform)
-    for mask in sam1_masks:
-        _, _, _ = pe3r.sam2.add_new_mask(
-            inference_state=inference_state,
-            frame_idx=0,
-            obj_id=mask_num,
-            mask=mask,
-        )
-        mask_num += 1
-
-    video_segments = {}  # video_segments contains the per-frame segmentation results
-    for out_frame_idx, out_obj_ids, out_mask_logits in pe3r.sam2.propagate_in_video(inference_state):
-        sam2_masks = (out_mask_logits > 0.0).squeeze(1)
-
-        video_segments[out_frame_idx] = {
-            out_obj_id: sam2_masks[i].cpu().numpy()
-            for i, out_obj_id in enumerate(out_obj_ids)
-        }
-
-        if out_frame_idx == 0:
-            continue
-
-        sam1_masks = get_mask_from_img_sam1(pe3r.mobilesamv2, pe3r.yolov8, sam1_images[out_frame_idx], np_images[out_frame_idx], np_images_size[out_frame_idx], sam1_images_size[out_frame_idx], images.sam1_transform)
-
-        for sam1_mask in sam1_masks:
-            flg = 1
-            for sam2_mask in sam2_masks:
-                # print(sam1_mask.shape, sam2_mask.shape)
-                area1 = sam1_mask.sum()
-                area2 = sam2_mask.sum()
-                intersection = (sam1_mask & sam2_mask).sum()
-                if min(intersection / area1, intersection / area2) > 0.25:
-                    flg = 0
-                    break
-            if flg:
-                video_segments[out_frame_idx][mask_num] = sam1_mask.cpu().numpy()
-                mask_num += 1
-
-    multi_view_clip_feats = torch.zeros((mask_num+1, 1024))
-    multi_view_clip_feats_map = {}
-    multi_view_clip_area_map = {}
-    for now_frame in range(0, len(video_segments), 1):
-        image = np_images[now_frame]
-
-        seg_img_list = []
-        out_obj_id_list = []
-        out_obj_mask_list = []
-        out_obj_area_list = []
-        # NOTE: background: -1
-        rev_seg_map = -np.ones(image.shape[:2], dtype=np.int64)
-        sorted_dict_items = sorted(video_segments[now_frame].items(), key=lambda x: np.count_nonzero(x[1]), reverse=False)
-        for out_obj_id, mask in sorted_dict_items:
-            if mask.sum() == 0:
-                continue
-            rev_seg_map[mask] = out_obj_id
-        rev_cog_seg_maps.append(rev_seg_map)
-
-        seg_map = -np.ones(image.shape[:2], dtype=np.int64)
-        sorted_dict_items = sorted(video_segments[now_frame].items(), key=lambda x: np.count_nonzero(x[1]), reverse=True)
-        for out_obj_id, mask in sorted_dict_items:
-            if mask.sum() == 0:
-                continue
-            box = np.int32(box_xyxy_to_xywh(mask_to_box(mask)))
+#     sam1_masks = get_mask_from_img_sam1(pe3r.mobilesamv2, pe3r.yolov8, sam1_images[0], np_images[0], np_images_size[0], sam1_images_size[0], images.sam1_transform)
+#     for mask in sam1_masks:
+#         _, _, _ = pe3r.sam2.add_new_mask(
+#             inference_state=inference_state,
+#             frame_idx=0,
+#             obj_id=mask_num,
+#             mask=mask,
+#         )
+#         mask_num += 1
+
+#     video_segments = {}  # video_segments contains the per-frame segmentation results
+#     for out_frame_idx, out_obj_ids, out_mask_logits in pe3r.sam2.propagate_in_video(inference_state):
+#         sam2_masks = (out_mask_logits > 0.0).squeeze(1)
+
+#         video_segments[out_frame_idx] = {
+#             out_obj_id: sam2_masks[i].cpu().numpy()
+#             for i, out_obj_id in enumerate(out_obj_ids)
+#         }
+
+#         if out_frame_idx == 0:
+#             continue
+
+#         sam1_masks = get_mask_from_img_sam1(pe3r.mobilesamv2, pe3r.yolov8, sam1_images[out_frame_idx], np_images[out_frame_idx], np_images_size[out_frame_idx], sam1_images_size[out_frame_idx], images.sam1_transform)
+
+#         for sam1_mask in sam1_masks:
+#             flg = 1
+#             for sam2_mask in sam2_masks:
+#                 # print(sam1_mask.shape, sam2_mask.shape)
+#                 area1 = sam1_mask.sum()
+#                 area2 = sam2_mask.sum()
+#                 intersection = (sam1_mask & sam2_mask).sum()
+#                 if min(intersection / area1, intersection / area2) > 0.25:
+#                     flg = 0
+#                     break
+#             if flg:
+#                 video_segments[out_frame_idx][mask_num] = sam1_mask.cpu().numpy()
+#                 mask_num += 1
+
+#     multi_view_clip_feats = torch.zeros((mask_num+1, 1024))
+#     multi_view_clip_feats_map = {}
+#     multi_view_clip_area_map = {}
+#     for now_frame in range(0, len(video_segments), 1):
+#         image = np_images[now_frame]
+
+#         seg_img_list = []
+#         out_obj_id_list = []
+#         out_obj_mask_list = []
+#         out_obj_area_list = []
+#         # NOTE: background: -1
+#         rev_seg_map = -np.ones(image.shape[:2], dtype=np.int64)
+#         sorted_dict_items = sorted(video_segments[now_frame].items(), key=lambda x: np.count_nonzero(x[1]), reverse=False)
+#         for out_obj_id, mask in sorted_dict_items:
+#             if mask.sum() == 0:
+#                 continue
+#             rev_seg_map[mask] = out_obj_id
+#         rev_cog_seg_maps.append(rev_seg_map)
+
+#         seg_map = -np.ones(image.shape[:2], dtype=np.int64)
+#         sorted_dict_items = sorted(video_segments[now_frame].items(), key=lambda x: np.count_nonzero(x[1]), reverse=True)
+#         for out_obj_id, mask in sorted_dict_items:
+#             if mask.sum() == 0:
+#                 continue
+#             box = np.int32(box_xyxy_to_xywh(mask_to_box(mask)))
             
-            if box[2] == 0 and box[3] == 0:
-                continue
-            # print(box)
-            seg_img = get_seg_img(mask, box, image)
-            pad_seg_img = cv2.resize(pad_img(seg_img), (256,256))
-            seg_img_list.append(pad_seg_img)
-            seg_map[mask] = out_obj_id
-            out_obj_id_list.append(out_obj_id)
-            out_obj_area_list.append(np.count_nonzero(mask))
-            out_obj_mask_list.append(mask)
-
-        if len(seg_img_list) == 0:
-            cog_seg_maps.append(seg_map)
-            continue
-
-        seg_imgs = np.stack(seg_img_list, axis=0) # b,H,W,3
-        seg_imgs = torch.from_numpy(seg_imgs).permute(0,3,1,2) # / 255.0
+#             if box[2] == 0 and box[3] == 0:
+#                 continue
+#             # print(box)
+#             seg_img = get_seg_img(mask, box, image)
+#             pad_seg_img = cv2.resize(pad_img(seg_img), (256,256))
+#             seg_img_list.append(pad_seg_img)
+#             seg_map[mask] = out_obj_id
+#             out_obj_id_list.append(out_obj_id)
+#             out_obj_area_list.append(np.count_nonzero(mask))
+#             out_obj_mask_list.append(mask)
+
+#         if len(seg_img_list) == 0:
+#             cog_seg_maps.append(seg_map)
+#             continue
+
+#         seg_imgs = np.stack(seg_img_list, axis=0) # b,H,W,3
+#         seg_imgs = torch.from_numpy(seg_imgs).permute(0,3,1,2) # / 255.0
         
-        inputs = pe3r.siglip_processor(images=seg_imgs, return_tensors="pt")
-        inputs = {key: value.to(device) for key, value in inputs.items()}
+#         inputs = pe3r.siglip_processor(images=seg_imgs, return_tensors="pt")
+#         inputs = {key: value.to(device) for key, value in inputs.items()}
         
-        image_features = pe3r.siglip.get_image_features(**inputs)
-        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
-        image_features = image_features.detach().cpu()
-
-        for i in range(len(out_obj_mask_list)):
-            for j in range(i + 1, len(out_obj_mask_list)):
-                mask1 = out_obj_mask_list[i]
-                mask2 = out_obj_mask_list[j]
-                intersection = np.logical_and(mask1, mask2).sum()
-                area1 = out_obj_area_list[i]
-                area2 = out_obj_area_list[j]
-                if min(intersection / area1, intersection / area2) > 0.025:
-                    conf1 = area1 / (area1 + area2)
-                    # conf2 = area2 / (area1 + area2)
-                    image_features[j] = slerp(image_features[j], image_features[i], conf1)
-
-        for i, clip_feat in enumerate(image_features):
-            id = out_obj_id_list[i]
-            if id in multi_view_clip_feats_map.keys():
-                multi_view_clip_feats_map[id].append(clip_feat)
-                multi_view_clip_area_map[id].append(out_obj_area_list[i])
-            else:
-                multi_view_clip_feats_map[id] = [clip_feat]
-                multi_view_clip_area_map[id] = [out_obj_area_list[i]]
-
-        cog_seg_maps.append(seg_map)
-        del image_features
+#         image_features = pe3r.siglip.get_image_features(**inputs)
+#         image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+#         image_features = image_features.detach().cpu()
+
+#         for i in range(len(out_obj_mask_list)):
+#             for j in range(i + 1, len(out_obj_mask_list)):
+#                 mask1 = out_obj_mask_list[i]
+#                 mask2 = out_obj_mask_list[j]
+#                 intersection = np.logical_and(mask1, mask2).sum()
+#                 area1 = out_obj_area_list[i]
+#                 area2 = out_obj_area_list[j]
+#                 if min(intersection / area1, intersection / area2) > 0.025:
+#                     conf1 = area1 / (area1 + area2)
+#                     # conf2 = area2 / (area1 + area2)
+#                     image_features[j] = slerp(image_features[j], image_features[i], conf1)
+
+#         for i, clip_feat in enumerate(image_features):
+#             id = out_obj_id_list[i]
+#             if id in multi_view_clip_feats_map.keys():
+#                 multi_view_clip_feats_map[id].append(clip_feat)
+#                 multi_view_clip_area_map[id].append(out_obj_area_list[i])
+#             else:
+#                 multi_view_clip_feats_map[id] = [clip_feat]
+#                 multi_view_clip_area_map[id] = [out_obj_area_list[i]]
+
+#         cog_seg_maps.append(seg_map)
+#         del image_features
         
-    for i in range(mask_num):
-        if i in multi_view_clip_feats_map.keys():
-            clip_feats = multi_view_clip_feats_map[i]
-            mask_area = multi_view_clip_area_map[i]
-            multi_view_clip_feats[i] = slerp_multiple(torch.stack(clip_feats), np.stack(mask_area))
-        else:
-            multi_view_clip_feats[i] = torch.zeros((1024))
-    multi_view_clip_feats[mask_num] = torch.zeros((1024))
+#     for i in range(mask_num):
+#         if i in multi_view_clip_feats_map.keys():
+#             clip_feats = multi_view_clip_feats_map[i]
+#             mask_area = multi_view_clip_area_map[i]
+#             multi_view_clip_feats[i] = slerp_multiple(torch.stack(clip_feats), np.stack(mask_area))
+#         else:
+#             multi_view_clip_feats[i] = torch.zeros((1024))
+#     multi_view_clip_feats[mask_num] = torch.zeros((1024))
 
-    return cog_seg_maps, rev_cog_seg_maps, multi_view_clip_feats
+#     return cog_seg_maps, rev_cog_seg_maps, multi_view_clip_feats
 
 @spaces.GPU(duration=180)
 def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
@@ -452,16 +452,16 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
     images = Images(filelist=filelist, device=device)
     
     # try:
-    cog_seg_maps, rev_cog_seg_maps, cog_feats = get_cog_feats(images)
-    imgs = load_images(images, rev_cog_seg_maps, size=512, verbose=not silent)
+    # cog_seg_maps, rev_cog_seg_maps, cog_feats = get_cog_feats(images)
+    # imgs = load_images(images, rev_cog_seg_maps, size=512, verbose=not silent)
     # except Exception as e:
-        # rev_cog_seg_maps = []
-        # for tmp_img in images.np_images:
-        #     rev_seg_map = -np.ones(tmp_img.shape[:2], dtype=np.int64)
-        #     rev_cog_seg_maps.append(rev_seg_map)
-        # cog_seg_maps = rev_cog_seg_maps
-        # cog_feats = torch.zeros((1, 1024))
-        # imgs = load_images(images, rev_cog_seg_maps, size=512, verbose=not silent)
+    rev_cog_seg_maps = []
+    for tmp_img in images.np_images:
+        rev_seg_map = -np.ones(tmp_img.shape[:2], dtype=np.int64)
+        rev_cog_seg_maps.append(rev_seg_map)
+    cog_seg_maps = rev_cog_seg_maps
+    cog_feats = torch.zeros((1, 1024))
+    imgs = load_images(images, rev_cog_seg_maps, size=512, verbose=not silent)
 
     if len(imgs) == 1:
         imgs = [imgs[0], copy.deepcopy(imgs[0])]
@@ -499,11 +499,8 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
         scene.ori_imgs = ori_imgs
         print(e)
 
-    print('a')
-
     outfile = get_3D_model_from_scene(outdir, scene, min_conf_thr, as_pointcloud, mask_sky,
                                       clean_depth, transparent_cams, cam_size)
-    print('b')
     # also return rgb, depth and confidence imgs
     # depth is normalized with the max value for all images
     # we apply the jet colormap on the confidence maps