将对象重新引入视频目标分割

Ho Kei Cheng, Seoung Wug Oh, Brian Price, Joon-Young Lee, Alexander Schwing

伊利诺伊大学厄巴纳-香槟分校和Adobe

CVPR 2024，亮点

[arXiV] [PDF] [项目页面]

亮点

Cutie是一个视频目标分割框架——它是XMem的后续工作，具有更好的一致性、鲁棒性和速度。 该仓库包含标准视频目标分割的代码以及用于交互式视频分割的GUI工具。 GUI工具还额外提供了“永久记忆”（来自XMem++）选项，以实现更好的可控性。

演示视频

https://github.com/hkchengrex/Cutie/assets/7107196/83a8abd5-369e-41a9-bb91-d9cc1289af70

来源：https://raw.githubusercontent.com/hkchengrex/Cutie/main/docs/sources.txt

安装

仅在Ubuntu上测试过。

先决条件：

• Python 3.8+
• PyTorch 1.12+及对应的torchvision

克隆我们的仓库：

git clone https://github.com/hkchengrex/Cutie.git

使用pip安装：

cd Cutie
pip install -e .

（如果遇到“setup.py文件未找到”的错误，请使用pip install --upgrade pip升级pip）

下载预训练模型：

python cutie/utils/download_models.py

快速入门

脚本演示

如果你想在自己的项目中使用Cutie，这可能是最好的起点。希望脚本能够自明（scripting_demo.py中有额外注释）。如果仍有疑问，欢迎提交问题。更高级的用法，例如添加或删除对象，请参阅scripting_demo_add_del_objects.py。

@torch.inference_mode()
@torch.cuda.amp.autocast()
def main():

    cutie = get_default_model()
    processor = InferenceCore(cutie, cfg=cutie.cfg)
    # 处理器会将输入图像的短边调整为这个尺寸
    # 你可以尝试不同的尺寸，-1表示保持原始尺寸
    processor.max_internal_size = 480

    image_path = './examples/images/bike'
    images = sorted(os.listdir(image_path))  # 顺序很重要
    mask = Image.open('./examples/masks/bike/00000.png')
    palette = mask.getpalette()
    objects = np.unique(np.array(mask))
    objects = objects[objects != 0].tolist()  # 背景“0”不算作对象
    mask = torch.from_numpy(np.array(mask)).cuda()

    for ti, image_name in enumerate(images):
        image = Image.open(os.path.join(image_path, image_name))
        image = to_tensor(image).cuda().float()

        if ti == 0:
            output_prob = processor.step(image, mask, objects=objects)
        else:
            output_prob = processor.step(image)

        # 将输出概率转换为对象掩码
        mask = processor.output_prob_to_mask(output_prob)

        # 可视化预测结果
        mask = Image.fromarray(mask.cpu().numpy().astype(np.uint8))
        mask.putpalette(palette)
        mask.show()  # 或者使用mask.save(...)将其保存到某处


main()

交互式演示

通过以下命令启动交互式演示：

python interactive_demo.py --video ./examples/example.mp4 --num_objects 1

更多说明请见这里。 如果你在远程服务器上运行此程序，可以启用X11转发功能。首先使用ssh -X。可能还需要一些额外配置，不过Google可能会比我能提供更多的帮助。

（对于单个视频的评估，请参阅来自https://github.com/hkchengrex/Cutie/pull/16的非官方脚本`scripts/process_video.py`）

训练与评估

1. 在视频目标分割数据上运行Cutie。
2. 训练Cutie。

引用

@inproceedings{cheng2023putting,
  title={Putting the Object Back into Video Object Segmentation},
  author={Cheng, Ho Kei and Oh, Seoung Wug and Price, Brian and Lee, Joon-Young and Schwing, Alexander},
  booktitle={arXiv},
  year={2023}
}

参考文献

• GUI工具使用了RITM进行交互式图像分割。该仓库也在gui/ritm中重新发布了他们的代码。这部分代码遵循RITM的许可协议。

• 对于自动视频分割或与外部检测器的集成，请参阅DEVA。

• 交互式演示基于IVS、MiVOS和XMem开发而成。

• 我们在视频修复演示中使用了ProPainter。

• 感谢RITM和XMem++，使这一切成为可能。

Cutie 快速上手指南

Cutie 是一个先进的视频对象分割（Video Object Segmentation, VOS）框架，作为 XMem 的继任者，它在一致性、鲁棒性和速度上均有显著提升。本工具支持标准视频分割脚本调用，并提供带有“永久记忆”功能的交互式 GUI 工具。

环境准备

系统要求：

• 操作系统：仅限 Ubuntu (已在 Ubuntu 上测试通过)
• Python 版本：3.8 或更高
• PyTorch 版本：1.12 或更高（需安装对应的 torchvision）

硬件建议：

• 推荐使用 NVIDIA GPU 以获得最佳推理速度。

安装步骤

1. 克隆仓库

   git clone https://github.com/hkchengrex/Cutie.git
   

2. 进入目录并安装依赖

   cd Cutie
   pip install -e .
   

   注意：如果遇到 File "setup.py" not found 错误，请先升级 pip：

   pip install --upgrade pip
   

3. 下载预训练模型 运行以下脚本自动下载所需的预训练权重：

   python cutie/utils/download_models.py
   

基本使用

方式一：脚本化演示（推荐用于项目开发）

这是将 Cutie 集成到自己项目中的最佳起点。以下示例展示了如何加载模型并对图像序列进行推理。

@torch.inference_mode()
@torch.cuda.amp.autocast()
def main():

    cutie = get_default_model()
    processor = InferenceCore(cutie, cfg=cutie.cfg)
    # 处理器会将输入图像的短边调整为此大小
    # 您可以尝试不同的大小，-1 表示保持原始尺寸
    processor.max_internal_size = 480

    image_path = './examples/images/bike'
    images = sorted(os.listdir(image_path))  # 顺序非常重要
    mask = Image.open('./examples/masks/bike/00000.png')
    palette = mask.getpalette()
    objects = np.unique(np.array(mask))
    objects = objects[objects != 0].tolist()  # 背景 "0" 不计入对象
    mask = torch.from_numpy(np.array(mask)).cuda()

    for ti, image_name in enumerate(images):
        image = Image.open(os.path.join(image_path, image_name))
        image = to_tensor(image).cuda().float()

        if ti == 0:
            # 第一帧需要传入初始掩码和对象列表
            output_prob = processor.step(image, mask, objects=objects)
        else:
            # 后续帧仅需传入图像
            output_prob = processor.step(image)

        # 将输出概率转换为对象掩码
        mask = processor.output_prob_to_mask(output_prob)

        # 可视化预测结果
        mask = Image.fromarray(mask.cpu().numpy().astype(np.uint8))
        mask.putpalette(palette)
        mask.show()  # 或使用 mask.save(...) 保存结果


main()

方式二：交互式演示（GUI 工具）

如果您希望通过图形界面手动标注首帧并进行自动跟踪，可以使用交互式 Demo。

启动命令：

python interactive_demo.py --video ./examples/example.mp4 --num_objects 1

提示：如果在远程服务器上运行且需要显示界面，请使用 ssh -X 开启 X11 转发。更多详细操作说明请参阅项目文档 docs/INTERACTIVE.md。