【小贪】项目实战——Zero-shot根据文字提示分割出图片目标掩码

贪钱算法还我头发

已于 2024-07-01 19:42:56 修改

阅读量1k

点赞数 20

分类专栏： # Deep Learning # Computer Vision AI 文章标签：目标检测深度学习 python 语义分割 Zero-shot

于 2024-07-01 19:32:09 首次发布

本文链接：https://blog.youkuaiyun.com/keiven_/article/details/140108422

版权

目标描述

给定RGB视频或图片，目标是分割出图像中的指定目标掩码。我们需要复现两个Zero-shot的开源项目，分别为IDEA研究院的GroundingDINO和Facebook的SAM。首先使用目标检测方法GroundingDINO，输入想检测目标的文字提示，可以获得目标的anchor box。将上一步获得的box信息作为SAM的提示，分割出目标mask。具体效果如下（测试数据来自VolumeDeform数据集）：

在这里插入图片描述

其中GroundingDINO根据white shirt的文字输入计算的box信息为："shirt_000500": "[194.23726, 2.378189, 524.09503, 441.5135]"。项目实测下来单张图片的预测速度GroundingDINO要慢于SAM。GroundingDINO和SAM均会给出多个预测结果，当选择置信度最高的结果时两个模型也会存在预测不准确的情况。

论文简介

GroundingDINO

GroundingDINO extends a closedset detector DINO by performing vision-language modality fusion at multiple phases, including a feature enhancer, a language-guided query selection module, and a cross-modality decoder. Such a deep fusion strategy effectively improves open-set object detection.

在这里插入图片描述

SAM

简介：使用三个组件建立图像分割的foundation model，解决一系列下游分割问题，可zero-shot生成
关键技术：
1. promptable分割任务：使用prompt engineering，prompt不确定时输出多目标mask
2. 分割模型：image encoder + prompt encoder -> mask decoder
3. 数据驱动：SA-1B（1B masks from 11M imgs）手工标注->半自动->全自动
Limitation：存在不连贯不精细的mask结果；交互式实时mask生成但是img encoder耗时；text-to-mask任务效果不鲁棒

在这里插入图片描述

项目实战

两个项目的复现很简单，按照github的readme配置相关环境并运行程序。当然也可以直接使用一站式项目Grounded Segment Anything等。当需要分割的图片较多时，可以修改GroundingDINO的demo.sh和demo/inference_on_a_image.py文件将检测结果保存至json文件。

demo/inference_on_a_image.py文件

# 修改plot_boxes_to_image函数输出box信息
image_with_box, mask, box_coor = plot_boxes_to_image(image_pil, pred_dict)
# obj为目标名称，i为当前图片的索引
obj = 'shirt'
data = {
   f'{
     obj}_{
     str(i).zfill(6)}': str(list(box_coor.cpu().detach().numpy()))}
with open("box.json", "r", encoding="utf-8") as f:
    old_data = json.load(f)
    old_data.update(data)
with open("box.json", "w", encoding="utf-8") as f:
    json.dump(old_data, f, indent=4)
    # f.write(json.dumps(old_data, indent=4, ensure_ascii=False))
f.close()

然后SAM再读取json文件获取box信息，将SAM的输入提示改为box。

测试代码

import os
import numpy as np
import matplotlib.pyplot as plt
import cv2
import glob
import json

coords = []

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)

最低0.47元/天解锁文章