YOLOV3-简单的单张图片识别

一、代码段

import numpy as np
import cv2
import matplotlib.pyplot as plt
def look_img(img):
    img_RGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    plt.imshow(img_RGB)
    plt.show()



net = cv2.dnn.readNet('yolov3.weights', 'yolov3.cfg')


with open('coco.names', 'r') as f:
    classes = f.read().splitlines()

img = cv2.imread('jiejing.png')

height , width, _ = img.shape
blob = cv2.dnn.blobFromImage(img, 1/255, (416, 416), (0, 0, 0), swapRB=True, crop=False)
net.setInput(blob)
layersNames = net.getLayerNames()
output_layers_names = [layersNames[i - 1] for i in net.getUnconnectedOutLayers()]

prediction = net.forward(output_layers_names)


boxes=[]
objectness=[]
class_probs = []
class_ids = []
class_names = []

for scale in prediction: #遍历三种尺度
    for bbox in scale:
        obj = bbox[4]
        class_scores = bbox[5:]
        class_id =np.argmax(class_scores)
        class_name = classes[class_id]
        class_prob = class_scores[class_id]

        center_x = int(bbox[0]*width)
        center_y =int(bbox[1]*height)
        w = int(bbox[2]*width)
        h=int(bbox[3]*height)
        x=int(center_x-w/2)
        y=int(center_y-h/2)

        boxes.append([x,y,w,h])
        objectness.append(float(obj))
        class_ids.append((class_id))
        class_names.append(class_name)
        class_probs.append(class_prob)

confidences = np.array(class_probs) * np.array(objectness)
CONF_THRES = 0.4
NMS_THRES = 0.1

indexes = cv2.dnn.NMSBoxes(boxes, confidences, CONF_THRES, NMS_THRES)
colors = np.random.uniform(0, 255, size=(len(boxes),3))

for i in indexes.flatten():
    x, y, w, h = boxes[i]
    confidence = str(round(confidences[i], 2))
    color = colors[i % len(colors)]
    cv2.rectangle(img, (x, y), (x+w, y+h), color, 8)
    string = '{} {}'.format(class_names[i], confidence)
    cv2.putText(img, string, (x, y+20), cv2.FONT_HERSHEY_PLAIN, 4, (255, 255, 255), 4)


look_img(img)

二、实际结果