torchvision库中已存在目标检测的相应模型，只需要调用相应的函数即可。
- 单文件：
import numpy as np
import torchvision
import torch
import torchvision.transforms as transforms
from PIL import Image, ImageDraw, ImageFont
import matplotlib.pyplot as plt
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()
image = Image.open("/mnt/COCO2017/val2017/000000000285.jpg")
transform_d = transforms.Compose([transforms.ToTensor()])
image_t = transform_d(image)
pred = model([image_t])
print(pred)
"""
fire hydrant 消防栓，stop sign 停车标志， parking meter 停车收费器， bench 长椅。
zebra 斑马， giraffe 长颈鹿， handbag 手提包， suitcase 手提箱， frisbee （游戏用）飞盘（flying disc）。
skis 滑雪板（ski的复数），snowboard 滑雪板（ski是单板滑雪，snowboarding 是双板滑雪。）
kite 风筝， baseball bat 棒球棍， baseball glove 棒球手套， skateboard 滑板， surfboard 冲浪板， tennis racket 网球拍。
broccoli 西蓝花，donut甜甜圈，炸面圈(doughnut，空心的油炸面包), cake 蛋糕、饼, couch 长沙发（靠chi)。
potted plant 盆栽植物。 dining table 餐桌。 laptop 笔记本电脑，remote 遥控器(=remote control),
cell phone 移动电话(=mobile phone)(cellular 细胞的、蜂窝状的)， oven 烤炉、烤箱。 toaster 烤面包器（toast 烤面包片）
sink 洗碗池, refrigerator 冰箱。（=fridge）， scissor剪刀(see, zer), teddy bear 泰迪熊。 hair drier 吹风机。
toothbrush 牙刷。
"""
COCO_INSTANCE_CATEGORY_NAMES = [
'__BACKGROUND__', 'person', 'bicycle', 'car', 'motorcycle',
'airplane', 'bus', 'train', 'trunk', 'boat', 'traffic light',
'fire hydrant', 'N/A', 'stop sign', 'parking meter', 'bench',
'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant',
'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A',
'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass',
'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A',
'dining table', 'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop',
'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
'toaster', 'toaster', 'sink', 'refrigerator', 'N/A', 'book', 'clock',
'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
pred_class = [COCO_INSTANCE_CATEGORY_NAMES[ii] for ii in list(pred[0]['labels'].numpy())]
pred_score = list(pred[0]['scores'].detach().numpy())
pred_boxes = [[ii[0], ii[1], ii[2], ii[3]] for ii in list(pred[0]['boxes'].detach().numpy())]
pred_index = [pred_score.index(x) for x in pred_score if x > 0.5]
fontsize = np.int16(image.size[1] / 20)
font1 = ImageFont.truetype("/usr/share/fonts/gnu-free/FreeMono.ttf", fontsize)
draw = ImageDraw.Draw(image)
for index in pred_index:
box = pred_boxes[index]
draw.rectangle(box, outline="blue")
texts = pred_class[index]+":"+str(np.round(pred_score[index], 2))
draw.text((box[0], box[1]), texts, fill="blue", font=font1)
plt.imshow(image)
plt.show()
- 对该文件夹下随机10个文件进行目标检测。
"""
对给定文件夹下随机数量的文件进行目标检测。
"""
import numpy as np
import torchvision
import torch
import torchvision.transforms as transforms
from PIL import Image, ImageDraw, ImageFont
import matplotlib.pyplot as plt
import os
"""
fire hydrant 消防栓，stop sign 停车标志， parking meter 停车收费器， bench 长椅。
zebra 斑马， giraffe 长颈鹿， handbag 手提包， suitcase 手提箱， frisbee （游戏用）飞盘（flying disc）。
skis 滑雪板（ski的复数），snowboard 滑雪板（ski是单板滑雪，snowboarding 是双板滑雪。）
kite 风筝， baseball bat 棒球棍， baseball glove 棒球手套， skateboard 滑板， surfboard 冲浪板， tennis racket 网球拍。
broccoli 西蓝花，donut甜甜圈，炸面圈(doughnut，空心的油炸面包), cake 蛋糕、饼, couch 长沙发（靠chi)。
potted plant 盆栽植物。 dining table 餐桌。 laptop 笔记本电脑，remote 遥控器(=remote control),
cell phone 移动电话(=mobile phone)(cellular 细胞的、蜂窝状的)， oven 烤炉、烤箱。 toaster 烤面包器（toast 烤面包片）
sink 洗碗池, refrigerator 冰箱。（=fridge）， scissor剪刀(see, zer), teddy bear 泰迪熊。 hair drier 吹风机。
toothbrush 牙刷。
"""
COCO_INSTANCE_CATEGORY_NAMES = [
'__BACKGROUND__', 'person', 'bicycle', 'car', 'motorcycle',
'airplane', 'bus', 'train', 'trunk', 'boat', 'traffic light',
'fire hydrant', 'N/A', 'stop sign', 'parking meter', 'bench',
'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant',
'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A',
'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass',
'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A',
'dining table', 'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop',
'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
'toaster', 'toaster', 'sink', 'refrigerator', 'N/A', 'book', 'clock',
'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()
image_folder_path = "/mnt/COCO2017/val2017/"
number = 10
transform_d = transforms.Compose([transforms.ToTensor()])
dirs = os.listdir(image_folder_path)
idx = np.random.randint(0, len(dirs), number)
for i in idx:
image_path = os.path.join(image_folder_path, dirs[i])
image = Image.open(image_path)
image_t = transform_d(image)
pred = model([image_t])
print(pred)
pred_class = [COCO_INSTANCE_CATEGORY_NAMES[ii] for ii in list(pred[0]['labels'].numpy())]
pred_score = list(pred[0]['scores'].detach().numpy())
pred_boxes = [[ii[0], ii[1], ii[2], ii[3]] for ii in list(pred[0]['boxes'].detach().numpy())]
pred_index = [pred_score.index(x) for x in pred_score if x > 0.5]
fontsize = np.int16(image.size[1] / 20)
font1 = ImageFont.truetype("/usr/share/fonts/gnu-free/FreeMono.ttf", fontsize)
draw = ImageDraw.Draw(image)
for index in pred_index:
box = pred_boxes[index]
draw.rectangle(box, outline="blue")
texts = pred_class[index]+":"+str(np.round(pred_score[index], 2))
draw.text((box[0], box[1]), texts, fill="blue", font=font1)
plt.imshow(image)
plt.show()
- 我们在COCO2017数据集上进行了一些测试，结果如下：
上图识别出了”人”，概率为100%，识别出了”棒球手套”，概率没瞅见哈哈。
上图识别出了”摩托”，概率100%，”长椅”，概率为74%。
这玩意认为有66%的概率存在”马桶”，87%的概率存在”杯子”，这个识别就有点问题啊。
上面识别出了”披萨（pizza）”，”勺子（spoon）”，”叉子（fork）”。
妈耶，我的Trypophobia都出来了。不过其识别出了 bowl，cup，dinning table，person、spoon，可以调优下哈哈。
下面的图我就不描述了。
识别出了西蓝花(broccoli)，三明治、叉子。
消防栓，fire hydrant
; 4. 参考书籍：
《PyTorch深度学习入门与实践》孙玉林、余本国著。
Original: https://blog.csdn.net/csdnliwenqi/article/details/121686129
Author: 爱学习的大白菜
Title: Pytorch目标检测实现
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/711976/
转载文章受原作者版权保护。转载请注明原作者出处！