1from ctypes import * 2import math 3import random 4import os 5import cv2 6import numpy as np 7import time 8import darknet 9 10def convertBack(x, y, w, h): 11 xmin = int(round(x - (w / 2))) 12 xmax = int(round(x + (w / 2))) 13 ymin = int(round(y - (h / 2))) 14 ymax = int(round(y + (h / 2))) 15 return xmin, ymin, xmax, ymax 16 17 18def cvDrawBoxes(detections, img): 19 for detection in detections: 20 x, y, w, h = detection[2][0],\ 21 detection[2][1],\ 22 detection[2][2],\ 23 detection[2][3] 24 xmin, ymin, xmax, ymax = convertBack( 25 float(x), float(y), float(w), float(h)) 26 pt1 = (xmin, ymin) 27 pt2 = (xmax, ymax) 28 cv2.rectangle(img, pt1, pt2, (0, 255, 0), 1) 29 cv2.putText(img, 30 detection[0].decode() + 31 " [" + str(round(detection[1] * 100, 2)) + "]", 32 (pt1[0], pt1[1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, 33 [0, 255, 0], 2) 34 return img 35 36 37netMain = None 38metaMain = None 39altNames = None 40 41 42def YOLO(): 43 44 global metaMain, netMain, altNames 45 configPath = "./cfg/yolov4.cfg" 46 weightPath = "./yolov4.weights" 47 metaPath = "./cfg/coco.data" 48 if not os.path.exists(configPath): 49 raise ValueError("Invalid config path `" + 50 os.path.abspath(configPath)+"`") 51 if not os.path.exists(weightPath): 52 raise ValueError("Invalid weight path `" + 53 os.path.abspath(weightPath)+"`") 54 if not os.path.exists(metaPath): 55 raise ValueError("Invalid data file path `" + 56 os.path.abspath(metaPath)+"`") 57 if netMain is None: 58 netMain = darknet.load_net_custom(configPath.encode( 59 "ascii"), weightPath.encode("ascii"), 0, 1) # batch size = 1 60 if metaMain is None: 61 metaMain = darknet.load_meta(metaPath.encode("ascii")) 62 if altNames is None: 63 try: 64 with open(metaPath) as metaFH: 65 metaContents = metaFH.read() 66 import re 67 match = re.search("names *= *(.*)$", metaContents, 68 re.IGNORECASE | re.MULTILINE) 69 if match: 70 result = match.group(1) 71 else: 72 result = None 73 try: 74 if os.path.exists(result): 75 with open(result) as namesFH: 76 namesList = namesFH.read().strip().split("\n") 77 altNames = [x.strip() for x in namesList] 78 except TypeError: 79 pass 80 except Exception: 81 pass 82 #cap = cv2.VideoCapture(0) 83 cap = cv2.VideoCapture("test.mp4") 84 cap.set(3, 1280) 85 cap.set(4, 720) 86 out = cv2.VideoWriter( 87 "output.avi", cv2.VideoWriter_fourcc(*"MJPG"), 10.0, 88 (darknet.network_width(netMain), darknet.network_height(netMain))) 89 print("Starting the YOLO loop...") 90 91 # Create an image we reuse for each detect 92 darknet_image = darknet.make_image(darknet.network_width(netMain), 93 darknet.network_height(netMain),3) 94 while True: 95 prev_time = time.time() 96 ret, frame_read = cap.read() 97 frame_rgb = cv2.cvtColor(frame_read, cv2.COLOR_BGR2RGB) 98 frame_resized = cv2.resize(frame_rgb, 99 (darknet.network_width(netMain), 100 darknet.network_height(netMain)), 101 interpolation=cv2.INTER_LINEAR) 102 103 darknet.copy_image_from_bytes(darknet_image,frame_resized.tobytes()) 104 105 detections = darknet.detect_image(netMain, metaMain, darknet_image, thresh=0.25) 106 image = cvDrawBoxes(detections, frame_resized) 107 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 108 print(1/(time.time()-prev_time)) 109 cv2.imshow('Demo', image) 110 cv2.waitKey(3) 111 cap.release() 112 out.release() 113 114if __name__ == "__main__": 115 YOLO() 116