Update README.md

README.md

@@ -18,18 +18,27 @@ import os
 import torch
 import matplotlib.pyplot as plt
 from PIL import Image
+from tqdm import tqdm
+from natsort import natsorted
 
 from sam2_plus.build_sam import build_sam2_video_predictor_plus
 
 from tools.visualization import show_mask, show_box, show_points
 from tools.vos_inference import load_ann_png, get_per_obj_mask, DAVIS_PALETTE, save_masks_to_dir
 
-checkpoint = "./checkpoints/SAM2-Plus/checkpoint_phase123.pt"
-model_cfg = "configs/sam2.1/sam2.1_hiera_b+_predmasks_decoupled_MAME.yaml"
-predictor = build_sam2_video_predictor_plus(model_cfg, checkpoint, task="mask")
-
-input_video_dir = "./examples/JPEGImages/00001"
-input_mask_path = "./examples/Annotations/00001/00000.png"
+predictor = build_sam2_video_predictor_plus(
+    config_file="configs/sam2.1/sam2.1_hiera_b+_predmasks_decoupled_MAME.yaml",
+    ckpt_path="./checkpoints/SAM2-Plus/checkpoint_phase123.pt",
+    apply_postprocessing=False,
+    hydra_overrides_extra=[
+        "++model.non_overlap_masks=" + ("false")
+    ],
+    vos_optimized=False,
+    task='mask'
+)
+
+input_video_dir = "./examples/JPEGImages/horsejump-low"
+input_mask_path = "./examples/Annotations/horsejump-low/00000.png"
 output_mask_dir = "./output/Annotations/"
 
 score_thresh = 0
@@ -41,9 +50,9 @@ with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
     frame_names = [
         os.path.splitext(p)[0]
         for p in os.listdir(input_video_dir)
-        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
+        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG", ".png", ".PNG"]
     ]
-    frame_names.sort()
+    frame_names = natsorted(frame_names)
     height = inference_state["video_height"]
     width = inference_state["video_width"]
 
@@ -89,7 +98,7 @@ with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
         )
     
     # visualize the tracking results
-    for out_frame_idx in range(0, len(frame_names)):
+    for out_frame_idx in tqdm(range(0, len(frame_names)), desc="Visualization Results"):
         plt.clf()
         plt.figure()
         # plt.title(f"frame {out_frame_idx}")
@@ -99,6 +108,7 @@ with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
         plt.axis('off')
         plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
         plt.savefig(f"{output_mask_dir}/{video_name}/{out_frame_idx:05d}_withMask.png", dpi=300, bbox_inches='tight', pad_inches=0)
+        plt.close()
 ```
 
 **[Video Object Tracking (Box Granularity)](https://github.com/MCG-NJU/SAM2-Plus/blob/main/tools/sot_inference_plus.sh)**
@@ -107,6 +117,8 @@ import os
 import torch
 import matplotlib.pyplot as plt
 from PIL import Image
+from tqdm import tqdm
+from natsort import natsorted
 import numpy as np
 import logging
 
@@ -118,12 +130,19 @@ from tools.sot_inference import save_boxes_to_dir, save_masks_and_boxes_to_dir
 from training.dataset_plus.box.utils import np_box_xywh_to_xyxy, np_box_xyxy_to_xywh, np_masks_to_boxes, np_box_clamp_xywh
 from benchmarks.sot_benchmark.datasets.utils import load_text
 
-checkpoint = "./checkpoints/SAM2-Plus/checkpoint_phase123.pt"
-model_cfg = "configs/sam2.1/sam2.1_hiera_b+_predmasks_decoupled_MAME.yaml"
-predictor = build_sam2_video_predictor_plus(model_cfg, checkpoint, task="box")
-
-input_video_dir = "./examples/JPEGImages/00001"
-input_box_path = "./examples/Boxes/00001.txt"
+predictor = build_sam2_video_predictor_plus(
+    config_file="configs/sam2.1/sam2.1_hiera_b+_predmasks_decoupled_MAME.yaml",
+    ckpt_path="./checkpoints/SAM2-Plus/checkpoint_phase123.pt",
+    apply_postprocessing=False,
+    hydra_overrides_extra=[
+        "++model.non_overlap_masks=" + ("false")
+    ],
+    vos_optimized=False,
+    task='box'
+)
+
+input_video_dir = "./examples/JPEGImages/horsejump-low"
+input_box_path = "./examples/Boxes/horsejump-low.txt"
 output_box_dir = "./output/Boxes/"
 
 score_thresh = 0
@@ -135,9 +154,9 @@ with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
     frame_names = [
         os.path.splitext(p)[0]
         for p in os.listdir(input_video_dir)
-        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
+        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG", ".png", ".PNG"]
     ]
-    frame_names.sort()
+    frame_names = natsorted(frame_names)
     height = inference_state["video_height"]
     width = inference_state["video_width"]
 
@@ -148,7 +167,8 @@ with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
     if os.path.isfile(input_box_path):
         input_box_xywh = load_text(str(input_box_path), delimiter=',', dtype=np.float64, backend='numpy').reshape(-1, 4)[0]
     else:
-        input_box_xywh = [1026,361,222,169]
+        print(f"Box file {input_box_path} not found. Using default box.")
+        input_box_xywh = [316,385,742,488]
     per_obj_input_box_xyxy = {1: np_box_xywh_to_xyxy(np.array(input_box_xywh))}
     object_box_xyxy = per_obj_input_box_xyxy[object_id]
 
@@ -189,7 +209,7 @@ with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
     
     # visualize the tracking results
     os.makedirs(os.path.join(output_box_dir, video_name), exist_ok=True)
-    for out_frame_idx in range(0, len(frame_names)):
+    for out_frame_idx in tqdm(range(0, len(frame_names)), desc="Visualization Results"):
         plt.clf()
         plt.figure()
         # plt.title(f"frame {out_frame_idx}")
@@ -201,6 +221,7 @@ with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
         plt.axis('off')
         plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
         plt.savefig(os.path.join(output_box_dir, video_name, f"{out_frame_idx:05d}_withbox.png"), dpi=300, bbox_inches='tight', pad_inches=0)
+        plt.close()
 ```
 
 **[Point Tracking (Point Granularity)](https://github.com/MCG-NJU/SAM2-Plus/blob/main/tools/pt_inference_plus.sh)**
@@ -209,51 +230,53 @@ import os
 import torch
 import matplotlib.pyplot as plt
 from PIL import Image
+from tqdm import tqdm
 import numpy as np
-import logging
+from natsort import natsorted
 
 from sam2_plus.build_sam import build_sam2_video_predictor_plus
 
 from tools.visualization import show_mask, show_box, show_points
 from tools.vos_inference import load_ann_png, get_per_obj_mask, DAVIS_PALETTE, save_masks_to_dir
-from tools.sot_inference import save_boxes_to_dir, save_masks_and_boxes_to_dir
 from tools.pt_inference_plus import load_visible_points_from_npz
-from training.dataset_plus.box.utils import np_box_xywh_to_xyxy, np_box_xyxy_to_xywh, np_masks_to_boxes, np_box_clamp_xywh
-from benchmarks.sot_benchmark.datasets.utils import load_text
-
-checkpoint = "./checkpoints/SAM2-Plus/checkpoint_phase123.pt"
-model_cfg = "configs/sam2.1/sam2.1_hiera_b+_predmasks_decoupled_MAME.yaml"
-predictor = build_sam2_video_predictor_plus(model_cfg, checkpoint, task="point")
 
-input_video_dir = "./examples/JPEGImages/00001"
-input_point_path = "./examples/Points/00001.npz"
+predictor = build_sam2_video_predictor_plus(
+    config_file="configs/sam2.1/sam2.1_hiera_b+_predmasks_decoupled_MAME.yaml",
+    ckpt_path="./checkpoints/SAM2-Plus/checkpoint_phase123.pt",
+    apply_postprocessing=False,
+    hydra_overrides_extra=[
+        "++model.non_overlap_masks=" + ("false")
+    ],
+    vos_optimized=False,
+    task='point'
+)
+
+input_video_dir = "./examples/JPEGImages/horsejump-low"
+input_point_path = "./examples/Points/horsejump-low.npz"
 output_point_dir = "./output/Points/"
 
 radius, sigma = 5, 2
-
 score_thresh = 0
 
 with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
-    inference_state = predictor.init_state(video_path=input_video_dir)
-
     video_name = os.path.basename(input_video_dir)
     frame_names = [
         os.path.splitext(p)[0]
         for p in os.listdir(input_video_dir)
-        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
+        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG", ".png", ".PNG"]
     ]
-    frame_names.sort()
+    frame_names = natsorted(frame_names)
+
+    inference_state = predictor.init_state(video_path=input_video_dir)
     height = inference_state["video_height"]
     width = inference_state["video_width"]
 
     input_frame_idx = 0     # the frame index we interact with
-    object_id = 1           # give a unique id to each object we interact with (it can be any integers)
+    object_id = 0           # give a unique id to each object we interact with (it can be any integers)
+    num_frames, num_points = len(frame_names), 1
 
-    input_palette = None
-    gt_data = np.load(input_point_path, allow_pickle=True)
-    trajs = gt_data['trajs_2d'].astype(np.float32)    # ndarray [N_frames, N_points, 2], xyxy
-    visible = gt_data['visibs'].astype(bool)    # ndarray [N_frames, N_points], bool
-    input_point, input_visible = torch.tensor(trajs), torch.tensor(visible)
+    input_data = np.load(input_point_path, allow_pickle=True)
+    input_point, input_visible = torch.tensor(input_data['trajs_2d'].astype(np.float32)), torch.tensor(input_data['visibs'].astype(bool))
     per_obj_input_point = load_visible_points_from_npz(
         input_points=input_point,
         input_visibles=input_visible,
@@ -272,7 +295,6 @@ with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
     )
 
     # run propagation throughout the video and collect the results in a dict
-    num_frames, num_points = len(frame_names), trajs.shape[1]
     point_array = -np.ones((num_frames, num_points, 2), dtype=np.float32)
     visible_array = np.zeros((num_frames, num_points), dtype=bool)
     for out_frame_idx, out_obj_ids, out_mask_logits, out_box_xyxys, out_obj_score_logits in predictor.propagate_in_video(
@@ -284,14 +306,14 @@ with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
             max_score_y, max_score_x = torch.unravel_index(max_index, out_mask_logit.shape)
             point_array[out_frame_idx, out_obj_id] = np.array([max_score_x.cpu(), max_score_y.cpu()])
             visible_array[out_frame_idx, out_obj_id] = (out_obj_score_logit > score_thresh).cpu().numpy()
-    
+
     # write the output masks as palette PNG files to output_mask_dir
     os.makedirs(output_point_dir, exist_ok=True)
     np.savez(os.path.join(output_point_dir, f"{video_name}.npz"), trajs_2d=point_array, visibs=visible_array, size=(width, height))
-    
+
     # visualize the tracking results
     os.makedirs(os.path.join(output_point_dir, video_name), exist_ok=True)
-    for out_frame_idx in range(0, len(frame_names)):
+    for out_frame_idx in tqdm(range(0, len(frame_names)), desc="Visualization Results"):
         plt.clf()
         plt.figure()
         # plt.title(f"frame {out_frame_idx}")
@@ -302,6 +324,7 @@ with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
         plt.axis('off')
         plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
         plt.savefig(os.path.join(output_point_dir, video_name, f"{out_frame_idx:05d}_withPoint.png"), dpi=300, bbox_inches='tight', pad_inches=0)
+        plt.close()
 ```
 
 ### Load from 🤗 Hugging Face