TORCH09-06:Yolov3与OpenCV

import cv2
import numpy as np
import os
# 配置文件(为了部署方便，采用相对包路径的方式读取配置文件)
current_dir = os.path.dirname(__file__)
configs_file = os.path.join(current_dir, "conf/yolov3.cfg")
weights_file = os.path.join(current_dir, "conf/yolov3.weights")

# 通过加载网络模型创建yolo网络
net = cv2.dnn.readNetFromDarknet(configs_file, weights_file)
# Yolo网络的输出层
ln = net.getUnconnectedOutLayersNames()   # unconnect就是没有下一层链接的层，其实就是输出层

# 准备图像
img = cv2.imread("yolo.jpg")
# 处理为yolo能接受的图像格式（4维:NCHW）
img_darknet = cv2.dnn.blobFromImage(img, 1 / 255.0, (608, 608), swapRB=True, crop=False)

# 设置需要处理的图像
net.setInput(img_darknet)
# c使用yolo网络处理
detect_out = net.forward(ln)    # 输出的数据长度由ln的长度确定
# print(len(detect_out))
# print(detect_out[0].shape)   # 1083, 85（5:中心位置，大小，置信度 + 80：80个目标的概率）  
# 下面三个参数是最大化抑制需要的参数
boxes = []
probs = []
cidxs = []
for layer_out in  detect_out:
    for obj_out in  layer_out:
        # 侦测的目标的中心为主与大小
        box =  obj_out[:4]
        # 属于目标的概率
        confidence = obj_out[4]
        # 目标类别分类
        pred = obj_out[5:]
        # 所属类别
        cls_id = np.argmax(pred)
        # 所属类别的概率
        prob = pred[cls_id]  
        if confidence > 0.0 and prob > 0.0:
            # print(F"{box},{confidence},{cls_id}, {prob}")
            # 对box进行放大
            box = box * np.array([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
            box = box.astype(np.int)
            
            # print(F"{box},{confidence},{cls_id}, {prob}")
            # 为了方便，需要计算矩形的左上角坐标
            c_x, c_y, w, h = box
            x = int(c_x - w / 2)
            y = int(c_y - h / 2)
            boxes.append([x, y, int(w), int(h)])
            probs.append(float(prob))  # prob是np.float32需要转换下
            cidxs.append(cls_id)
            # 在原图像上绘制目标

idxs = cv2.dnn.NMSBoxes(boxes, probs, 0.5, 0.3)

if len(idxs) >0:
    for [idx] in idxs:
        (x, y, w, h)= boxes[idx]
        cv2.rectangle(img, (x, y), (x + w, y + h), (255, 0, 0), 2)
        cv2.putText(img, F"{cidxs[idx]},{probs[idx]:3.2f}", (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

cv2.imwrite("out.jpg", img)