(Yolov 4テスト)Yolov 4モデル出力の検出ボックスを原図にマッピング

16041 ワード

(Yolov 4テスト)Yolov 4モデル出力の検出ボックスを原図にマッピング


論理:
  • 原図を読みだします.プロファイルに基づいて、.Dataファイルとウェイトファイルをモデルにインポートします.
  • ピクチャ、トレーニングされたネットワークモデル、ラベルなどをdetection関数に一緒に送り込んで検出出力検出ピクチャとdetectionsリストを行い、リストにはカテゴリ、信頼度、およびボックスの座標(x,y,w,h)などが含まれる.
  • 原図と検出ピクチャの寸法変換関係を計算し、この変換関係は座標変換にも適用される.検出画像上のボックスの座標を変換関係に基づいて原図にマッピングします.
  • 変換後の検出枠の座標を得た後、原図に枠を描く.
  • 保存します.
  • 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)