(Yolov4 테스트) Yolov4 모델이 출력한 검사 상자를 원도에 비추기

(Yolov4 테스트) Yolov4 모델이 출력한 검사 상자를 원도에 비추기
논리:
4

원도를 읽습니다.구성 파일에 따라,.데이터 파일과 권중 파일 가져오기 모델입니다

4

그림, 훈련된 네트워크 모델, 라벨 등을 detection 함수에 함께 보내서 검측 출력 검측 그림과 detections 목록을 검사한다. 목록에는 분류, 신뢰도, 상자의 좌표(x, y, w, h) 등이 포함된다

4

원도와 검측 그림의 사이즈 변환 관계를 계산하는데 이 변환 관계는 좌표 변환에도 적용된다.검출 그림의 상자의 좌표를 변환 관계에 따라 원도에 비추십시오

4

변환된 검측틀의 좌표를 받은 후 원도에 액자를 그립니다

4

보존을 진행한다

import os
import cv2
import darknet

config_file = './train/data/yolov4.cfg'
data_file = "./train/data/train.data"
weights = "./backup/yolov4_last.weights"

def image_detection(image_path, network, class_names, class_colors, thresh):
    # Darknet doesn't accept numpy images.
    # Create one with image we reuse for each detect
    width = darknet.network_width(network)
    height = darknet.network_height(network)
    darknet_image = darknet.make_image(width, height, 3)

    image = cv2.imread(image_path)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image_resized = cv2.resize(image_rgb, (width, height),
                               interpolation=cv2.INTER_LINEAR)

    darknet.copy_image_from_bytes(darknet_image, image_resized.tobytes())
    detections = darknet.detect_image(network, class_names, darknet_image, thresh=thresh)
    darknet.free_image(darknet_image)
    image = darknet.draw_boxes(detections, image_resized, class_colors)
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB), detections


network, class_names, class_colors = darknet.load_network(config_file,
                                                          data_file,
                                                          weights,
                                                          batch_size=1)

def bbox2points(bbox):
    x, y, w, h = bbox
    xmin = int(round(x - (w / 2)))
    xmax = int(round(x + (w / 2)))
    ymin = int(round(y - (h / 2)))
    ymax = int(round(y + (h / 2)))
    return xmin, ymin, xmax, ymax

colors = {
     'truck': (0, 0, 255), 'car': (0, 255, 0), 'people': (255, 0, 0)}
save_path = "save"
images_path = "test_demo1"
img_names = os.listdir(images_path)
for img_name in img_names:
    img = os.path.join(images_path, img_name)
    image, detections = image_detection(img, network, class_names, class_colors, thresh=0.25)

    im = cv2.imread(img)
    w = im.shape[1]
    h = im.shape[0]

    for label, confidence, bbox in detections:
        xmin, ymin, xmax, ymax = bbox2points(bbox)
        xmin = int(float((xmin * w) / 608))
        xmax = int(float((xmax * w) / 608))
        ymin = int(float((ymin * h) / 608))
        ymax = int(float((ymax * h) / 608))

        cv2.rectangle(im, (xmin, ymax), (xmax, ymin), colors[label], 1)
        cv2.putText(im, "{} [{:.2f}]".format(label, float(confidence)),
                    (xmin, ymax - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                    colors[label], 2)

    cv2.imwrite(os.path.join(save_path, img_name + '__.jpg'), im)