darknetで処理して枠を付けた画像を保存するやつ

darknet-test.py

from ctypes import *
import random
import os
import argparse
import shutil
from pathlib import Path

global DARKNET_FORCE_CPU
DARKNET_FORCE_CPU = True
import cv2
import time
import darknet

def parser():
    parser = argparse.ArgumentParser(description="YOLO Streamer Test Application")

    return parser.parse_args()


if __name__ == '__main__':
    
    network, class_names, class_colors = darknet.load_network(
            config_file="cfg/yolov4.cfg",
            data_file="cfg/coco.data",
            weights="yolov4.weights",
        )

    width = darknet.network_width(network)
    height = darknet.network_height(network)

    darknet_image = darknet.make_image(width, height, 3) 
    # 3 is color channels count https://github.com/AlexeyAB/darknet/blob/master/include/yolo_v2_class.hpp
    
    input_dir = "input"
    output_dir = "output"

    currentPath = Path(".")
    inputPath = currentPath.joinpath(input_dir)
    outputPath = currentPath.joinpath(output_dir)
    cv2.namedWindow("hoge", cv2.WINDOW_NORMAL)
    
    for child in inputPath.iterdir():
        if(child.is_file()):
            print(str(child))
            prev_time = time.time()
            image = cv2.imread(str(child))
            print("imread")

            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            resized_image = cv2.resize(
                rgb_image, (width, height), interpolation=cv2.INTER_LINEAR)
            print("resized")
            
            darknet.copy_image_from_bytes(darknet_image, resized_image.tobytes())
            print("copyed")
            
            detections = darknet.detect_image(
                network, class_names, darknet_image, )
            print("detected")

            prev_image = darknet.draw_boxes(detections, resized_image, class_colors)
            passed_time_ms = ( time.time() - prev_time ) * 1000
            print("processing time: {}ms".format(passed_time_ms))
            cv2.imshow("hoge", prev_image)
            bgr_prev_image = cv2.cvtColor(prev_image, cv2.COLOR_RGB2BGR)
            cv2.imwrite(str(outputPath.joinpath(child.name)), bgr_prev_image)
            #cv2.waitKey(3000)