项目链接： github.com/kivenyangmi...
参考论文： arxiv.org/pdf/2306.14...

前言

这段时间大模型是大家研究的重点，其中以【SegmentAnything】为首的模型让我们眼睛一亮，该模型可以对绝大多数图像能够做到较为细致的分割，其中还会返回目标框的bbox坐标。由此我们是否可以利用这一特征为自己写一脚本赋能标注工作呢？这就是本篇博客以BBOX参数为例子进行筛选为主要阐述的内容。

SAM回顾

它能够自动识别和提取图像中的各种形状和对象，包括但不限于线条、区域、轮廓、圆形、多边形等。这项技术的核心优势在于其强大的通用性，能够适应各种不同的图像类型和场景。这很符合官方宣传SAM模型已经学会了关于物体是什么的一般性概念--这种理解使其能够对不熟悉的物体和图像进行概括，而不需要额外的训练。

但是考虑到【SegmentAnything】中最小的vit-l模型仍然是比较大的，这对于普通选手不是怎么友好，因此我在这个项目中采用了【mobilesam】进行推理。【mobilesam】是经过一系列蒸溜得到的模型，其推理方式同【SegmentAnything】，在效率方面都得到了大大提升，精度损失的在日常的标注工作中尚可接受。

赋能工作

为了能够将大模型的推理与标注的工作更好的融合，我们第一步需要将大模型的推理结果获取到，根据推理结果筛选出"目标框"。这里面有两个重点：

通过大模型的推理过程获取图像的输出值
能够通过鼠标交互性操作获取到需要自己的"目标框"

mobileSam推理

在这一步中我们可以参考官网给出的推理过程，也可以参考其他各大博主发布的推理过程，这里我们定义函数加载mobile_sam模型并获取mask输出值：

ini 复制代码

from mobile_sam import SamAutomaticMaskGenerator, sam_model_registry


def load_mask():
    model_type = "vit_t"
    sam_checkpoint = "./weights/mobile_sam.pt"
    mobile_sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
    mobile_sam.to(device="cuda")
    mobile_sam.eval()
    mask_generator = SamAutomaticMaskGenerator(mobile_sam)
    return mask_generator

读取源码部分关于图像推理输入模块，可以得到输入图像的格式需要为【np.ndarray】且为RGB格式。故而后续的调用网络的推理过程中我们选择使用cv2调用网络并将BGR格式数据更改为RGB进行输入。

ini 复制代码

if __name__ == "__main__":
    mask_generator = load_mask()  # 加载网络
    src_img = cv2.imread("car.jpg")
    masks = mask_generator.generate(cv2.cvtColor(src_img, cv2.COLOR_BGR2RGB))

交互操作获取目标框

在这个里面笔者面临的最大问题为如何根据简单的交互操作实现"目标框"的获取？思考如下：

将所有的bbox全部绘制到图像中检验是否能较为合适的匹配我们需要标注的目标；
如何才能将"非目标框"剔除

在笔者通过思考1中的输出检验得知：在大模型的推理下我们是可以获取到较为合适的bbox框，因此通过大模型推理是具有一定的可行性的。在具备了可行性了后我们接着进行下一步：分析每个bbox覆盖的范围，我们【是否可以通过两个"待检测"的目标上的点进行筛选 】出大模型推理的框呢？由于在网络的输出里面含有bbox参数以及area参数参数，那么我们可以根据包含了鼠标指针点击两个点的最小bbox,也就是包含了这两个点的最小面积的bbox。反向整理这句话：获取目标框的最小外接矩形

需要实现上述的需求，需要对大模型的输出进行筛选，获取bbox参数以及area参数，并对area参数进行由小到大的排序并且需要将排序后的指针位置同步给bbox进行排序。

获取推理结果中的坐标框的面积集合
获取推理结果中的bbox坐标集此时是xywh
根据面积由小到大对xywh排序

python 复制代码

def mask_infer(masks):
    area_list = [mask['area'] for mask in masks] 
    xywh_list = [masks[i]['bbox'] for i in range(len(masks))]  # 
    xywh_list = [b for _, b in sorted(zip(area_list, xywh_list))]  # 
    return xywh_list

经过排序后我们可得到了bbox的坐标数组，在后续中我们通过鼠标点击图像中的待标注的目标点，通过点击的点进行筛查返回【第一个】bbox即可，此时我们仅计算了一次面积最小值，避免后续多个目标框多次筛选带来的复杂计算。由于在后面判断指定点是否在bbox内我们需要xyxy格式的坐标因此我们还需要xyxy_list数据：

arduino 复制代码

def build_box_tree(xyxy_list):
    box_tree = SortedDict()
    for i, box in enumerate(xyxy_list):
        box_tree[i] = box
    return box_tree

def is_point_in_box(point, box):
    xmin, ymin, xmax, ymax = box
    x, y = point
    return xmin <= x <= xmax and ymin <= y <= ymax

def is_box_containing_point(box, point):
    return is_point_in_box(point[0], box) and is_point_in_box(point[1], box)

交互性操作

在opencv中可通过cv2.setMouseCallback函数可以捕捉到鼠标指针的操作，我们可以定义如下操作：

左键单击可以获取鼠标指针所在位置的像素坐标
右键单击获取点击目标的最小外接矩形
左键双击清空点击的点以及控制台

定义了这三大操作的含义，后续工作就是逐渐修饰完善这三大操作的细节。我们可定义最多点击两个点（两点成线），当点击第三个点的时候并不会同第二个点进行连接。需要注意的是其实最初在判断指定的两个点是否在大模型推理得到的bbox数组内时采用的为for循环判断，总感觉这种方式很Low，是否有更好的筛选及方式呢？在寻求大佬的帮助后，指出红黑树进行筛查。或许本人技术有限并没有将红黑树写的很好，只是简单的实现了下。

python 复制代码

class monitor_mouse:
    def __init__(self, img, xyxy_list=None, xyxy_tree=None):
        self.lsPointsChoose = []  # 存入选择的点
        self.tpPointsChoose = []  # 存入选择的点
        self.pointsCount = 0  # 对鼠标按下的点计数(初始化)
        self.count = 0  # 统计量：鼠标按下的点计数(初始化)
        self.pointsMax = 2  # 最多点击多少下
        self.wim_name = 'src'  # 显示图像框名称
        self.img = img  # 转录输入图像
        self.img2 = img.copy()  # 对输入图像进行copy
        self.xyxy_list = xyxy_list  # 根据面积由小到大存储坐标框xyxy格式
        self.xyxy_tree = xyxy_tree  # sam结果存储到红黑树中便于查找

    def on_mouse(self, event, x, y, flags, param):
        # -------------------------左键单击事件---------------------------------------
        if event == cv2.EVENT_LBUTTONDOWN:  # 左键点击
            if self.pointsCount == self.pointsMax:  # 设定了只允许两点成线
                self.clear_info()

            self.pointsCount += 1  # 当左键点击后，记录点击的次数
            cv2.circle(self.img2, (x, y), 2, (255, 0, 0), 2)  # 绘制了当前鼠标选取的点
            # 将选取的点保存到list列表里
            self.lsPointsChoose.append([x, y])  #
            self.tpPointsChoose.append((x, y))  # 用于画点
            # ----------------------------------------------------------------------
            # # 将鼠标选的点用直线连起来
            for i in range(len(self.tpPointsChoose) - 1):
                cv2.line(self.img2, self.tpPointsChoose[i], self.tpPointsChoose[i + 1], (0, 0, 255), 2)

            cv2.imshow(self.wim_name, self.img2)

        # -------------------------右键单击事件-----------------------------
        if event == cv2.EVENT_RBUTTONDOWN:  # 右键点击
            self.axis_in_boxs()  # 输出单击选择的坐标集
            self.clear_info()
            cv2.imshow(self.wim_name, self.img2)

        # -------------------------左键双击事件：清空画布及控制台-----------------------------
        if event == cv2.EVENT_LBUTTONDBLCLK:
            self.img2 = self.img.copy()
            self.clear_info()
            self.clear_win()
            cv2.imshow(self.wim_name, self.img2)

    def clear_info(self):
        self.pointsCount = 0
        self.tpPointsChoose.clear()
        self.lsPointsChoose.clear()
        self.count = 0

    def clear_win(self):
        os.system('cls')

    def axis_in_boxs(self):
        for _, box in self.xyxy_tree.items():
            if is_box_containing_point(box, self.lsPointsChoose):
                cv2.rectangle(self.img2, (box[0], box[1]), (box[2], box[3]), (0, 0, 0), 2)
                print("pt1:", (box[0], box[1]), " ", "pt2:", (box[2], box[3]))
                break

使用过程示意图

请求共同开发

由于该项目笔者已开源在github,希望志同道合之友可以合力完善上述项目。在这里我整理了下待完善的部分：

推理模块可以加速，尽可能的降低资源的消耗；
对其已成熟的标注程序如LabelImg进行二创，将此项目融合或根据其他的标注程序进行二创，能力优秀可以独立创作
~~3. 在1中的加速中可选择将网络推理结果中的xyxy_tree、xyxy_list存储为txt到本地，后期可调用txt即可（已实现）~~

结语

欢迎来到我的开源项目！我非常希望能找到一些志同道合的朋友来一起完成这个项目。这个项目对我来说非常重要，我相信与有相同兴趣和热情的人合作，我们可以把这个项目做得更好。如果你对这个项目感兴趣，欢迎加入我们的开发团队，一起为这个项目做出贡献。无论你是经验丰富的开发者还是新手，只要你有热情和愿意学习，都可以成为我们团队的一部分。

在这里衷心寻求志同道合之友，齐心协力实现上述功能点。顿首再拜，以表万分感激之情！

【mobileSam】使用大模型推理赋能标注工作，让标注工作不再困难

前言