Function SAM2_Refine

preprocess/track_FrameINO.py:215–398

(sam2_predictor, all_frames, GPU_offset, visual_store_folder, obj_Track_Traj, obj_Track_Visibility, 
                    original_start_frame_idx, potential_useful_region_box, min_area_required, debug)

def SAM2_Refine(sam2_predictor, all_frames, GPU_offset, visual_store_folder, obj_Track_Traj, obj_Track_Visibility, 
                    original_start_frame_idx, potential_useful_region_box, min_area_required, debug):


    # Prepare the folder needed
    tmp_store_folder = "tmp_SAM2/process" + str(GPU_offset)
    if os.path.exists(tmp_store_folder):
        shutil.rmtree(tmp_store_folder)
    os.makedirs(tmp_store_folder)
    if os.path.exists(visual_store_folder):
        shutil.rmtree(visual_store_folder)
    os.makedirs(visual_store_folder)


    # Init the tracking
    panoptic_track_points_start = obj_Track_Traj[original_start_frame_idx]
    point_valid_status = [True for _ in range(len(panoptic_track_points_start))]


    # Write to a temp storage available
    sam_frames = all_frames[original_start_frame_idx : ]
    for frame_idx, frame in enumerate(sam_frames):
        store_img_path = os.path.join(tmp_store_folder, str(frame_idx).zfill(4) + ".jpg")
        cv2.imwrite(store_img_path, cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))


    # SAM Process
    info, bbox_info = [], []
    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):

        # Init the state
        state = sam2_predictor.init_state(tmp_store_folder)
        sam2_predictor.reset_state(state)

        # Add new prompts and instantly get the output on the same frame
        labels = np.array([1] * len(panoptic_track_points_start), np.int32)  # All are the same label
        frame_idx, object_ids, masks = sam2_predictor.add_new_points_or_box(
                                                                                state, 
                                                                                frame_idx = 0,
                                                                                obj_id = 1,     # Only consider single isntance now
                                                                                points = panoptic_track_points_start,  # Use points in the first frame
                                                                                labels = labels,
                                                                            )

        # Iterate all frames and Recognize multiple masks
        for frame_idx, object_ids, masks in sam2_predictor.propagate_in_video(state, start_frame_idx=0):
            for obj_idx, out_obj_id in enumerate(object_ids):

                # Convert to boolean mask and 3 channels
                segmentation_mask_raw = (masks[obj_idx] > 0.0).cpu().numpy().astype(np.uint8)
                _, height, width = segmentation_mask_raw.shape
                segmentation_mask_cat = np.stack([segmentation_mask_raw, segmentation_mask_raw, segmentation_mask_raw], axis=-1)[0]     

                # Prepare masked segmentation images
                segmentation_img = segmentation_mask_cat * cv2.cvtColor(sam_frames[frame_idx], cv2.COLOR_BGR2RGB)
                segmentation_mask = segmentation_mask_raw * 255