255 lines
8.1 KiB
Python
255 lines
8.1 KiB
Python
# pylint: disable=R, W0401, W0614, W0703
|
|
from ctypes import *
|
|
import random
|
|
import os
|
|
import cv2
|
|
import platform
|
|
from typing import List, Tuple
|
|
|
|
# C-structures from Darknet lib
|
|
|
|
class BOX(Structure):
|
|
_fields_ = [("x", c_float),
|
|
("y", c_float),
|
|
("w", c_float),
|
|
("h", c_float)]
|
|
|
|
|
|
class DETECTION(Structure):
|
|
_fields_ = [("bbox", BOX),
|
|
("classes", c_int),
|
|
("best_class_idx", c_int),
|
|
("prob", POINTER(c_float)),
|
|
("mask", POINTER(c_float)),
|
|
("objectness", c_float),
|
|
("sort_class", c_int),
|
|
("uc", POINTER(c_float)),
|
|
("points", c_int),
|
|
("embeddings", POINTER(c_float)),
|
|
("embedding_size", c_int),
|
|
("sim", c_float),
|
|
("track_id", c_int)]
|
|
|
|
class IMAGE(Structure):
|
|
_fields_ = [("w", c_int),
|
|
("h", c_int),
|
|
("c", c_int),
|
|
("data", POINTER(c_float))]
|
|
|
|
|
|
class METADATA(Structure):
|
|
_fields_ = [("classes", c_int),
|
|
("names", POINTER(c_char_p))]
|
|
|
|
class YoloNet:
|
|
"""Darknet-based detector implementation"""
|
|
net: c_void_p
|
|
meta: METADATA
|
|
|
|
def __init__(self, weight_path: str, meta_path: str, config_path: str, asked_to_use_gpu: bool):
|
|
if not os.path.exists(config_path):
|
|
raise ValueError("Invalid config path `"+os.path.abspath(config_path)+"`")
|
|
if not os.path.exists(weight_path):
|
|
raise ValueError("Invalid weight path `"+os.path.abspath(weight_path)+"`")
|
|
if not os.path.exists(meta_path):
|
|
raise ValueError("Invalid data file path `"+os.path.abspath(meta_path)+"`")
|
|
if not lib:
|
|
raise ImportError(f"Unable to load darknet module.")
|
|
|
|
if asked_to_use_gpu and not using_gpu:
|
|
raise Exception('I respectfully decline to load the net as I am asked to use GPU but the loaded darknet module does NOT have GPU support')
|
|
|
|
self.net = load_net_custom(config_path.encode("ascii"), weight_path.encode("ascii"), 0, 1) # batch size = 1
|
|
self.meta = load_meta(meta_path.encode("ascii"))
|
|
|
|
def detect(self, meta, image, alt_names, thresh=.5, hier_thresh=.5, nms=.45, debug=False) -> List[Tuple[str, float, Tuple[float, float, float, float]]]:
|
|
#pylint: disable= C0321
|
|
custom_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
im, arr = array_to_image(custom_image) # you should comment line below: free_image(im)
|
|
if debug:
|
|
print("Loaded image")
|
|
num = c_int(0)
|
|
if debug:
|
|
print("Assigned num")
|
|
pnum = pointer(num)
|
|
if debug:
|
|
print("Assigned pnum")
|
|
predict_image(self.net, im)
|
|
if debug:
|
|
print("did prediction")
|
|
dets = get_network_boxes(self.net, custom_image.shape[1], custom_image.shape[0], thresh, hier_thresh, None, 0, pnum, 0) # OpenCV
|
|
if debug:
|
|
print("Got dets")
|
|
num = pnum[0]
|
|
if debug:
|
|
print("got zeroth index of pnum")
|
|
if nms:
|
|
do_nms_sort(dets, num, meta.classes, nms)
|
|
if debug:
|
|
print("did sort")
|
|
res = []
|
|
if debug:
|
|
print("about to range")
|
|
for j in range(num):
|
|
if debug:
|
|
print("Ranging on "+str(j)+" of "+str(num))
|
|
if debug:
|
|
print("Classes: "+str(meta), meta.classes, meta.names)
|
|
for i in range(meta.classes):
|
|
if debug:
|
|
print("Class-ranging on "+str(i)+" of "+str(meta.classes)+"= "+str(dets[j].prob[i]))
|
|
if dets[j].prob[i] > 0:
|
|
b = dets[j].bbox
|
|
if alt_names is None:
|
|
nameTag = meta.names[i]
|
|
else:
|
|
nameTag = alt_names[i]
|
|
if debug:
|
|
print("Got bbox", b)
|
|
print(nameTag)
|
|
print(dets[j].prob[i])
|
|
print((b.x, b.y, b.w, b.h))
|
|
res.append((nameTag, dets[j].prob[i], (b.x, b.y, b.w, b.h)))
|
|
if debug:
|
|
print("did range")
|
|
res = sorted(res, key=lambda x: -x[1])
|
|
if debug:
|
|
print("did sort")
|
|
free_detections(dets, num)
|
|
if debug:
|
|
print("freed detections")
|
|
return res
|
|
|
|
# Loads darknet shared library. May fail if some dependencies like OpenCV not installed
|
|
# libdarknet_gpu.so needs Cuda + Cudnn and other libraries in path, which may not exist
|
|
# For the such case, it will try to load libdarknet.so instead
|
|
lib = None
|
|
using_gpu = False
|
|
|
|
print('\n')
|
|
so_path = os.path.join('/darknet', "libdarknet_cpu.so")
|
|
lib = CDLL(so_path, RTLD_GLOBAL)
|
|
print(f" Darknet is now running on CPU.")
|
|
print('\n')
|
|
|
|
if lib:
|
|
lib.network_width.argtypes = [c_void_p]
|
|
lib.network_width.restype = c_int
|
|
lib.network_height.argtypes = [c_void_p]
|
|
lib.network_height.restype = c_int
|
|
|
|
predict = lib.network_predict
|
|
predict.argtypes = [c_void_p, POINTER(c_float)]
|
|
predict.restype = POINTER(c_float)
|
|
|
|
if using_gpu:
|
|
set_gpu = lib.cuda_set_device
|
|
set_gpu.argtypes = [c_int]
|
|
|
|
make_image = lib.make_image
|
|
make_image.argtypes = [c_int, c_int, c_int]
|
|
make_image.restype = IMAGE
|
|
|
|
get_network_boxes = lib.get_network_boxes
|
|
get_network_boxes.argtypes = [c_void_p, c_int, c_int, c_float, c_float, POINTER(c_int), c_int, POINTER(c_int), c_int]
|
|
get_network_boxes.restype = POINTER(DETECTION)
|
|
|
|
make_network_boxes = lib.make_network_boxes
|
|
make_network_boxes.argtypes = [c_void_p]
|
|
make_network_boxes.restype = POINTER(DETECTION)
|
|
|
|
free_detections = lib.free_detections
|
|
free_detections.argtypes = [POINTER(DETECTION), c_int]
|
|
|
|
free_ptrs = lib.free_ptrs
|
|
free_ptrs.argtypes = [POINTER(c_void_p), c_int]
|
|
|
|
network_predict = lib.network_predict
|
|
network_predict.argtypes = [c_void_p, POINTER(c_float)]
|
|
|
|
reset_rnn = lib.reset_rnn
|
|
reset_rnn.argtypes = [c_void_p]
|
|
|
|
load_net = lib.load_network
|
|
load_net.argtypes = [c_char_p, c_char_p, c_int]
|
|
load_net.restype = c_void_p
|
|
|
|
load_net_custom = lib.load_network_custom
|
|
load_net_custom.argtypes = [c_char_p, c_char_p, c_int, c_int]
|
|
load_net_custom.restype = c_void_p
|
|
|
|
do_nms_obj = lib.do_nms_obj
|
|
do_nms_obj.argtypes = [POINTER(DETECTION), c_int, c_int, c_float]
|
|
|
|
do_nms_sort = lib.do_nms_sort
|
|
do_nms_sort.argtypes = [POINTER(DETECTION), c_int, c_int, c_float]
|
|
|
|
free_image = lib.free_image
|
|
free_image.argtypes = [IMAGE]
|
|
|
|
letterbox_image = lib.letterbox_image
|
|
letterbox_image.argtypes = [IMAGE, c_int, c_int]
|
|
letterbox_image.restype = IMAGE
|
|
|
|
load_meta = lib.get_metadata
|
|
lib.get_metadata.argtypes = [c_char_p]
|
|
lib.get_metadata.restype = METADATA
|
|
|
|
load_image = lib.load_image_color
|
|
load_image.argtypes = [c_char_p, c_int, c_int]
|
|
load_image.restype = IMAGE
|
|
|
|
rgbgr_image = lib.rgbgr_image
|
|
rgbgr_image.argtypes = [IMAGE]
|
|
|
|
predict_image = lib.network_predict_image
|
|
predict_image.argtypes = [c_void_p, IMAGE]
|
|
predict_image.restype = POINTER(c_float)
|
|
|
|
def sample(probs):
|
|
s = sum(probs)
|
|
probs = [a/s for a in probs]
|
|
r = random.uniform(0, 1)
|
|
for i in range(len(probs)):
|
|
r = r - probs[i]
|
|
if r <= 0:
|
|
return i
|
|
return len(probs)-1
|
|
|
|
|
|
def c_array(ctype, values):
|
|
arr = (ctype*len(values))()
|
|
arr[:] = values
|
|
return arr
|
|
|
|
def array_to_image(arr):
|
|
import numpy as np
|
|
# need to return old values to avoid python freeing memory
|
|
arr = arr.transpose(2, 0, 1)
|
|
c = arr.shape[0]
|
|
h = arr.shape[1]
|
|
w = arr.shape[2]
|
|
arr = np.ascontiguousarray(arr.flat, dtype=np.float32) / 255.0
|
|
data = arr.ctypes.data_as(POINTER(c_float))
|
|
im = IMAGE(w, h, c, data)
|
|
return im, arr
|
|
|
|
|
|
def classify(net, meta, im):
|
|
global alt_names
|
|
|
|
out = predict_image(net, im)
|
|
res = []
|
|
for i in range(meta.classes):
|
|
if alt_names is None:
|
|
nameTag = meta.names[i]
|
|
else:
|
|
nameTag = alt_names[i]
|
|
res.append((nameTag, out[i]))
|
|
res = sorted(res, key=lambda x: -x[1])
|
|
return res
|
|
|
|
|
|
|
|
|