当前位置: 代码迷 >> 综合 >> Python实现yolov4(tiny)(摄像头实现)
  详细解决方案

Python实现yolov4(tiny)(摄像头实现)

热度:88   发布时间:2023-11-01 00:13:47.0

环境要求:基本的Pytorch及相关模块

模型文件:见https://blog.csdn.net/qq_37249793/article/details/115467880

代码:

# -*- coding: utf-8 -*-
"""
Created on Sat Mar 13 10:24:50 2021@author: S
"""import numpy as np
import cv2
import os
import timedef video_demo():# 加载已经训练好的模型路径,可以是绝对路径或者相对路径weightsPath = "yolov4.weights"configPath = "yolov4.cfg"labelsPath = "coco.names"# 初始化一些参数LABELS = open(labelsPath).read().strip().split("\n")  # 物体类别COLORS = np.random.randint(0, 255, size=(len(LABELS), 3), dtype="uint8")  # 颜色boxes = []confidences = []classIDs = []net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)# 读入待检测的图像# 0是代表摄像头编号,只有一个的话默认为0capture = cv2.VideoCapture(0)while (True):ref, image = capture.read()(H, W) = image.shape[:2]# 得到 YOLO需要的输出层ln = net.getLayerNames()ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]# 从输入图像构造一个blob,然后通过加载的模型,给我们提供边界框和相关概率blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (416, 416), swapRB=True, crop=False)net.setInput(blob)layerOutputs = net.forward(ln)# 在每层输出上循环for output in layerOutputs:# 对每个检测进行循环for detection in output:scores = detection[5:]classID = np.argmax(scores)confidence = scores[classID]# 过滤掉那些置信度较小的检测结果if confidence > 0.5:# 框后接框的宽度和高度box = detection[0:4] * np.array([W, H, W, H])(centerX, centerY, width, height) = box.astype("int")# 边框的左上角x = int(centerX - (width / 2))y = int(centerY - (height / 2))# 更新检测出来的框boxes.append([x, y, int(width), int(height)])confidences.append(float(confidence))classIDs.append(classID)# 极大值抑制idxs = cv2.dnn.NMSBoxes(boxes, confidences, 0.2, 0.3)if len(idxs) > 0:for i in idxs.flatten():(x, y) = (boxes[i][0], boxes[i][1])(w, h) = (boxes[i][2], boxes[i][3])# 在原图上绘制边框和类别color = [int(c) for c in COLORS[classIDs[i]]]cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)text = "{}: {:.4f}".format(LABELS[classIDs[i]], confidences[i])cv2.putText(image, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)cv2.imshow("Image", image)# 等待30ms显示图像,若过程中按“ESC”退出c = cv2.waitKey(30) & 0xffif c == 27:capture.release()breakvideo_demo()