当前位置:网站首页>基于PyTorch和Fast RCNN快速实现目标识别
基于PyTorch和Fast RCNN快速实现目标识别
2022-07-06 06:35:00 【GIS开发者】
Faster RCNN,相对于R-CNN在结构上,Faster RCNN已经将特征抽取(feature extraction),proposal提取,bounding box regression(rect refine),classification都整合在了一个网络中,使得综合性能有较大提高,在检测速度方面尤为明显。
PyTorch是一个开源的Python机器学习库,基于Torch,用于自然语言处理等应用程序。相对于TensorFlow更加轻量化,更适合于科研和小型项目的使用。
这里介绍一个简单的例子,基于Fast RCNN算法和PyTorch快速实现目标识别。这里使用的是coco数据集已经训练好的在线的模型,运行起来比较简单。
代码
from PIL import Image
import matplotlib.pyplot as plt
# pip install -U matplotlib
import torch
# pip install pytorch
import torchvision.transforms as T
import torchvision
# pip install torchvision
import numpy as np
import cv2
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
# pip install opencv-python
# 下载已经训练好的模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()
COCO_INSTANCE_CATEGORY_NAMES = [
'__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
'train', 'truck', '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', 'sink', 'refrigerator', 'N/A', 'book',
'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
def get_prediction(img_path, threshold):
img = Image.open(img_path)
# 转换一个PIL库的图片或者numpy的数组为tensor张量类型;转换从[0,255]->[0,1]
transform = T.Compose([T.ToTensor()])
img = transform(img)
pred = model([img])
print(pred[0]['labels'].numpy())
# 类别提取
pred_class = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in list(pred[0]['labels'].numpy())]
# 坐标提取
pred_boxes = [[(i[0], i[1]), (i[2], i[3])] for i in list(pred[0]['boxes'].detach().numpy())]
# 找出符合相似度要求的
pred_score = list(pred[0]['scores'].detach().numpy())
pred_t = [pred_score.index(x) for x in pred_score if x > threshold][-1]
pred_boxes = pred_boxes[:pred_t + 1]
pred_class = pred_class[:pred_t + 1]
print("pred_class:", pred_class)
print("pred_boxes:", pred_boxes)
return pred_boxes, pred_class
def object_detection_api(img_path, threshold=0.5, rect_th=3, text_size=3, text_th=3):
boxes, pred_cls = get_prediction(img_path, threshold)
img = cv2.imread(img_path)
# 转换为RGB图像
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
for i in range(len(boxes)):
# 根据坐标圈出目标
cv2.rectangle(img, (int(boxes[i][0][0]), int(boxes[i][0][1])), (int(boxes[i][1][0]), int(boxes[i][1][1])),
color=(0, 255, 0),
thickness=rect_th)
# 标注类别
cv2.putText(img, pred_cls[i], (int(boxes[i][0][0]), int(boxes[i][0][1])), cv2.FONT_HERSHEY_SIMPLEX, text_size,
(0, 255, 0), thickness=text_th)
plt.imshow(img)
plt.show()
if __name__ == '__main__':
object_detection_api(img_path=r"C:\Users\hanbo\Pictures\dog.jpg")
结果
示例1
示例2
边栏推荐
- Luogu p2141 abacus mental arithmetic test
- Oscp raven2 target penetration process
- Lecture 8: 1602 LCD (Guo Tianxiang)
- Day 248/300 关于毕业生如何找工作的思考
- P5706 [deep foundation 2. Example 8] redistributing fat house water -- February 13, 2022
- The internationalization of domestic games is inseparable from professional translation companies
- 成功解决TypeError: data type ‘category‘ not understood
- 【软件测试进阶第1步】自动化测试基础知识
- Modify the list page on the basis of jeecg boot code generation (combined with customized components)
- Attributeerror successfully resolved: can only use cat accessor with a ‘category‘ dtype
猜你喜欢
[brush questions] how can we correctly meet the interview?
女生学软件测试难不难 入门门槛低,学起来还是比较简单的
Office doc add in - Online CS
26岁从财务转行软件测试,4年沉淀我已经是25k的测开工程师...
[English] Verb Classification of grammatical reconstruction -- English rabbit learning notes (2)
How effective is the Chinese-English translation of international economic and trade contracts
LeetCode 732. My schedule III
Apple has open source, but what about it?
翻译影视剧字幕,这些特点务必要了解
How to do a good job in financial literature translation?
随机推荐
What are the characteristics of trademark translation and how to translate it?
[Yu Yue education] flower cultivation reference materials of Weifang Vocational College
自动化测试环境配置
LeetCode 731. My schedule II
How to translate biomedical instructions in English
Cobalt strike feature modification
Leetcode daily question (971. flip binary tree to match preorder traversal)
今日夏至 Today‘s summer solstice
Use shortcut LNK online CS
机器学习植物叶片识别
同事上了个厕所,我帮产品妹子轻松完成BI数据产品顺便得到奶茶奖励
Biomedical English contract translation, characteristics of Vocabulary Translation
[ 英語 ] 語法重塑 之 動詞分類 —— 英語兔學習筆記(2)
Biomedical localization translation services
The ECU of 21 Audi q5l 45tfsi brushes is upgraded to master special adjustment, and the horsepower is safely and stably increased to 305 horsepower
CS通过(CDN+证书)powershell上线详细版
SAP SD发货流程中托盘的管理
Drug disease association prediction based on multi-scale heterogeneous network topology information and multiple attributes
生物医学英文合同翻译,关于词汇翻译的特点
CS-证书指纹修改