使用Python实现yolov5分类识别猫
时间: 2023-05-17 19:03:34 浏览: 113
基于Python实现的猫狗分类识别系统【100011201】
5星 · 资源好评率100%
可以使用以下代码实现:
```
import torch
import cv2
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.torch_utils import select_device
# Load YOLOv5s model
model = attempt_load('yolov5s.pt', map_location=torch.device('cpu'))
# Set device
device = select_device('')
# Define class names
class_names = ['cat']
# Load image
img = cv2.imread('cat.jpg')
# Resize image
img = cv2.resize(img, (640, 640))
# Convert image to tensor
img = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0) / 255.0
# Detect objects
pred = model(img.to(device))[0]
# Apply non-maximum suppression
pred = non_max_suppression(pred, conf_thres=0.5, iou_thres=0.5)
# Draw bounding boxes and labels
for det in pred:
if det is not None and len(det):
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img.shape[2:]).round()
for *xyxy, conf, cls in reversed(det):
label = f'{class_names[int(cls)]} {conf:.2f}'
cv2.rectangle(img, (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3])), (0, 255, 0), 2)
cv2.putText(img, label, (int(xyxy[0]), int(xyxy[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
# Display image
cv2.imshow('image', img)
cv2.waitKey(0)
cv2.destroyAllWindows()
```
请注意,这只是一个简单的示例,实际上,要实现更准确的猫分类识别,需要更多的训练数据和调整模型参数。
阅读全文