torchvision内置模型maskrcnn

maskrcnn_resnet50.py：

import sys

from PIL import Image
import torch
import torchvision
import torchvision.transforms as T
import cv2
import numpy as np
import random
import matplotlib.pyplot as plt

COCO_INSTANCE_CATEGORY_NAMES = ["Person", "Bicycle", "Car", "Motorcycle", "Airplane", "Bus",
                                "Train", "Truck", "Boat", "Traffic light", "Fire hydrant", "Stop sign", "Parking meter",
                                "Bench", "Bird", "Cat", "Dog", "Horse", "Sheep", "Cow", "Elephant", "Bear", "Zebra",
                                "Giraffe", "Backpack", "Umbrella", "Handbag", "Tie", "Suitcase", "Frisbee", "Skis",
                                "Snowboard", "Sports ball", "Kite", "Baseball bat", "Baseball glove", "Skateboard",
                                "Surfboard", "Tennis racket", "Bottle", "Wine glass", "Cup", "Fork", "Knife", "Spoon",
                                "Bowl", "Banana", "Apple", "Sandwich", "Orange", "Broccoli", "Carrot", "Hot dog",
                                "Pizza", "Donut", "Cake", "Chair", "Couch", "Potted plant", "Bed", "Dining table",
                                "Toilet", "TV", "Laptop", "Mouse", "Remote", "Keyboard", "Cell phone", "Microwave",
                                "Oven",
                                "Toaster", "Sink", "Refrigerator", "Book", "Clock", "Vase", "Scissors", "Teddy bear",
                                "Hair drier", "Toothbrush"]

COCO_INSTANCE_CATEGORY_NAMES = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
                                'bus', 'train', 'truck', 'boat', 'traffic light',
                                'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
                                'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
                                'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
                                'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
                                'kite', 'baseball bat', 'baseball glove', 'skateboard',
                                'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
                                'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
                                'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
                                'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
                                'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
                                'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
                                'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
                                'teddy bear', 'hair drier', 'toothbrush']

model = torchvision.models.detection.maskrcnn_resnet50_fpn()
model.load_state_dict(torch.load("./maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth"))
model.eval()
print(model)


# 定义模型，加载权重，并根据阈值过滤结果
def getPrediction(img_path, threshold):
    transform = T.Compose([T.Resize((256, 256)), T.CenterCrop(224),
                           T.ToTensor()])  # T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    transform = T.Compose([T.ToTensor()])
    img = Image.open(img_path)
    img = transform(img).unsqueeze(0)

    pred = model(img)[0]
    pred_boxes = pred['boxes'].detach().cpu().numpy()  # 边框信息
    pred_labels = pred['labels'].detach().cpu().numpy()  # 分类信息
    pred_scores = list(pred['scores'].detach().cpu().numpy())  # 分类置信度
    pred_masks = (pred['masks'] > 0.5).squeeze().detach().cpu().numpy()  # 像素掩码
    pred_class = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in pred_labels]
    pred_t = [pred_scores.index(x) for x in pred_scores if x > threshold]
    if len(pred_t) == 0:
        return None, None, None
    pred_t = pred_t[-1]

    out_masks = pred_masks[:pred_t + 1]
    out_boxes = pred_boxes[:pred_t + 1]
    out_class = pred_class[:pred_t + 1]

    return out_masks, out_boxes, out_class


def random_colour_masks(image, mask):
    # 获取掩码的形状
    height, width = mask.shape
    # 创建一个与原图像尺寸相同的黑色背景
    colored_image = np.zeros_like(image)
    # 为每个实例生成一个随机颜色 (RGB)
    color = [random.randint(0, 255) for _ in range(3)]
    # 为该实例的掩码位置填充颜色
    colored_image[mask == 1] = color
    return colored_image, color


if __name__ == "__main__":
    # img_pth = r'./dog.2.jpg'
    img_pth = r'./b42ff7ee-00000000.jpg'
    # img_pth = r'./cat.3.jpg'
    masks, boxes, classes = getPrediction(img_pth, 0.5)
    if masks is None:
        print("No masks")
        sys.exit()

    img_show = cv2.imread(img_pth)
    img_show = cv2.cvtColor(img_show, cv2.COLOR_BGR2RGB)
    for i in range(len(boxes)):
        colored_image, color = random_colour_masks(img_show, masks[i])
        img_show = cv2.addWeighted(img_show, 1, colored_image, 0.4, 0)
        pt1 = (int(boxes[i][0]), int(boxes[i][1]))
        pt2 = (int(boxes[i][2]), int(boxes[i][3]))
        cv2.rectangle(img_show, pt1, pt2, tuple(color), 2)
        cv2.putText(img_show, classes[i], (int((pt1[0] + pt2[0])/2), int((pt1[1] + pt2[1])/2)), cv2.FONT_HERSHEY_SIMPLEX, 2, color, 2)

    plt.figure()
    plt.imshow(img_show)
    plt.xticks([])
    plt.yticks([])
    plt.show()
    img_save = cv2.cvtColor(img_show, cv2.COLOR_RGB2BGR)
    cv2.imwrite('result_b42ff7ee-00000000.jpg', img_save)
    print('Done.')

原图b42ff7ee-00000000.jpg

结果图result_b42ff7ee-00000000.jpg：