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：
