SAM2：使用mask作为提示输入，实现VOS视频分割

8k50o45u_seg目录1. 引言2. 使用SAM2实现VOS任务2.1 数据集2.2 主要函数2.3 主要代码3. 结果展示1. 引言本文尝试使用SAM2模型来实现VOS任务。由于在官方的github代码中只找到了point或者box作为提示但是论文中却说明是可以输入mask作为提示的所以打算自己尝试一下。官方githubhttps://github.com/facebookresearch/sam22. 使用SAM2实现VOS任务2.1 数据集本文是基于MOSE数据集来进行实验。MOSE 数据集是一个用于视频目标分割Video Object Segmentation, VOS的数据集主要用于评估像 Segment Anything Model 2 这类模型在复杂视频场景中的分割与跟踪能力。MOSE数据集https://mose.video/数据集的格式如下每一个视频都处理一个帧列表train/valid.tar.gz │ ├── Annotations │ ├── video_name_1 │ │ ├── 00000.png │ │ ├── 00001.png │ │ └── ... │ └── video_name_... │ └── ... │ └── JPEGImages ├── video_name_1 │ ├── 00000.png │ ├── 00001.png │ └── ... └── video_name_... └── ...2.2 主要函数在VOS任务中通常会使用第一帧的mask图像作为提示输入然而MOSE数据中的标注是彩色的mask每一种颜色代表一种类别如下图所示而SAM2中所需要的mask提示需要为 Binary Mask所以需要对MOSE数据集中的mask进行处理提取出每一个类别的掩码图。将每种颜色拆分成为单独的 Binary mask图像并保存到列表中。def split_multiclass_mask(mask_path, ignore_backgroundTrue): 读取多类别 mask并将每种颜色拆分成单独的二值 mask 参数: mask_path: mask图像路径 ignore_background: 是否忽略背景(默认忽略黑色) 返回: binary_masks: list每个元素是一个 H×W 的二值 mask colors: 每个 mask 对应的颜色 mask cv2.imread(mask_path) if mask is None: raise ValueError(f无法读取 mask: {mask_path}) # 获取所有颜色 colors np.unique(mask.reshape(-1, 3), axis0) binary_masks [] valid_colors [] for color in colors: # 跳过背景 if ignore_background and np.all(color [0, 0, 0]): continue # 找到该颜色的位置 m np.all(mask color, axis-1) binary_mask (m * 255).astype(np.uint8) binary_masks.append(binary_mask) valid_colors.append(color) return binary_masks, valid_colors虽然输入需要binary mask但是最终生成的掩码图应该是彩色的mask所以使用generate_mask函数生成彩色mask为每一个id生成一种颜色。def generate_mask(obj_id, mask): mask np.squeeze(mask) H, W mask.shape # 生成颜色 cmap plt.get_cmap(tab10) color np.array(cmap(obj_id % 10)[:3]) * 255 color color.astype(np.uint8) # 初始化黑图 mask_img np.zeros((H, W, 3), dtypenp.uint8) mask_bool mask.astype(bool) # 填充颜色 mask_img[mask_bool] color return mask_img本文中希望最终生成的结果是带有掩码的视频和彩色mask图像def merge_masks_and_overlay(frame, all_masks, alpha0.5): frame: 原图 (H,W,3), dtypeuint8 all_masks: list每个元素是 (H,W,3) 彩色mask, dtypeuint8 alpha: 叠加透明度 return: merged_mask : 合并后的mask图 (H,W,3) overlay_img : 原图mask叠加 (H,W,3) H, W frame.shape[:2] # 1 初始化 mask merged_mask np.zeros((H, W, 3), dtypenp.uint8) # 2 合并所有mask for m in all_masks: mask_bool np.any(m 0, axis2) merged_mask[mask_bool] m[mask_bool] # 3 叠加到原图 overlay frame.copy() mask_area np.any(merged_mask 0, axis2) # 使用 numpy 做叠加 overlay[mask_area] ( alpha * merged_mask[mask_area].astype(np.float32) (1 - alpha) * overlay[mask_area].astype(np.float32) ).astype(np.uint8) # cv2.imwrite(frame.png, overlay) return merged_mask, overlay2.3 主要代码配置模型与保存路径if torch.cuda.is_available(): device torch.device(cuda) elif torch.backends.mps.is_available(): device torch.device(mps) else: device torch.device(cpu) print(fusing device: {device}) if device.type cuda: # use bfloat16 for the entire notebook torch.autocast(cuda, dtypetorch.bfloat16).__enter__() # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices) if torch.cuda.get_device_properties(0).major 8: torch.backends.cuda.matmul.allow_tf32 True torch.backends.cudnn.allow_tf32 True elif device.type mps: print( \nSupport for MPS devices is preliminary. SAM 2 is trained with CUDA and might give numerically different outputs and sometimes degraded performance on MPS. See e.g. https://github.com/pytorch/pytorch/issues/84936 for a discussion. ) if __name__ __main__: sam2_checkpoint sam2/checkpoints/sam2.1_hiera_large.pt model_cfg configs/sam2.1/sam2.1_hiera_l.yaml video_path sam2/MOSE_val/JPEGImages anno_path sam2/MOSE_val/Annotations output_dir output/MOSE obj_num 1 os.makedirs(output_dir, exist_okTrue) shutil.rmtree(output_dir) file_count 1 predictor build_sam2_video_predictor(model_cfg, sam2_checkpoint, devicedevice) for file in os.listdir(video_path): if file_count 1: predictor.reset_state(inference_state) torch.cuda.empty_cache() frames_path os.path.join(video_path, file) mask_path os.path.join(anno_path, file, os.listdir(os.path.join(anno_path, file))[0]) video_save_path os.path.join(output_dir, video,file) mask_save_path os.path.join(output_dir, masks,file) os.makedirs(video_save_path, exist_okTrue) os.makedirs(mask_save_path, exist_okTrue)接下来对彩色mask图像进行处理并且将视频帧按顺序排列这一步非常重要不然可能会导致输入的提示mask与输入的帧不一致也是为了保证后续生成符合顺序的掩码视频# 处理mask图像 binary_masks, valid_colors split_multiclass_mask(mask_path) # 对帧按顺序排列 frame_names [ p for p in os.listdir(frames_path) if os.path.splitext(p)[-1].lower() in [.jpg, .jpeg, .png] ] frame_names.sort(keylambda p: int(os.path.splitext(p)[0]))初始化predictor后输入提示掩码注意每一种类别对应着一种obj_id每一个类别对应着的binary mask都需要输入一次不能一次性输入。inference_state predictor.init_state(video_pathframes_path) # take a look the first video frame frame_idx 0 obj_id 1 for mask in binary_masks: _, out_obj_ids, out_mask_logits predictor.add_new_mask( inference_state, frame_idx, obj_id, mask, ) obj_id 1用 SAM2 的视频传播propagation功能对视频每一帧生成目标的分割 mask并把结果保存到video_segments字典中。video_segments {} # video_segments contains the per-frame segmentation results for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state): video_segments[out_frame_idx] { out_obj_id: (out_mask_logits[i] 0.0).cpu().numpy() for i, out_obj_id in enumerate(out_obj_ids) }最后生成结果output_video os.path.join(video_save_path, f{file}_seg.mp4) # 读取第一帧获取尺寸 first_img cv2.imread(os.path.join(frames_path, frame_names[0])) h, w, _ first_img.shape # 创建 VideoWriter out iio.get_writer( output_video, fps5, # 保持 RGB macro_block_size1 ) for frame_idx in range(len(frame_names)): img_path os.path.join(frames_path, frame_names[frame_idx]) frame cv2.imread(img_path) all_masks [] if frame_idx in video_segments: for out_obj_id, out_mask in video_segments[frame_idx].items(): obj_mask_img generate_mask(out_obj_id, out_mask) all_masks.append(obj_mask_img) merge_mask, overlay_frame merge_masks_and_overlay(frame, all_masks) cv2.imwrite(os.path.join(mask_save_path, frame_names[frame_idx]), merge_mask) cv2.imwrite(overlay.png, overlay_frame) overlay_frame cv2.cvtColor(overlay_frame, cv2.COLOR_BGR2RGB) out.append_data(overlay_frame) out.close() print(Saved video:, output_video) file_count 13. 结果展示抽取了一个视频中的前几帧mask结果展示jxmcdk8k_seg