├── yolov4
    ├── __init__.py
    ├── helpers.py
    └── darknet.py
├── misc
    └── sample1.png
├── setup.py
├── README.md
└── test.py


/yolov4/__init__.py:
--------------------------------------------------------------------------------
1 | from yolov4.darknet import Detector, MultiGPU


--------------------------------------------------------------------------------
/misc/sample1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/philipperemy/python-darknet-yolo-v4/HEAD/misc/sample1.png


--------------------------------------------------------------------------------
/setup.py:
--------------------------------------------------------------------------------
 1 | from setuptools import setup, find_packages
 2 | 
 3 | setup(
 4 |     name='yolo-v4',
 5 |     version='0.5',
 6 |     author='Philippe Remy',
 7 |     description='Interface Darknet YOLOv4 with python',
 8 |     include_package_data=True,
 9 |     data_files=[
10 |     ],
11 |     install_requires=[
12 |         'numpy',
13 |         'opencv-python',
14 |         'scikit-image',
15 |         'attrs',
16 |         'Pillow'
17 |     ],
18 |     packages=find_packages(),
19 | )
20 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | ## YOLOv4 in Python
 2 | 
 3 | Python interface to Darknet Yolo V4. The multi GPU is supported (load balancer).
 4 | 
 5 | ### Installation
 6 | 
 7 | Compile the Darknet framework first.
 8 | 
 9 | ```bash
10 | sudo apt-get update 
11 | sudo apt-get install -y pkg-config git build-essential libopencv-dev wget cmake
12 | git clone https://github.com/AlexeyAB/darknet.git
13 | cd darknet
14 | make LIBSO=1 OPENCV=1 GPU=1 AVX=1 OPENMP=1 CUDNN=1 CUDNN_HALF=1 OPENMP=1 -j $(nproc)
15 | chmod +x darknet
16 | ```
17 | 
18 | Then, download the weights by following the instructions here: https://github.com/AlexeyAB/darknet.
19 | 
20 | From there, create a virtual environment with python3.6+ and run this command:
21 | 
22 | ```bash
23 | pip install yolo-v4
24 | ```
25 | 
26 | ### Run inference on images
27 | 
28 | To run inference on the GPU on an image `data/dog.jpg`, run this script:
29 | 
30 | ```python
31 | import numpy as np
32 | from PIL import Image
33 | 
34 | from yolov4 import Detector
35 | 
36 | img = Image.open('data/dog.jpg')
37 | d = Detector(gpu_id=0)
38 | img_arr = np.array(img.resize((d.network_width(), d.network_height())))
39 | detections = d.perform_detect(image_path_or_buf=img_arr, show_image=False)
40 | for detection in detections:
41 |     box = detection.left_x, detection.top_y, detection.width, detection.height
42 |     print(f'{detection.class_name.ljust(10)} | {detection.class_confidence * 100:.1f} % | {box}')
43 | ```
44 | 
45 | ```c
46 | dog          | 97.6 % | (100, 236, 147, 334)
47 | truck        | 93.0 % | (367, 81, 175, 98)
48 | bicycle      | 92.0 % | (90, 134, 362, 315)
49 | pottedplant  | 34.1 % | (538, 115, 29, 47)
50 | ```
51 | 
52 | <p align="center">
53 |   <img src="misc/sample1.png" width="500">
54 | </p>
55 | 
56 | 


--------------------------------------------------------------------------------
/test.py:
--------------------------------------------------------------------------------
 1 | from threading import Thread
 2 | from time import time
 3 | 
 4 | import cv2
 5 | import numpy as np
 6 | from PIL import Image
 7 | from tqdm import tqdm
 8 | 
 9 | from yolov4 import Detector, MultiGPU
10 | 
11 | c = 0
12 | 
13 | 
14 | def target(g, desc):
15 |     global c
16 |     for _ in range(1000):
17 |         # print(desc)
18 |         g.perform_detect(show_image=False)
19 |         if c % 100 == 0:
20 |             print(time(), c)
21 |         c += 1
22 |     print(time(), c)
23 | 
24 | 
25 | def batch_single():
26 |     img_samples = ['data/person.jpg', 'data/dog.jpg', 'data/person.jpg',
27 |                    'data/person.jpg', 'data/dog.jpg', 'data/person.jpg']
28 |     bs = len(img_samples)
29 |     image_list = [cv2.imread(k) for k in img_samples]
30 |     img_list = []
31 |     d = Detector(gpu_id=0, lib_darknet_path='lib/libdarknet.so', batch_size=bs)
32 |     for custom_image_bgr in image_list:
33 |         custom_image = cv2.cvtColor(custom_image_bgr, cv2.COLOR_BGR2RGB)
34 |         custom_image = cv2.resize(
35 |             custom_image, (d.network_width(), d.network_height()), interpolation=cv2.INTER_NEAREST)
36 |         custom_image = custom_image.transpose(2, 0, 1)
37 |         img_list.append(custom_image)
38 |     for _ in tqdm(range(900)):
39 |         d.perform_batch_detect(img_list, batch_size=bs)
40 | 
41 | 
42 | def main_single():
43 |     img = Image.open('data/dog.jpg')
44 |     d = Detector(gpu_id=0, lib_darknet_path='lib/libdarknet.so')
45 |     img_arr = np.array(img.resize((d.network_width(), d.network_height())))
46 |     for _ in tqdm(range(900)):
47 |         d.perform_detect(image_path_or_buf=img_arr, show_image=False)
48 | 
49 | 
50 | def main():
51 |     g = MultiGPU([
52 |         Detector(gpu_id=0, lib_darknet_path='lib/libdarknet.so'),
53 |         Detector(gpu_id=1, lib_darknet_path='lib/libdarknet.so')
54 |     ])
55 | 
56 |     thread1 = Thread(target=target, args=(g, 't1'))
57 |     thread2 = Thread(target=target, args=(g, 't2'))
58 | 
59 |     thread1.start()
60 |     thread2.start()
61 | 
62 |     thread1.join()
63 |     thread2.join()
64 | 
65 | 
66 | if __name__ == '__main__':
67 |     batch_single()
68 | 


--------------------------------------------------------------------------------
/yolov4/helpers.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | from ctypes import *
  3 | from pathlib import Path
  4 | 
  5 | import attr
  6 | import numpy as np
  7 | 
  8 | 
  9 | @attr.s
 10 | class DarkNetPredictionResult:
 11 |     class_name = attr.ib(type=str)  # name of the class detected. E.g. dog.
 12 |     class_confidence = attr.ib(type=float)  # probability of the class. E.g. 95%.
 13 |     left_x = attr.ib(type=int)  # box center x.
 14 |     top_y = attr.ib(type=int)  # box center y.
 15 |     width = attr.ib(type=int)  # box width.
 16 |     height = attr.ib(type=int)  # box height.
 17 |     info = attr.ib(type=dict, default={})  # extra info. Can be blank.
 18 | 
 19 | 
 20 | def c_array(ctype, values):
 21 |     arr = (ctype * len(values))()
 22 |     arr[:] = values
 23 |     return arr
 24 | 
 25 | 
 26 | class BOX(Structure):
 27 |     _fields_ = [("x", c_float),
 28 |                 ("y", c_float),
 29 |                 ("w", c_float),
 30 |                 ("h", c_float)]
 31 | 
 32 | 
 33 | class DETECTION(Structure):
 34 |     _fields_ = [("bbox", BOX),
 35 |                 ("classes", c_int),
 36 |                 ("prob", POINTER(c_float)),
 37 |                 ("mask", POINTER(c_float)),
 38 |                 ("objectness", c_float),
 39 |                 ("sort_class", c_int),
 40 |                 ("uc", POINTER(c_float)),
 41 |                 ("points", c_int),
 42 |                 ("embeddings", POINTER(c_float)),
 43 |                 ("embedding_size", c_int),
 44 |                 ("sim", c_float),
 45 |                 ("track_id", c_int)]
 46 | 
 47 | 
 48 | class DETNUMPAIR(Structure):
 49 |     _fields_ = [("num", c_int),
 50 |                 ("dets", POINTER(DETECTION))]
 51 | 
 52 | 
 53 | class IMAGE(Structure):
 54 |     _fields_ = [("w", c_int),
 55 |                 ("h", c_int),
 56 |                 ("c", c_int),
 57 |                 ("data", POINTER(c_float))]
 58 | 
 59 | 
 60 | class METADATA(Structure):
 61 |     _fields_ = [("classes", c_int),
 62 |                 ("names", POINTER(c_char_p))]
 63 | 
 64 | 
 65 | def array_to_image(arr):
 66 |     # need to return old values to avoid python freeing memory
 67 |     arr = arr.transpose(2, 0, 1)
 68 |     c = arr.shape[0]
 69 |     h = arr.shape[1]
 70 |     w = arr.shape[2]
 71 |     arr = np.ascontiguousarray(arr.flat, dtype=np.float32) / 255.0
 72 |     data = arr.ctypes.data_as(POINTER(c_float))
 73 |     im = IMAGE(w, h, c, data)
 74 |     return im, arr
 75 | 
 76 | 
 77 | def init_lib(lib_darknet_path):
 78 |     # lib = CDLL("/home/pjreddie/documents/darknet/libdarknet.so", RTLD_GLOBAL)
 79 |     # lib = CDLL("libdarknet.so", RTLD_GLOBAL)
 80 |     hasGPU = True
 81 |     if os.name == "nt":
 82 |         cwd = os.path.dirname(__file__)
 83 |         os.environ['PATH'] = cwd + ';' + os.environ['PATH']
 84 |         winGPUdll = os.path.join(cwd, "yolo_cpp_dll.dll")
 85 |         winNoGPUdll = os.path.join(cwd, "yolo_cpp_dll_nogpu.dll")
 86 |         envKeys = list()
 87 |         for k, v in os.environ.items():
 88 |             envKeys.append(k)
 89 |         try:
 90 |             try:
 91 |                 tmp = os.environ["FORCE_CPU"].lower()
 92 |                 if tmp in ["1", "true", "yes", "on"]:
 93 |                     raise ValueError("ForceCPU")
 94 |                 else:
 95 |                     print("Flag value '" + tmp + "' not forcing CPU mode")
 96 |             except KeyError:
 97 |                 # We never set the flag
 98 |                 if 'CUDA_VISIBLE_DEVICES' in envKeys:
 99 |                     if int(os.environ['CUDA_VISIBLE_DEVICES']) < 0:
100 |                         raise ValueError("ForceCPU")
101 |                 try:
102 |                     global DARKNET_FORCE_CPU
103 |                     if DARKNET_FORCE_CPU:
104 |                         raise ValueError("ForceCPU")
105 |                 except NameError:
106 |                     pass
107 |                 # print(os.environ.keys())
108 |                 # print("FORCE_CPU flag undefined, proceeding with GPU")
109 |             if not os.path.exists(winGPUdll):
110 |                 raise ValueError("NoDLL")
111 |             lib = CDLL(winGPUdll, RTLD_GLOBAL)
112 |         except (KeyError, ValueError):
113 |             hasGPU = False
114 |             if os.path.exists(winNoGPUdll):
115 |                 lib = CDLL(winNoGPUdll, RTLD_GLOBAL)
116 |                 print("Notice: CPU-only mode")
117 |             else:
118 |                 # Try the other way, in case no_gpu was
119 |                 # compile but not renamed
120 |                 lib = CDLL(winGPUdll, RTLD_GLOBAL)
121 |                 print(
122 |                     "Environment variables indicated a CPU run, but we didn't find `" + winNoGPUdll + "`. Trying a GPU run anyway.")
123 |     else:
124 |         lib = CDLL(Path(lib_darknet_path).resolve(), RTLD_GLOBAL)
125 | 
126 |     lib.network_width.argtypes = [c_void_p]
127 |     lib.network_width.restype = c_int
128 |     lib.network_height.argtypes = [c_void_p]
129 |     lib.network_height.restype = c_int
130 |     copy_image_from_bytes = lib.copy_image_from_bytes
131 |     copy_image_from_bytes.argtypes = [IMAGE, c_char_p]
132 |     return lib, hasGPU
133 | 
134 | 
135 | def read_alt_names(meta_path: str):
136 |     try:
137 |         with open(meta_path) as metaFH:
138 |             metaContents = metaFH.read()
139 |             import re
140 |             match = re.search("names *= *(.*)$", metaContents, re.IGNORECASE | re.MULTILINE)
141 |             if match:
142 |                 result = match.group(1)
143 |             else:
144 |                 result = None
145 |             try:
146 |                 if os.path.exists(result):
147 |                     with open(result) as namesFH:
148 |                         namesList = namesFH.read().strip().split("\n")
149 |                         altNames = [x.strip() for x in namesList]
150 |                         return altNames
151 |             except TypeError:
152 |                 pass
153 |     except Exception:
154 |         pass
155 |     return None
156 | 


--------------------------------------------------------------------------------
/yolov4/darknet.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import random
  3 | from concurrent.futures.thread import ThreadPoolExecutor
  4 | from ctypes import *
  5 | from pathlib import Path
  6 | from threading import Lock
  7 | from typing import List
  8 | 
  9 | import numpy as np
 10 | 
 11 | from yolov4.helpers import init_lib, DETECTION, DETNUMPAIR, METADATA, IMAGE, read_alt_names, DarkNetPredictionResult
 12 | 
 13 | 
 14 | class Detector:
 15 | 
 16 |     def __init__(
 17 |             self,
 18 |             config_path='cfg/yolov4.cfg',
 19 |             weights_path='yolov4.weights',
 20 |             meta_path='cfg/coco.data',
 21 |             lib_darknet_path='libdarknet.so',
 22 |             batch_size=1,
 23 |             gpu_id=None
 24 |     ):
 25 |         """
 26 |         :param config_path: Path to the configuration file. Raises ValueError if not found.
 27 |         :param weights_path: Path to the weights file. Raises ValueError if not found.
 28 |         :param meta_path: Path to the data file. Raises ValueError if not found.
 29 |         :param lib_darknet_path: Path to the darknet library (.so in linux).
 30 |         :param gpu_id: GPU on which to perform the inference.
 31 |         """
 32 |         self.config_path = config_path
 33 |         self.weights_path = weights_path
 34 |         self.meta_path = meta_path
 35 |         self.gpu_id = gpu_id
 36 |         # to make sure we have only one inference per GPU.
 37 |         self.lock = Lock()
 38 | 
 39 |         self.net_main = None
 40 |         self.meta_main = None
 41 |         self.alt_names = None
 42 | 
 43 |         self.lib, self.has_gpu = init_lib(lib_darknet_path)
 44 | 
 45 |         self.predict = self.lib.network_predict_ptr
 46 |         self.predict.argtypes = [c_void_p, POINTER(c_float)]
 47 |         self.predict.restype = POINTER(c_float)
 48 | 
 49 |         # definition of all the bindings (functions to interface with C).
 50 |         if self.has_gpu:
 51 |             self.set_gpu = self.lib.cuda_set_device
 52 |             self.set_gpu.argtypes = [c_int]
 53 | 
 54 |         self.init_cpu = self.lib.init_cpu
 55 | 
 56 |         self.make_image = self.lib.make_image
 57 |         self.make_image.argtypes = [c_int, c_int, c_int]
 58 |         self.make_image.restype = IMAGE
 59 | 
 60 |         self.get_network_boxes = self.lib.get_network_boxes
 61 |         self.get_network_boxes.argtypes = [c_void_p, c_int, c_int, c_float, c_float, POINTER(c_int), c_int,
 62 |                                            POINTER(c_int),
 63 |                                            c_int]
 64 |         self.get_network_boxes.restype = POINTER(DETECTION)
 65 | 
 66 |         self.make_network_boxes = self.lib.make_network_boxes
 67 |         self.make_network_boxes.argtypes = [c_void_p]
 68 |         self.make_network_boxes.restype = POINTER(DETECTION)
 69 | 
 70 |         self.free_detections = self.lib.free_detections
 71 |         self.free_detections.argtypes = [POINTER(DETECTION), c_int]
 72 | 
 73 |         self.free_batch_detections = self.lib.free_batch_detections
 74 |         self.free_batch_detections.argtypes = [POINTER(DETNUMPAIR), c_int]
 75 | 
 76 |         self.free_ptrs = self.lib.free_ptrs
 77 |         self.free_ptrs.argtypes = [POINTER(c_void_p), c_int]
 78 | 
 79 |         self.network_predict = self.lib.network_predict_ptr
 80 |         self.network_predict.argtypes = [c_void_p, POINTER(c_float)]
 81 | 
 82 |         self.reset_rnn = self.lib.reset_rnn
 83 |         self.reset_rnn.argtypes = [c_void_p]
 84 | 
 85 |         self.load_net = self.lib.load_network
 86 |         self.load_net.argtypes = [c_char_p, c_char_p, c_int]
 87 |         self.load_net.restype = c_void_p
 88 | 
 89 |         self.load_net_custom = self.lib.load_network_custom
 90 |         self.load_net_custom.argtypes = [c_char_p, c_char_p, c_int, c_int]
 91 |         self.load_net_custom.restype = c_void_p
 92 | 
 93 |         self.do_nms_obj = self.lib.do_nms_obj
 94 |         self.do_nms_obj.argtypes = [POINTER(DETECTION), c_int, c_int, c_float]
 95 | 
 96 |         self.do_nms_sort = self.lib.do_nms_sort
 97 |         self.do_nms_sort.argtypes = [POINTER(DETECTION), c_int, c_int, c_float]
 98 | 
 99 |         self.free_image = self.lib.free_image
100 |         self.free_image.argtypes = [IMAGE]
101 | 
102 |         self.letterbox_image = self.lib.letterbox_image
103 |         self.letterbox_image.argtypes = [IMAGE, c_int, c_int]
104 |         self.letterbox_image.restype = IMAGE
105 | 
106 |         self.load_meta = self.lib.get_metadata
107 |         self.lib.get_metadata.argtypes = [c_char_p]
108 |         self.lib.get_metadata.restype = METADATA
109 | 
110 |         self.load_image = self.lib.load_image_color
111 |         self.load_image.argtypes = [c_char_p, c_int, c_int]
112 |         self.load_image.restype = IMAGE
113 | 
114 |         self.rgbgr_image = self.lib.rgbgr_image
115 |         self.rgbgr_image.argtypes = [IMAGE]
116 | 
117 |         self.predict_image = self.lib.network_predict_image
118 |         self.predict_image.argtypes = [c_void_p, IMAGE]
119 |         self.predict_image.restype = POINTER(c_float)
120 | 
121 |         self.copy_image_from_bytes = self.lib.copy_image_from_bytes
122 |         self.copy_image_from_bytes.argtypes = [IMAGE, c_char_p]
123 | 
124 |         self.predict_image_letterbox = self.lib.network_predict_image_letterbox
125 |         self.predict_image_letterbox.argtypes = [c_void_p, IMAGE]
126 |         self.predict_image_letterbox.restype = POINTER(c_float)
127 | 
128 |         self.network_predict_batch = self.lib.network_predict_batch
129 |         self.network_predict_batch.argtypes = [c_void_p, IMAGE, c_int, c_int, c_int,
130 |                                                c_float, c_float, POINTER(c_int), c_int, c_int]
131 |         self.network_predict_batch.restype = POINTER(DETNUMPAIR)
132 | 
133 |         if not Path(self.config_path).exists():
134 |             raise ValueError('Invalid config path `' + os.path.abspath(self.config_path) + '`')
135 |         if not Path(self.weights_path).exists():
136 |             raise ValueError('Invalid weight path `' + os.path.abspath(self.weights_path) + '`')
137 |         if not Path(meta_path).exists():
138 |             raise ValueError('Invalid data file path `' + os.path.abspath(self.meta_path) + '`')
139 |         if self.gpu_id is not None:
140 |             print(f'GPU -> {self.gpu_id}.')
141 |             self.set_gpu(self.gpu_id)
142 |         # batch size = 1
143 |         self.net_main = self.load_net_custom(self.config_path.encode('ascii'), self.weights_path.encode('ascii'),
144 |                                              0, batch_size)
145 |         self.meta_main = self.load_meta(self.meta_path.encode('ascii'))
146 |         # In Python 3, the metafile default access craps out on Windows (but not Linux)
147 |         # Read the names file and create a list to feed to detect
148 |         self.alt_names = read_alt_names(self.meta_path)
149 | 
150 |         self.darknet_image = self.make_image(self.network_width(), self.network_height(), 3)
151 | 
152 |     def network_width(self):
153 |         net = self.net_main
154 |         return self.lib.network_width(net)
155 | 
156 |     def network_height(self):
157 |         net = self.net_main
158 |         return self.lib.network_height(net)
159 | 
160 |     def classify(self, im):
161 |         net, meta = self.net_main, self.meta_main
162 |         out = self.predict_image(net, im)
163 |         res = []
164 |         for i in range(meta.classes):
165 |             if self.alt_names is None:
166 |                 nameTag = meta.names[i]
167 |             else:
168 |                 nameTag = self.alt_names[i]
169 |             res.append((nameTag, out[i]))
170 |         res = sorted(res, key=lambda x: -x[1])
171 |         return res
172 | 
173 |     def detect(self, image, thresh=.5, hier_thresh=.5, nms=.45, debug=False):
174 |         #
175 |         # custom_image_bgr = cv2.imread(image.decode('utf8'))  # use: detect(,,imagePath,)
176 |         # custom_image = cv2.cvtColor(custom_image_bgr, cv2.COLOR_BGR2RGB)
177 |         # custom_image = cv2.resize(custom_image, (self.network_width(), self.network_height()),
178 |         #                           interpolation=cv2.INTER_LINEAR)
179 |         # img = Image.open(image.decode('utf8'))
180 |         # custom_image2 = np.array(img.resize((self.network_width(), self.network_height())))
181 |         # self.copy_image_from_bytes(self.darknet_image, custom_image2.tobytes())
182 |         # im2 = self.load_image(image, 0, 0)
183 |         if isinstance(image, str):
184 |             if not os.path.exists(image):
185 |                 raise ValueError('Invalid image path `' + os.path.abspath(image) + '`')
186 |             image = image.encode('ascii')
187 |             im = self.load_image(image, 0, 0)
188 |         else:
189 |             self.copy_image_from_bytes(self.darknet_image, image.tobytes())
190 |             im = self.darknet_image
191 | 
192 |         ret = self.detect_image(im, thresh, hier_thresh, nms, debug)
193 |         if isinstance(image, str):
194 |             self.free_image(im)  # just when we load a new image.
195 |         return ret
196 | 
197 |     def detect_image(self, im, thresh=.5, hier_thresh=.5, nms=.45, debug=False):
198 |         net, meta = self.net_main, self.meta_main
199 |         num = c_int(0)
200 |         if debug:
201 |             print('Assigned num')
202 |         pnum = pointer(num)
203 |         if debug:
204 |             print('Assigned pnum')
205 |         self.predict_image(net, im)
206 |         letter_box = 0
207 |         # predict_image_letterbox(net, im)
208 |         # letter_box = 1
209 |         if debug:
210 |             print('did prediction')
211 |         dets = self.get_network_boxes(net, im.w, im.h, thresh, hier_thresh, None, 0, pnum, letter_box)
212 |         if debug:
213 |             print('Got dets')
214 |         num = pnum[0]
215 |         if debug:
216 |             print('got zeroth index of pnum')
217 |         if nms:
218 |             self.do_nms_sort(dets, num, meta.classes, nms)
219 |         if debug:
220 |             print('did sort')
221 |         res = []
222 |         if debug:
223 |             print('about to range')
224 |         for j in range(num):
225 |             if debug:
226 |                 print('Ranging on ' + str(j) + ' of ' + str(num))
227 |             if debug:
228 |                 print('Classes: ' + str(meta), meta.classes, meta.names)
229 |             for i in range(meta.classes):
230 |                 if debug:
231 |                     print('Class-ranging on ' + str(i) + ' of ' + str(meta.classes) + '= ' + str(dets[j].prob[i]))
232 |                 if dets[j].prob[i] > 0:
233 |                     b = dets[j].bbox
234 |                     if self.alt_names is None:
235 |                         nameTag = meta.names[i]
236 |                     else:
237 |                         nameTag = self.alt_names[i]
238 |                     if debug:
239 |                         print('Got bbox', b)
240 |                         print(nameTag)
241 |                         print(dets[j].prob[i])
242 |                         print((b.x, b.y, b.w, b.h))
243 |                     res.append((nameTag, dets[j].prob[i], (b.x, b.y, b.w, b.h)))
244 |         if debug:
245 |             print('did range')
246 |         res = sorted(res, key=lambda x: -x[1])
247 |         if debug:
248 |             print('did sort')
249 |         self.free_detections(dets, num)
250 |         if debug:
251 |             print('freed detections')
252 |         return res
253 | 
254 |     def perform_detect(
255 |             self,
256 |             image_path_or_buf='data/dog.jpg',
257 |             thresh: float = 0.25,
258 |             show_image: bool = True,
259 |             make_image_only: bool = False,
260 |     ):
261 |         self.lock.acquire()
262 |         assert 0 < thresh < 1, 'Threshold should be a float between zero and one (non-inclusive)'
263 |         detections = self.detect(image_path_or_buf, thresh)
264 |         if show_image and isinstance(image_path_or_buf, str):
265 |             try:
266 |                 from skimage import io, draw
267 |                 image = io.imread(image_path_or_buf)
268 |                 print('*** ' + str(len(detections)) + ' Results, color coded by confidence ***')
269 |                 imcaption = []
270 |                 for detection in detections:
271 |                     label = detection[0]
272 |                     confidence = detection[1]
273 |                     pstring = label + ': ' + str(np.rint(100 * confidence)) + '%'
274 |                     imcaption.append(pstring)
275 |                     print(pstring)
276 |                     bounds = detection[2]
277 |                     shape = image.shape
278 |                     # x = shape[1]
279 |                     # xExtent = int(x * bounds[2] / 100)
280 |                     # y = shape[0]
281 |                     # yExtent = int(y * bounds[3] / 100)
282 |                     yExtent = int(bounds[3])
283 |                     xEntent = int(bounds[2])
284 |                     # Coordinates are around the center
285 |                     xCoord = int(bounds[0] - bounds[2] / 2)
286 |                     yCoord = int(bounds[1] - bounds[3] / 2)
287 |                     boundingBox = [
288 |                         [xCoord, yCoord],
289 |                         [xCoord, yCoord + yExtent],
290 |                         [xCoord + xEntent, yCoord + yExtent],
291 |                         [xCoord + xEntent, yCoord]
292 |                     ]
293 |                     # Wiggle it around to make a 3px border
294 |                     rr, cc = draw.polygon_perimeter([x[1] for x in boundingBox], [x[0] for x in boundingBox],
295 |                                                     shape=shape)
296 |                     rr2, cc2 = draw.polygon_perimeter([x[1] + 1 for x in boundingBox], [x[0] for x in boundingBox],
297 |                                                       shape=shape)
298 |                     rr3, cc3 = draw.polygon_perimeter([x[1] - 1 for x in boundingBox], [x[0] for x in boundingBox],
299 |                                                       shape=shape)
300 |                     rr4, cc4 = draw.polygon_perimeter([x[1] for x in boundingBox], [x[0] + 1 for x in boundingBox],
301 |                                                       shape=shape)
302 |                     rr5, cc5 = draw.polygon_perimeter([x[1] for x in boundingBox], [x[0] - 1 for x in boundingBox],
303 |                                                       shape=shape)
304 |                     boxColor = (int(255 * (1 - (confidence ** 2))), int(255 * (confidence ** 2)), 0)
305 |                     draw.set_color(image, (rr, cc), boxColor, alpha=0.8)
306 |                     draw.set_color(image, (rr2, cc2), boxColor, alpha=0.8)
307 |                     draw.set_color(image, (rr3, cc3), boxColor, alpha=0.8)
308 |                     draw.set_color(image, (rr4, cc4), boxColor, alpha=0.8)
309 |                     draw.set_color(image, (rr5, cc5), boxColor, alpha=0.8)
310 |                 if not make_image_only:
311 |                     io.imshow(image)
312 |                     io.show()
313 |                 detections = {
314 |                     'detections': detections,
315 |                     'image': image,
316 |                     'caption': '\n<br/>'.join(imcaption)
317 |                 }
318 |             except Exception as e:
319 |                 print('Unable to show image: ' + str(e))
320 | 
321 |         results = []
322 |         sub_detections = detections['detections'] if 'detections' in detections else detections
323 |         for detection in sub_detections:
324 |             class_name, class_confidence, bbox = detection
325 |             x, y, w, h = bbox
326 |             x_min = int(x - (w / 2))
327 |             y_min = int(y - (h / 2))
328 |             results.append(DarkNetPredictionResult(
329 |                 class_name=class_name,
330 |                 class_confidence=class_confidence,
331 |                 left_x=max(0, int(x_min)),
332 |                 top_y=max(0, int(y_min)),
333 |                 width=max(0, int(w)),
334 |                 height=max(0, int(h))
335 |             ))
336 |         self.lock.release()
337 |         return results
338 | 
339 |     def perform_batch_detect(self, image_list, thresh=0.25, hier_thresh=.5, nms=.45, batch_size=3):
340 |         net, meta = self.net_main, self.meta_main
341 |         # NB! Image sizes should be the same
342 |         # You can change the images, yet, be sure that they have the same width and height
343 |         pred_height, pred_width, c = image_list[0].shape
344 |         net_width, net_height = (self.network_width(), self.network_height())
345 |         arr = np.concatenate(image_list, axis=0)
346 |         arr = np.ascontiguousarray(arr.flat, dtype=np.float32) / 255.0
347 |         data = arr.ctypes.data_as(POINTER(c_float))
348 |         im = IMAGE(net_width, net_height, c, data)
349 | 
350 |         batch_dets = self.network_predict_batch(net, im, batch_size, pred_width,
351 |                                                 pred_height, thresh, hier_thresh, None, 0, 0)
352 |         batch_boxes = []
353 |         batch_scores = []
354 |         batch_classes = []
355 |         for b in range(batch_size):
356 |             num = batch_dets[b].num
357 |             dets = batch_dets[b].dets
358 |             if nms:
359 |                 self.do_nms_obj(dets, num, meta.classes, nms)
360 |             boxes = []
361 |             scores = []
362 |             classes = []
363 |             for i in range(num):
364 |                 det = dets[i]
365 |                 score = -1
366 |                 label = None
367 |                 for c in range(det.classes):
368 |                     p = det.prob[c]
369 |                     if p > score:
370 |                         score = p
371 |                         label = c
372 |                 if score > thresh:
373 |                     box = det.bbox
374 |                     left, top, right, bottom = map(int, (box.x - box.w / 2, box.y - box.h / 2,
375 |                                                          box.x + box.w / 2, box.y + box.h / 2))
376 |                     boxes.append((top, left, bottom, right))
377 |                     scores.append(score)
378 |                     classes.append(label)
379 |                     # boxColor = (int(255 * (1 - (score ** 2))), int(255 * (score ** 2)), 0)
380 |                     # cv2.rectangle(image_list[b], (left, top),
381 |                     #               (right, bottom), boxColor, 2)
382 |             # cv2.imwrite(os.path.basename(img_samples[b]), image_list[b])
383 | 
384 |             batch_boxes.append(boxes)
385 |             batch_scores.append(scores)
386 |             batch_classes.append(classes)
387 |         self.free_batch_detections(batch_dets, batch_size)
388 |         # batch_classes: use altNames to fetch the real names.
389 |         return batch_boxes, batch_scores, batch_classes
390 | 
391 | 
392 | class MultiGPU:
393 | 
394 |     def __init__(self, detectors_list: List[Detector]):
395 |         self.detectors = {}
396 |         for detector in detectors_list:
397 |             gpu_id = detector.gpu_id
398 |             if gpu_id is None:
399 |                 raise ValueError('To use MultiGPU, every gpu_id should be defined.')
400 |             self.detectors[gpu_id] = detector
401 |         self.num_gpus = len(self.detectors)
402 |         self.thread_pool = ThreadPoolExecutor(max_workers=self.num_gpus, thread_name_prefix='detector_')
403 |         self.counter = 0  # used to dispatch, assuming the load is same on each GPU.
404 |         self.locks = {}
405 |         for gpu_id in self.detectors.keys():
406 |             self.locks[gpu_id] = Lock()
407 | 
408 |     def find_available_gpu(self):
409 |         list_to_choose_from = []
410 |         for gpu_id, lock in self.locks.items():
411 |             if not lock.locked():
412 |                 list_to_choose_from.append(gpu_id)
413 |         if len(list_to_choose_from) == 0:
414 |             list_to_choose_from = list(self.locks.keys())
415 |         return random.choice(list_to_choose_from)
416 | 
417 |     def _perform_detect(self, gpu_id, *args):
418 |         return self.detectors[gpu_id].perform_detect(*args)
419 | 
420 |     def perform_detect(self, *args, **kwargs):
421 |         gpu_id = self.find_available_gpu()
422 |         lock = self.locks[gpu_id]
423 |         with lock:
424 |             # print(f'Lock for GPU {gpu_id} acquired.')
425 | 
426 |             def run(*arg, **kwarg):
427 |                 return self.detectors[gpu_id].perform_detect(*arg, **kwarg)
428 | 
429 |             future = self.thread_pool.submit(fn=run, *args, **kwargs)
430 |             return future.result()
431 | 


--------------------------------------------------------------------------------