分割标注中Sam类set_video_list方法形式参数与SAM2原始仓库不符 & sam_tools.py中forward_sam_multi_stage函数存在的问题

[FelixxLuo]

作者您好，感谢开源。我在尝试使用分割标注时遇到了一些报错，在排查的过程中发现了代码中的一些问题，想在此探讨。

1. Sam类set_video_list方法形式参数数量问题

这里传递了四个参数 (https://github.com/InternRobotics/RoboInter/blob/main/RoboInterTools/tools/sam.py#L96), 但SAM2原始仓库只需要三个参数 (https://github.com/facebookresearch/sam2/blob/main/sam2/sam2_video_predictor.py#L42)

2. sam_tools.py中forward_sam_multi_stage函数存在的问题

sam_tools.py中forward_sam_multi_stage函数调用了set_video_list方法(https://github.com/InternRobotics/RoboInter/blob/main/RoboInterTools/tools/sam_tools.py#L97), 因此并不能成功运行

3. 问题探讨与可能的解决方案 (一孔之见，仅供参考)

forward_sam_multi_stage函数根据forward或是backward将视频分了段, 但SAM2原始仓库只能通过读取video_path加载视频 (https://github.com/facebookresearch/sam2/blob/main/sam2/utils/misc.py#L172)， 无法将分好段的视频video_part直接传递sam_model。

因此，我考虑的做法是修改Sam类set_video_list方法，删去形式参数video_list。另外，修改sam_tools.py中forward_sam_multi_stage函数，将video_part暂存成.mp4文件，然后读取该文件。

修改Sam类set_video_list方法，删去形式参数video_list

def set_video_list(self, video_path):
    if video_path[-3:] !="mp4":
        video_path = video_path + ".mp4"
    self.inference_state = self.predictor.init_state(video_path, offload_video_to_cpu=True, offload_state_to_cpu=True)

修改sam_tools.py中forward_sam_multi_stage函数，将video_part暂存成.mp4文件，然后读取该文件 (需要在开头import os, shutil)

def forward_sam_multi_stage(model_config, model_sam):
    video_path = model_config["video_path"]
    is_video = model_config["is_video"]
    select_frame = model_config["select_frame"]
    direction = model_config["direction"]

    # ----------------------修改开始-----------------------------------------
    temp_save_dir = video_path.rsplit(".", 1)[0] + "_cache"
    if not os.path.exists(temp_save_dir):
        os.makedirs(temp_save_dir)
    # ----------------------修改结束-----------------------------------------
    
    if 'ann_human' in video_path:
        video_path = model_config['origin_video_path']
    video = extract_frames(video_path)
    
    if not is_video:
        video = video[select_frame:select_frame + 1]
    elif direction == "forward":
        video = video[select_frame:]     
    elif direction == "backward":
        video = video[:select_frame+1][::-1]
    
    positive_points_dict = model_config["positive_points"][select_frame]
    negative_points_dict = model_config["negative_points"][select_frame]
    labels_dict = model_config["labels"][select_frame]

    positive_points = [np.array(positive_points_dict[obj_idx]) for obj_idx in positive_points_dict.keys()]
    negative_points = [np.array(negative_points_dict[obj_idx]) for obj_idx in positive_points_dict.keys()]
    labels = [labels_dict[obj_idx] for obj_idx in positive_points_dict.keys()]

    for i in range(len(positive_points)):
        if len(positive_points[i]) == 0:
            raise ValueError("No positive points in the frame")
        if len(negative_points[i]) != 0:
            positive_points[i] = np.concatenate([positive_points[i], negative_points[i]], axis=0)
    
    # if length of video is larger than 800, we split the video into N/800 parts
    num_parts = int(len(video) / 800) + 1
    masks_all = np.zeros((len(positive_points), len(video), 1, video[0].shape[0], video[0].shape[1]))
    ind_all = np.arange(len(video))
    
    for i in range(num_parts):
        video_part = video[i::num_parts]
        ind_part = ind_all[i::num_parts]
        
        # must contain the first frame
        video_part = np.concatenate([video[:1], video_part], axis=0)
        ind_part = np.concatenate([np.zeros_like(ind_part[:1]), ind_part], axis=0)       
        
        
        # ----------------------修改开始-----------------------------------------
        video_part_path = os.path.join(temp_save_dir, f"part_{i}.mp4")
        height, width, _ = video_part[0].shape
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(video_part_path, fourcc, 20, (width, height))

        for frame in video_part:
            out.write(frame)
        out.release()
        
        model_sam.set_video_list(video_part_path)
        masks_all[:, ind_part] = model_sam(positive_points, labels, 0, list(positive_points_dict.keys()))
        
        if os.path.exists(temp_save_dir):
            shutil.rmtree(temp_save_dir)
        # ----------------------修改结束-----------------------------------------

    return masks_all

[hanyangyu1021]

Do you fix this problem?

[FelixxLuo]

Yeah, I fixed it as mentioned above. However, I haven't tested if it works for multi-object masks

[wtt0213]

Hi,

I apologize for the oversight in the previous code. To resolve this, please follow the steps below to replace the original predictor and update the configuration:

1. Replace the Predictor Class

Please add the following SAM2VideoPredictorFast class based on the original SAM2VideoPredictor in sam2/sam2_video_predictor.py

class SAM2VideoPredictorFast(SAM2VideoPredictor):
    
    def load_video_frames(
        self,
        video_list,
        image_size,
        offload_video_to_cpu,
        img_mean=(0.485, 0.456, 0.406),
        img_std=(0.229, 0.224, 0.225),
        compute_device=torch.device("cuda"),
    ):
        num_frames = len(video_list)
        img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None]
        img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None]
        images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float32)
        for n, img in enumerate(video_list):
            images[n], video_height, video_width = self._convert_img_as_tensor(img, image_size)
        if not offload_video_to_cpu:
            images = images.to(compute_device)
            img_mean = img_mean.to(compute_device)
            img_std = img_std.to(compute_device)
        # normalize by mean and std
        images -= img_mean
        images /= img_std
        return images, video_height, video_width
    
    def _convert_img_as_tensor(self, img, image_size):
        img_pil = Image.fromarray(img)
        img_np = np.array(img_pil.resize((image_size, image_size)))
        if img_np.dtype == np.uint8:  # np.uint8 is expected for JPEG images
            img_np = img_np / 255.0
        else:
            raise RuntimeError(f"Unknown image dtype: {img_np.dtype}")
        img = torch.from_numpy(img_np).permute(2, 0, 1)
        video_width, video_height = img_pil.size  # the original video size
        return img, video_height, video_width
    
    def forward_image(self, img_batch):
        """Get the image feature on the input batch."""
        backbone_out = self.image_encoder(img_batch)
        if self.use_high_res_features_in_sam:
            backbone_out["backbone_fpn"][0] = self.sam_mask_decoder.conv_s0(
                backbone_out["backbone_fpn"][0]
            )
            backbone_out["backbone_fpn"][1] = self.sam_mask_decoder.conv_s1(
                backbone_out["backbone_fpn"][1]
            )
        return backbone_out
    
    @torch.inference_mode()
    def init_state(
        self,
        video_list,
        video_path,
        offload_video_to_cpu=False,
        offload_state_to_cpu=False,
        async_loading_frames=False,
    ):
        """Initialize an inference state."""
        compute_device = self.device
        images, video_height, video_width = self.load_video_frames(
            video_list=video_list,
            image_size=self.image_size,
            offload_video_to_cpu=offload_video_to_cpu,
            compute_device=compute_device,
        )
        inference_state = {}
        inference_state["video_path"] = video_path
        inference_state["images"] = images
        inference_state["num_frames"] = len(images)
        inference_state["offload_video_to_cpu"] = offload_video_to_cpu
        inference_state["offload_state_to_cpu"] = offload_state_to_cpu
        inference_state["video_height"] = video_height
        inference_state["video_width"] = video_width
        inference_state["device"] = compute_device
        if offload_state_to_cpu:
            inference_state["storage_device"] = torch.device("cpu")
        else:
            inference_state["storage_device"] = compute_device
        
        inference_state["point_inputs_per_obj"] = {}
        inference_state["mask_inputs_per_obj"] = {}
        inference_state["cached_features"] = {}
        inference_state["constants"] = {}
        inference_state["obj_id_to_idx"] = OrderedDict()
        inference_state["obj_idx_to_id"] = OrderedDict()
        inference_state["obj_ids"] = []
        inference_state["output_dict"] = {
            "cond_frame_outputs": {},
            "non_cond_frame_outputs": {},
        }
        inference_state["output_dict_per_obj"] = {}
        inference_state["temp_output_dict_per_obj"] = {}
        inference_state["consolidated_frame_inds"] = {
            "cond_frame_outputs": set(),
            "non_cond_frame_outputs": set(),
        }
        inference_state["tracking_has_started"] = False
        inference_state["frames_already_tracked"] = {}
        inference_state["cached_image_features"] = {}
        
        self._get_image_feature(inference_state, frame_idx=0, batch_size=1)
        
        return inference_state

2. Update the Build Function

In sam2/build_sam.py, modify the function build_sam2_video_predictor by updating the hydra_overrides content to:

"++model._target_=sam2.sam2_video_predictor.SAM2VideoPredictorFast"

3. Update funtion _get_image_feature in class SAM2VideoPredictor to support Speed Up Feature.

    def _get_image_feature(self, inference_state, frame_idx, batch_size):
        """Compute the image features on a given frame."""
        # Look up in the cache first
        # image, backbone_out = inference_state["cached_features"].get(
        #     frame_idx, (None, None)
        # )
        device = inference_state["device"]
        image = inference_state["images"][frame_idx].to(device).float().unsqueeze(0)
        backbone_out = inference_state["cached_image_features"].get(-1, None)
        
        if backbone_out is None:
            # Cache miss -- we will run inference on a single image
            device = inference_state["device"]
            backbone_out = self.forward_image(image)
            # Cache the most recent frame's feature (for repeated interactions with
            # a frame; we can use an LRU cache for more frames in the future).
            inference_state["cached_features"] = dict(frame_idx=(image, backbone_out))
            inference_state["cached_image_features"][frame_idx] = backbone_out

        # expand the features to have the same dimension as the number of objects
        expanded_image = image.expand(batch_size, -1, -1, -1)
        expanded_backbone_out = {
            "backbone_fpn": backbone_out["backbone_fpn"].copy(),
            "vision_pos_enc": backbone_out["vision_pos_enc"].copy(),
        }
        for i, feat in enumerate(expanded_backbone_out["backbone_fpn"]):
            expanded_backbone_out["backbone_fpn"][i] = feat.expand(
                batch_size, -1, -1, -1
            )
        for i, pos in enumerate(expanded_backbone_out["vision_pos_enc"]):
            pos = pos.expand(batch_size, -1, -1, -1)
            expanded_backbone_out["vision_pos_enc"][i] = pos

        features = self._prepare_backbone_features(expanded_backbone_out)
        features = (expanded_image,) + features
        return features