maintainance fixes

.gitmodules

@@ -3,7 +3,7 @@
 	url = https://github.com/minghanqin/segment-anything-langsplat
 [submodule "submodules/simple-knn"]
 	path = submodules/simple-knn
-	url = https://gitlab.inria.fr/bkerbl/simple-knn.git
+	url = https://github.com/camenduru/simple-knn.git
 [submodule "submodules/langsplat-rasterization"]
 	path = submodules/langsplat-rasterization
 	url = https://github.com/minghanqin/langsplat-rasterization

autoencoder/test.py

@@ -23,7 +23,7 @@
                     default=[16, 32, 64, 128, 256, 256, 512],
                     )
     args = parser.parse_args()
-    
+
     dataset_name = args.dataset_name
     encoder_hidden_dims = args.encoder_dims
     decoder_hidden_dims = args.decoder_dims
@@ -33,7 +33,7 @@
     data_dir = f"{dataset_path}/language_features"
     output_dir = f"{dataset_path}/language_features_dim3"
     os.makedirs(output_dir, exist_ok=True)
-    
+
     # copy the segmentation map
     for filename in os.listdir(data_dir):
         if filename.endswith("_s.npy"):
@@ -42,15 +42,15 @@
             shutil.copy(source_path, target_path)
 
 
-    checkpoint = torch.load(ckpt_path)
+    checkpoint = torch.load(ckpt_path, weights_only=False)
     train_dataset = Autoencoder_dataset(data_dir)
 
     test_loader = DataLoader(
-        dataset=train_dataset, 
+        dataset=train_dataset,
         batch_size=256,
-        shuffle=False, 
-        num_workers=16, 
-        drop_last=False   
+        shuffle=False,
+        num_workers=16,
+        drop_last=False
     )
 
 
@@ -62,15 +62,15 @@
     for idx, feature in tqdm(enumerate(test_loader)):
         data = feature.to("cuda:0")
         with torch.no_grad():
-            outputs = model.encode(data).to("cpu").numpy()  
+            outputs = model.encode(data).to("cpu").numpy()
         if idx == 0:
             features = outputs
         else:
             features = np.concatenate([features, outputs], axis=0)
 
     os.makedirs(output_dir, exist_ok=True)
     start = 0
-    
+
     for k,v in train_dataset.data_dic.items():
         path = os.path.join(output_dir, k)
         np.save(path, features[start:start+v])

environment.yml

@@ -5,13 +5,12 @@ channels:
   - defaults
 dependencies:
   - open-clip-torch
-  - cudatoolkit=11.6
-  - plyfile=0.8.1
-  - python=3.7.13
-  - pip=22.3.1
-  - pytorch=1.12.1
-  - torchaudio=0.12.1
-  - torchvision=0.13.1
+  - plyfile
+  - python
+  - pip
+  - pytorch
+  - torchaudio
+  - torchvision
   - tqdm
   - opencv
   - tensorboard
@@ -23,3 +22,4 @@ dependencies:
     - submodules/segment-anything-langsplat
     - submodules/langsplat-rasterization
     - submodules/simple-knn
+    - nvidia-cuda-runtime-cu12

eval/evaluate_iou_loc.py

@@ -63,9 +63,9 @@ def eval_gt_lerfdata(json_folder: Union[str, Path] = None, ouput_path: Path = No
         img_ann = defaultdict(dict)
         with open(js_path, 'r') as f:
             gt_data = json.load(f)
-        
+
         h, w = gt_data['info']['height'], gt_data['info']['width']
-        idx = int(gt_data['info']['name'].split('_')[-1].split('.jpg')[0]) - 1 
+        idx = int(gt_data['info']['name'].split('_')[-1].split('.jpg')[0]) - 1
         for prompt_data in gt_data["objects"]:
             label = prompt_data['category']
             box = np.asarray(prompt_data['bbox']).reshape(-1)           # x1y1x2y2
@@ -77,7 +77,7 @@ def eval_gt_lerfdata(json_folder: Union[str, Path] = None, ouput_path: Path = No
             else:
                 img_ann[label]['bboxes'] = box
             img_ann[label]['mask'] = mask
-            
+
             # # save for visulsization
             save_path = ouput_path / 'gt' / gt_data['info']['name'].split('.jpg')[0] / f'{label}.jpg'
             save_path.parent.mkdir(exist_ok=True, parents=True)
@@ -87,12 +87,12 @@ def eval_gt_lerfdata(json_folder: Union[str, Path] = None, ouput_path: Path = No
     return gt_ann, (h, w), img_paths
 
 
-def activate_stream(sem_map, 
-                    image, 
-                    clip_model, 
+def activate_stream(sem_map,
+                    image,
+                    clip_model,
                     image_name: Path = None,
-                    img_ann: Dict = None, 
-                    thresh : float = 0.5, 
+                    img_ann: Dict = None,
+                    thresh : float = 0.5,
                     colormap_options = None):
     valid_map = clip_model.get_max_across(sem_map)                 # 3xkx832x1264
     n_head, n_prompt, h, w = valid_map.shape
@@ -110,12 +110,12 @@ def activate_stream(sem_map,
             avg_filtered = cv2.filter2D(np_relev, -1, kernel)
             avg_filtered = torch.from_numpy(avg_filtered).to(valid_map.device)
             valid_map[i][k] = 0.5 * (avg_filtered + valid_map[i][k])
-            
+
             output_path_relev = image_name / 'heatmap' / f'{clip_model.positives[k]}_{i}'
             output_path_relev.parent.mkdir(exist_ok=True, parents=True)
             colormap_saving(valid_map[i][k].unsqueeze(-1), colormap_options,
                             output_path_relev)
-            
+
             # NOTE 与lerf一致，激活值低于0.5的认为是背景
             p_i = torch.clip(valid_map[i][k] - 0.5, 0, 1).unsqueeze(-1)
             valid_composited = colormaps.apply_colormap(p_i / (p_i.max() + 1e-6), colormaps.ColormapOptions("turbo"))
@@ -124,7 +124,7 @@ def activate_stream(sem_map,
             output_path_compo = image_name / 'composited' / f'{clip_model.positives[k]}_{i}'
             output_path_compo.parent.mkdir(exist_ok=True, parents=True)
             colormap_saving(valid_composited, colormap_options, output_path_compo)
-            
+
             # truncate the heatmap into mask
             output = valid_map[i][k]
             output = output - torch.min(output)
@@ -136,7 +136,7 @@ def activate_stream(sem_map,
             mask_pred = smooth(mask_pred)
             mask_lvl[i] = mask_pred
             mask_gt = img_ann[clip_model.positives[k]]['mask'].astype(np.uint8)
-            
+
             # calculate iou
             intersection = np.sum(np.logical_and(mask_gt, mask_pred))
             union = np.sum(np.logical_or(mask_gt, mask_pred))
@@ -148,10 +148,10 @@ def activate_stream(sem_map,
             score = valid_map[i, k].max()
             score_lvl[i] = score
         chosen_lvl = torch.argmax(score_lvl)
-        
+
         chosen_iou_list.append(iou_lvl[chosen_lvl])
         chosen_lvl_list.append(chosen_lvl.cpu().numpy())
-        
+
         # save for visulsization
         save_path = image_name / f'chosen_{clip_model.positives[k]}.png'
         vis_mask_save(mask_lvl[chosen_lvl], save_path)
@@ -165,19 +165,19 @@ def lerf_localization(sem_map, image, clip_model, image_name, img_ann):
 
     valid_map = clip_model.get_max_across(sem_map)                 # 3xkx832x1264
     n_head, n_prompt, h, w = valid_map.shape
-    
+
     # positive prompts
     acc_num = 0
     positives = list(img_ann.keys())
     for k in range(len(positives)):
         select_output = valid_map[:, k]
-        
+
         # NOTE 平滑后的激活值图中找最大值点
         scale = 30
         kernel = np.ones((scale,scale)) / (scale**2)
         np_relev = select_output.cpu().numpy()
         avg_filtered = cv2.filter2D(np_relev.transpose(1,2,0), -1, kernel)
-        
+
         score_lvl = np.zeros((n_head,))
         coord_lvl = []
         for i in range(n_head):
@@ -188,29 +188,29 @@ def lerf_localization(sem_map, image, clip_model, image_name, img_ann):
 
         selec_head = np.argmax(score_lvl)
         coord_final = coord_lvl[selec_head]
-        
+
         for box in img_ann[positives[k]]['bboxes'].reshape(-1, 4):
             flag = 0
             x1, y1, x2, y2 = box
             x_min, x_max = min(x1, x2), max(x1, x2)
             y_min, y_max = min(y1, y2), max(y1, y2)
             for cord_list in coord_final:
-                if (cord_list[0] >= x_min and cord_list[0] <= x_max and 
+                if (cord_list[0] >= x_min and cord_list[0] <= x_max and
                     cord_list[1] >= y_min and cord_list[1] <= y_max):
                     acc_num += 1
                     flag = 1
                     break
             if flag != 0:
                 break
-        
+
         # NOTE 将平均后的结果与原结果相加，抑制噪声并保持激活边界清晰
         avg_filtered = torch.from_numpy(avg_filtered[..., selec_head]).unsqueeze(-1).to(select_output.device)
         torch_relev = 0.5 * (avg_filtered + select_output[selec_head].unsqueeze(-1))
         p_i = torch.clip(torch_relev - 0.5, 0, 1)
         valid_composited = colormaps.apply_colormap(p_i / (p_i.max() + 1e-6), colormaps.ColormapOptions("turbo"))
         mask = (torch_relev < 0.5).squeeze()
         valid_composited[mask, :] = image[mask, :] * 0.3
-        
+
         save_path = output_path_loca / f"{positives[k]}.png"
         show_result(valid_composited.cpu().numpy(), coord_final,
                     img_ann[positives[k]]['bboxes'], save_path)
@@ -238,7 +238,7 @@ def evaluate(feat_dir, output_path, ae_ckpt_path, json_folder, mask_thresh, enco
 
     # instantiate autoencoder and openclip
     clip_model = OpenCLIPNetwork(device)
-    checkpoint = torch.load(ae_ckpt_path, map_location=device)
+    checkpoint = torch.load(ae_ckpt_path, map_location=device, weights_only=False)
     model = Autoencoder(encoder_hidden_dims, decoder_hidden_dims).to(device)
     model.load_state_dict(checkpoint)
     model.eval()
@@ -248,7 +248,7 @@ def evaluate(feat_dir, output_path, ae_ckpt_path, json_folder, mask_thresh, enco
     for j, idx in enumerate(tqdm(eval_index_list)):
         image_name = Path(output_path) / f'{idx+1:0>5}'
         image_name.mkdir(exist_ok=True, parents=True)
-        
+
         sem_feat = compressed_sem_feats[:, j, ...]
         sem_feat = torch.from_numpy(sem_feat).float().to(device)
         rgb_img = cv2.imread(image_paths[j])[..., ::-1]
@@ -259,10 +259,10 @@ def evaluate(feat_dir, output_path, ae_ckpt_path, json_folder, mask_thresh, enco
             lvl, h, w, _ = sem_feat.shape
             restored_feat = model.decode(sem_feat.flatten(0, 2))
             restored_feat = restored_feat.view(lvl, h, w, -1)           # 3x832x1264x512
-        
+
         img_ann = gt_ann[f'{idx}']
         clip_model.set_positives(list(img_ann.keys()))
-        
+
         c_iou_list, c_lvl = activate_stream(restored_feat, rgb_img, clip_model, image_name, img_ann,
                                             thresh=mask_thresh, colormap_options=colormap_options)
         chosen_iou_all.extend(c_iou_list)
@@ -290,8 +290,8 @@ def seed_everything(seed_value):
     np.random.seed(seed_value)
     torch.manual_seed(seed_value)
     os.environ['PYTHONHASHSEED'] = str(seed_value)
-    
-    if torch.cuda.is_available(): 
+
+    if torch.cuda.is_available():
         torch.cuda.manual_seed(seed_value)
         torch.cuda.manual_seed_all(seed_value)
         torch.backends.cudnn.deterministic = True
@@ -301,7 +301,7 @@ def seed_everything(seed_value):
 if __name__ == "__main__":
     seed_num = 42
     seed_everything(seed_num)
-    
+
     parser = ArgumentParser(description="prompt any label")
     parser.add_argument("--dataset_name", type=str, default=None)
     parser.add_argument('--feat_dir', type=str, default=None)

process.sh

@@ -1,24 +1,28 @@
 #!/bin/bash
 
+set -e
+
 # get the language feature of the scene
-python preprocess.py --dataset_name $dataset_path
+python preprocess.py --dataset_path $dataset_path
 
 # train the autoencoder
 cd autoencoder
-python train.py --dataset_path $dataset_path --encoder_dims 256 128 64 32 3 --decoder_dims 16 32 64 128 256 256 512 --lr 0.0007 --dataset_name ae_ckpt
+python train.py --dataset_path $dataset_path --encoder_dims 256 128 64 32 3 --decoder_dims 16 32 64 128 256 256 512 --lr 0.0007 --dataset_name $dataset_name
 # e.g. python train.py --dataset_path ../data/sofa --encoder_dims 256 128 64 32 3 --decoder_dims 16 32 64 128 256 256 512 --lr 0.0007 --dataset_name sofa
 
 # get the 3-dims language feature of the scene
-python test.py --dataset_name $dataset_path --dataset_name $dataset_name
+python test.py --dataset_path $dataset_path --dataset_name $dataset_name
 # e.g. python test.py --dataset_path ../data/sofa --dataset_name sofa
 
+cd ..
+
 # ATTENTION: Before you train the LangSplat, please follow https://github.com/graphdeco-inria/gaussian-splatting
 # to train the RGB 3D Gaussian Splatting model.
 # put the path of your RGB model after '--start_checkpoint'
 
 for level in 1 2 3
 do
-    python train.py -s $dataset_path -m output/${casename} --start_checkpoint $dataset_path/$casename/chkpnt30000.pth --feature_level ${level}
+    python train.py -s $dataset_path -m output/${casename} --start_checkpoint $dataset_path/output/$casename/chkpnt30000.pth --feature_level ${level}
     # e.g. python train.py -s data/sofa -m output/sofa --start_checkpoint data/sofa/sofa/chkpnt30000.pth --feature_level 3
 done
 

render.py

@@ -3,7 +3,7 @@
 # GRAPHDECO research group, https://team.inria.fr/graphdeco
 # All rights reserved.
 #
-# This software is free for non-commercial, research and evaluation use 
+# This software is free for non-commercial, research and evaluation use
 # under the terms of the LICENSE.md file.
 #
 # For inquiries contact  george.drettakis@inria.fr
@@ -39,26 +39,26 @@ def render_set(model_path, source_path, name, iteration, views, gaussians, pipel
             rendering = output["render"]
         else:
             rendering = output["language_feature_image"]
-            
+
         if not args.include_feature:
             gt = view.original_image[0:3, :, :]
-            
+
         else:
             gt, mask = view.get_language_feature(os.path.join(source_path, args.language_features_name), feature_level=args.feature_level)
 
         np.save(os.path.join(render_npy_path, '{0:05d}'.format(idx) + ".npy"),rendering.permute(1,2,0).cpu().numpy())
         np.save(os.path.join(gts_npy_path, '{0:05d}'.format(idx) + ".npy"),gt.permute(1,2,0).cpu().numpy())
         torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(idx) + ".png"))
         torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ".png"))
-               
+
 def render_sets(dataset : ModelParams, iteration : int, pipeline : PipelineParams, skip_train : bool, skip_test : bool, args):
     with torch.no_grad():
         gaussians = GaussianModel(dataset.sh_degree)
         scene = Scene(dataset, gaussians, shuffle=False)
         checkpoint = os.path.join(args.model_path, 'chkpnt30000.pth')
-        (model_params, first_iter) = torch.load(checkpoint)
+        (model_params, first_iter) = torch.load(checkpoint, weights_only=False)
         gaussians.restore(model_params, args, mode='test')
-        
+
         bg_color = [1,1,1] if dataset.white_background else [0, 0, 0]
         background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
 
@@ -70,7 +70,7 @@ def render_sets(dataset : ModelParams, iteration : int, pipeline : PipelineParam
 
 if __name__ == "__main__":
     # Set up command line argument parser
-    
+
     parser = ArgumentParser(description="Testing script parameters")
     model = ModelParams(parser, sentinel=True)
     pipeline = PipelineParams(parser)