SAM分割一切模型下游任务测试：基于检测框的目标分割

import argparse
import cv2
import numpy as np
import  matplotlib.pyplot as plt
from segment_anything import sam_model_registry, SamPredictor
import time
from v5lite import yolov5_lite


def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)


def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))    




if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--detectorModel', type=str, default='best_person.onnx', help="onnx filepath")
    parser.add_argument('--classFile', type=str, default='coco.names', help="classname filepath")
    parser.add_argument('--confThreshold', default=0.5, type=float, help='class confidence')
    parser.add_argument('--nmsThreshold', default=0.6, type=float, help='nms iou thresh')
    parser.add_argument('--sam_checkpoint', type = str, default="sam_vit_b_01ec64.pth",  help='sam checkpoint')
    parser.add_argument('--sam_model_type', type = str, default="vit_b",  help='sam model type')
    parser.add_argument("--device", default="cuda", help="device")

    args = parser.parse_args()


    # 初始化目标检测器
    net = yolov5_lite(args.detectorModel, args.classFile, confThreshold=args.confThreshold, nmsThreshold=args.nmsThreshold)

    # 初始化SAM
    sam = sam_model_registry[args.sam_model_type](checkpoint=args.sam_checkpoint)
    sam.to(device=args.device)


    # 读取视频
    # cap = cv2.VideoCapture('D:/012-PersonDetection/100try.mp4')


    predictor = SamPredictor(sam)


    # start_time = time.time()

    image = cv2.imread('00000.jpg')
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)


  

    srcimg, m_roi = net.detect(image.copy())



    print(m_roi[0])

    start_time = time.time()

    predictor.set_image(image)


    input_box = np.array(m_roi[0])

    masks, _, _ = predictor.predict(
        point_coords=None,
        point_labels=None,
        box=input_box[None, :],
        multimask_output=False,
    )


    det_time = time.time() - start_time
    print("time:{:.3f}ms".format(det_time*1e3))


    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    show_mask(masks[0], plt.gca())
    show_box(input_box, plt.gca())
    plt.axis('off')
    plt.show()