Yolov5-detect.py代码简化(便于移植)
# -*- coding: UTF-8 -*-
# @Time : 2023/11/1 18:23
# @File : detect_iter.py
# @Software: PyCharm
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
import argparse
import os
import sys
import time
from pathlib import Path
import cv2
import numpy as np
import torch
torch.cuda.current_device()
import torch.backends.cudnn as cudnn
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # YOLOv5 root directory
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
from models.experimental import attempt_load
from utils.datasets import LoadImages, LoadStreams
from utils.general import apply_classifier, check_img_size, check_imshow, check_requirements, check_suffix, colorstr, \
increment_path, non_max_suppression, print_args, save_one_box, scale_coords, set_logging, \
strip_optimizer, xyxy2xywh
from utils.plots import Annotator, colors, plot_one_box
from utils.torch_utils import load_classifier, select_device, time_sync
from utils.augmentations import letterbox
@torch.no_grad()
def run():
# Initialize
weights = './yolov5s.pt' # model.pt path(s)
device = 'cuda:0' # cuda device, i.e. 0 or 0,1,2,3 or cpu
save_conf = False
imgsz = 640
line_thickness = 3 # bounding box thickness (pixels)
hide_labels = False # hide labels
hide_conf = False # hide confidences
half = False
device = select_device(device)
half &= device.type != 'cpu' # half precision only supported on CUDA
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
stride = 32 # model stride
names = model.module.names if hasattr(model, 'module') else model.names # get class names
if half:
model.half() # to FP16
# 导入图片
img = cv2.imread("./data/images/image2.jpg")
im0 = img.copy()
# 处理图片
img = letterbox(img, new_shape=(imgsz, imgsz), stride=stride)[0]
img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
img = np.ascontiguousarray(img)
# 个数统计
statistic_dic = {name: 0 for name in names}
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img = img / 255.0 # 0 - 255 to 0.0 - 1.0
if len(img.shape) == 3:
img = img[None] # expand for batch dim
pred = model(img, augment=False, visualize=False)[0]
# classes决定检测类别
pred = non_max_suppression(pred, conf_thres=0.6, iou_thres=0.45, classes=None, max_det=1000)
# Process predictions
for i, det in enumerate(pred): # per image
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# Write results
for *xyxy, conf, cls in reversed(det):
c = int(cls)
statistic_dic[names[c]] += 1
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4))).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
plot_one_box(xyxy, im0, label=label, color=colors(c, True), line_thickness=line_thickness)
print(statistic_dic)
cv2.imshow("img", im0)
cv2.waitKey(0)
cv2.destroyAllWindows()
if __name__ == "__main__":
run()
文章来源地址https://www.toymoban.com/news/detail-737216.html
文章来源:https://www.toymoban.com/news/detail-737216.html
到了这里,关于Yolov5-detect.py代码简化(便于移植)的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!