Update demo_utils.py

demo_utils.py

@@ -11,21 +11,12 @@ import sam2point.dataset as dataset
 import sam2point.configs as configs
 from sam2point.voxelizer import Voxelizer
 from sam2point.utils import cal
-
 import matplotlib.pyplot as plt
 import plotly.graph_objects as go
 
 print("Torch CUDA:", torch.cuda.is_available())
-# use bfloat16 for the entire notebook
 torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
 
-
-# if torch.cuda.get_device_properties(0).major >= 8:
-#    # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
-#    torch.backends.cuda.matmul.allow_tf32 = True
-#    torch.backends.cudnn.allow_tf32 = True
-
-
 def run_demo(dataset_name, prompt_type, sample_idx, prompt_idx, voxel_size, theta, mode, ret_prompt):
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', choices=['S3DIS', 'ScanNet', 'Objaverse', 'KITTI', 'Semantic3D'], default='Objaverse', help='dataset selected')
@@ -33,19 +24,18 @@ def run_demo(dataset_name, prompt_type, sample_idx, prompt_idx, voxel_size, thet
    parser.add_argument('--sample_idx', type=int, default=2, help='the index of the scene or object')
    parser.add_argument('--prompt_idx', type=int, default=0, help='the index of the prompt')   
    parser.add_argument('--voxel_size', type=float, default=0.02, help='voxel size')   
-   parser.add_argument('--theta', type=float, default=0.5)  # indoor NOTE
-   parser.add_argument('--mode', type=str, default='bilinear')  # indoor NOTE
+   parser.add_argument('--theta', type=float, default=0.5)  
+   parser.add_argument('--mode', type=str, default='bilinear') 
    parser.add_argument("--ret_prompt", action="store_true")
    args = parser.parse_args()
    args.dataset, args.prompt_type, args.sample_idx, args.prompt_idx = dataset_name, prompt_type, sample_idx, prompt_idx
    args.voxel_size, args.theta, args.mode, args.ret_prompt = voxel_size, theta, mode, ret_prompt
    print(args)
 
-   #cache
    name_list = [args.dataset, "sample" + str(args.sample_idx), args.prompt_type + "-prompt" + str(args.prompt_idx)]
    name = '_'.join(name_list)
 
-   # hf
+   # use cache result for speeding up
    repo_id = "ZiyuG/Cache"
    result_name = "cache_results/" + name + '.npy'
    prompt_name = "cache_prompt/" + name + '.npy'
@@ -64,7 +54,7 @@ def run_demo(dataset_name, prompt_type, sample_idx, prompt_idx, voxel_size, thet
            PROMPT = np.load("./cache_prompt/" + name + '.npy')
            if not args.ret_prompt: return new_color, PROMPT
            else:   return PROMPT
-   #########
+
    if args.dataset == 'S3DIS':
        info = configs.S3DIS_samples[args.sample_idx]
        # early return
@@ -99,14 +89,12 @@ def run_demo(dataset_name, prompt_type, sample_idx, prompt_idx, voxel_size, thet
        point, color = dataset.load_Semantic3D_sample(info['path'], args.sample_idx)
        args.voxel_size = info[configs.VOXEL[args.prompt_type]][args.prompt_idx]
    
-   
    point_color = np.concatenate([point, color], axis=1)
    voxelizer = Voxelizer(voxel_size=args.voxel_size, clip_bound=None)
   
    labels_in = point[:, :1].astype(int)
    locs, feats, labels, inds_reconstruct = voxelizer.voxelize(point, color, labels_in)
 
-
    if args.prompt_type == 'point':
        if args.ret_prompt:     return list(np.array(info['point_prompts'])[args.prompt_idx])
        mask = seg_point(locs, feats, info['point_prompts'], args)
@@ -145,9 +133,8 @@ def run_demo(dataset_name, prompt_type, sample_idx, prompt_idx, voxel_size, thet
 
    name_list = [args.dataset, "sample" + str(args.sample_idx), args.prompt_type + "-prompt" + str(args.prompt_idx)]
    name = '_'.join(name_list) + 'frames'
-   # os.system('rm -rf ' + name)
 
-   #cache
+   #cache for speeding up
    name_list = [args.dataset, "sample" + str(args.sample_idx), args.prompt_type + "-prompt" + str(args.prompt_idx)]
    name = '_'.join(name_list)
    os.makedirs("cache_results", exist_ok=True)
@@ -156,9 +143,6 @@ def run_demo(dataset_name, prompt_type, sample_idx, prompt_idx, voxel_size, thet
    np.save("./cache_prompt/" + name + '.npy', PROMPT)
    return new_color, PROMPT
 
-
-
-
 def create_box(prompt):
    x_min, y_min, z_min, x_max, y_max, z_max = tuple(prompt)
    bbox_points = np.array([
@@ -171,15 +155,11 @@ def create_box(prompt):
        [x_max, y_max, z_max],
        [x_min, y_max, z_max]
    ])
-
-
    edges = [
        (0, 1), (1, 2), (2, 3), (3, 0), # Bottom face
        (4, 5), (5, 6), (6, 7), (7, 4), # Top face
        (0, 4), (1, 5), (2, 6), (3, 7)  # Vertical edges
    ]
-
-
    bbox_lines = []
    f = 1
    for start, end in edges:
@@ -188,13 +168,9 @@ def create_box(prompt):
            y=[bbox_points[start, 1], bbox_points[end, 1]],
            z=[bbox_points[start, 2], bbox_points[end, 2]],
            mode='lines',
-        #    line=dict(color='red', width=2),  # Customize color and width
-        #    line=dict(color='rgb(255, 140, 0)', width=4),  # Customize color and width
-           line=dict(color='rgb(220, 20, 60)', width=6),  # Customize color and width
+           line=dict(color='rgb(220, 20, 60)', width=6),  
            name="Box Prompt" if f == 1 else "",
            showlegend=True if f == 1 else False
        ))
        f = 0
-   return bbox_lines
-
-
+   return bbox_lines