yolov7推理代码 python版

点击查看代码

import time

import cv2
import numpy as np
from datetime import datetime  # 导入 datetime 模块


##cfg_path是配置文件路径，使用的时候用这个yolov7-tiny.cfg就可以
##weights_path是权重路径，yolov7-tiny_10000.weights是训练了10000轮的权重，效果不行的话也可以使用剩余的两个权重文件试试
##data_path是存放标签名字的路径，classes.names是放标签名字的文件，一行一个标签，使用的时候不要空行
##image_path是待测试的图片的路径，要更换测试图片在这个路径下更换就可以了
# 文件路径
cfg_path= r'./sit-page/yolov7-tiny.cfg'
weights_path= r'./sit-page/yolov7-tiny_10000.weights'
data_path= r'./sit-page/classes.names'
image_path= r'./ng/TRIODE_4_19_20240718165226_OK_.png'
# output_image_path = r'./result/testpic/test_rec/pre.jpg'

# 加载类别名称
with open(data_path, 'r') as f:
    class_names = f.read().strip().split('\n')

# 加载网络
net = cv2.dnn.readNetFromDarknet(cfg_path, weights_path)
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)

# 读取输入图片
image = cv2.imread(image_path)
height, width = image.shape[:2]

start = time.time()
# 创建 blob，并设置为网络的输入
blob = cv2.dnn.blobFromImage(image, scalefactor=1/255.0, size=(256, 256), swapRB=True, crop=False)
net.setInput(blob)

# 获取输出层的名称
layer_names = net.getLayerNames()
output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]

# 前向传递，获得预测结果
layer_outputs = net.forward(output_layers)

# 初始化检测结果
conf_threshold = 0.5
nms_threshold = 0.4
boxes, confidences, class_ids = [], [], []

# 处理检测结果
for output in layer_outputs:
    for detection in output:
        scores = detection[5:]
        class_id = np.argmax(scores)
        confidence = scores[class_id]

        if confidence > conf_threshold:
            box = detection[:4] * np.array([width, height, width, height])
            center_x, center_y, w, h = box.astype('int')

            x = int(center_x - w / 2)
            y = int(center_y - h / 2)

            boxes.append([x, y, int(w), int(h)])
            confidences.append(float(confidence))
            class_ids.append(class_id)

# 应用非极大值抑制（NMS）去除重叠框
indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)
print("indices:",indices)
# 绘制检测框
for i in indices.flatten():
    x, y, w, h = boxes[i]
    print("x:", x)
    print("y:", y)
    print("w:", w)
    print("h:", h)
    color = (0, 255, 0)
    label = f'{class_names[class_ids[i]]}: {confidences[i]:.2f}'

    cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
    cv2.putText(image, label, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

# 生成动态文件名，output_filename是绘制了框图的图像，图像大小没有发生改变，仍然是原图大小
timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
output_filename = f"./test_rec/rgb_image_{timestamp}.png"
end = time.time()

print(f"cost time :{end - start}")
# 保存结果图像
cv2.imwrite(output_filename, image)