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