当前位置:网站首页>基于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
边栏推荐
- 云上有AI,让地球科学研究更省力
- 机器学习植物叶片识别
- Leetcode daily question (971. flip binary tree to match preorder traversal)
- [ 英語 ] 語法重塑 之 動詞分類 —— 英語兔學習筆記(2)
- ECS accessKey key disclosure and utilization
- Use shortcut LNK online CS
- Number of query fields
- Making interactive page of "left tree and right table" based on jeecg-boot
- [English] Grammar remodeling: the core framework of English Learning -- English rabbit learning notes (1)
- Financial German translation, a professional translation company in Beijing
猜你喜欢
端午节快乐Wish Dragon Boat Festival is happy
机器学习植物叶片识别
Changes in the number of words in English papers translated into Chinese
专业论文翻译,英文摘要如何写比较好
Market segmentation of supermarket customers based on purchase behavior data (RFM model)
Biomedical English contract translation, characteristics of Vocabulary Translation
Financial German translation, a professional translation company in Beijing
[English] Verb Classification of grammatical reconstruction -- English rabbit learning notes (2)
[mqtt from getting started to improving series | 01] quickly build an mqtt test environment from 0 to 1
Is it difficult for girls to learn software testing? The threshold for entry is low, and learning is relatively simple
随机推荐
Automated test environment configuration
Cobalt strike feature modification
Thesis abstract translation, multilingual pure human translation
Lecture 8: 1602 LCD (Guo Tianxiang)
雲上有AI,讓地球科學研究更省力
How much is it to translate Chinese into English for one minute?
Today's summer solstice
LeetCode每日一题(971. Flip Binary Tree To Match Preorder Traversal)
[English] Grammar remodeling: the core framework of English Learning -- English rabbit learning notes (1)
Luogu p2141 abacus mental arithmetic test
LeetCode 731. My schedule II
Bitcoinwin (BCW): 借贷平台Celsius隐瞒亏损3.5万枚ETH 或资不抵债
mysql的基础命令
Esp32 esp-idf watchdog twdt
基於JEECG-BOOT的list頁面的地址欄參數傳遞
SQL Server Manager studio (SSMS) installation tutorial
Tms320c665x + Xilinx artix7 DSP + FPGA high speed core board
Cobalt Strike特征修改
Day 246/300 ssh连接提示“REMOTE HOST IDENTIFICATION HAS CHANGED! ”
Wish Dragon Boat Festival is happy