diff --git a/ml_api/.gitignore b/ml_api/.gitignore index ac23329fa..76d4d767e 100644 --- a/ml_api/.gitignore +++ b/ml_api/.gitignore @@ -1 +1,2 @@ model/model.weights +model/*.onnx diff --git a/ml_api/detect.py b/ml_api/detect.py new file mode 100755 index 000000000..3deb24195 --- /dev/null +++ b/ml_api/detect.py @@ -0,0 +1,119 @@ +#!python3 +import cv2 +from dataclasses import asdict +import json +from lib.geometry import compare_detections, Detection +import os +from lib.detection_model import * + +KNOWN_IMAGE_EXTENSIONS = ('.jpg', '.png') +KNOWN_VIDEO_EXTENSIONS = ('.mp4', '.avi') + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("image", type=str, help="Image file path") + parser.add_argument("--weights", type=str, default="model/model.weights", help="Model weights file") + parser.add_argument("--det-threshold", type=float, default=0.25, help="Detection threshold") + parser.add_argument("--nms-threshold", type=float, default=0.4, help="NMS threshold") + parser.add_argument("--preheat", action='store_true', help="Make a dry run of NN for initlalization") + parser.add_argument("--cpu", action='store_true', help="Force use CPU") + parser.add_argument("--save-detections-to", type=str, help="Save detections into this file") + parser.add_argument("--compare-detections-with", type=str, help="Load detections from this file and compare with result") + parser.add_argument("--render-to", type=str, help="Save detections into this file or directory") + parser.add_argument("--print", action='store_true', help="Print detections") + opt = parser.parse_args() + + + net_main_1, meta_main_1 = load_net("model/model.cfg", opt.weights, "model/model.meta") + + # force use CPU, only implemented for ONNX + if opt.cpu and onnx_ready and isinstance(net_main_1, OnnxNet): + net_main_1.force_cpu() + + assert os.path.exists(opt.image) + filename = os.path.basename(opt.image) + filename, extension = os.path.splitext(filename) + + is_image = extension in KNOWN_IMAGE_EXTENSIONS + is_video = extension in KNOWN_VIDEO_EXTENSIONS + frame_number = 0 + vwr = None + if is_video: + cap = cv2.VideoCapture(opt.image) + fps = cap.get(cv2.CAP_PROP_FPS) + frame_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + frame_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + reading_success, custom_image_bgr = cap.read() + if opt.render_to: + fourcc = cv2.VideoWriter_fourcc("m", "p", "4", "v") + vwr = cv2.VideoWriter(opt.render_to, fourcc, fps, (frame_w, frame_h)) + else: + cap = None + fps = 0.0 + custom_image_bgr = cv2.imread(opt.image) + reading_success = True + + + # this will make library initialize all the required resources at the first run + # then the following runs will be much faster + if opt.preheat: + detections = detect(net_main_1, meta_main_1, custom_image_bgr, thresh=opt.det_threshold, nms=opt.nms_threshold) + + while reading_success: + started_at = time.time() + detections = detect(net_main_1, meta_main_1, custom_image_bgr, thresh=opt.det_threshold, nms=opt.nms_threshold) + finished_at = time.time() + execution_time = finished_at - started_at + print(f"Frame #{frame_number} execution time: {execution_time:.3} sec, detection count: {len(detections)}") + + detections = Detection.from_tuple_list(detections) + # dump detections into some file + if opt.save_detections_to: + output_filename, output_extension = os.path.splitext(opt.save_detections_to) + if is_video and not output_extension and not os.path.exists(opt.save_detections_to): + os.makedirs(opt.save_detections_to) + if os.path.isdir(opt.save_detections_to): + if is_video: + output_file_name = f"{filename}#{frame_number:04}.json" + else: + output_file_name = f"{filename}.json" + output_file_name = os.path.join(opt.save_detections_to, output_file_name) + else: + output_file_name = opt.save_detections_to + + with open(output_file_name, "w") as f: + json.dump([asdict(d) for d in detections], f) + + # load detections from some file and compare with detection result + if opt.compare_detections_with: + if is_video: + read_file_name = os.path.join(opt.compare_detections_with, f"{filename}#{frame_number:04}.json") + else: + read_file_name = opt.compare_detections_with + + with open(read_file_name) as f: + items = json.load(f) + loaded = [Detection.from_dict(d) for d in items] + compare_result = compare_detections(loaded, detections) + if not compare_result: + print(f"Frame #{frame_number} loaded detections and resulting are different") + if opt.render_to: + for d in detections: + cv2.rectangle(custom_image_bgr, + (int(d.box.left()), int(d.box.top())), (int(d.box.right()), int(d.box.bottom())), + (0, 255, 0), 2) + if vwr: + vwr.write(custom_image_bgr) + else: + cv2.imwrite(opt.render_to, custom_image_bgr) + + + if opt.print: + print(detections) + + if is_image: + reading_success = False + elif cap: + reading_success, custom_image_bgr = cap.read() + frame_number += 1 + diff --git a/ml_api/lib/backend_darknet.py b/ml_api/lib/backend_darknet.py new file mode 100644 index 000000000..3652489e3 --- /dev/null +++ b/ml_api/lib/backend_darknet.py @@ -0,0 +1,239 @@ +# 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), + ("prob", POINTER(c_float)), + ("mask", POINTER(c_float)), + ("objectness", c_float), + ("sort_class", 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, config_path: str, weight_path: str, meta_path: str): + 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)+"`") + 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 + +DIRNAME = os.path.abspath( + os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "bin") +) +# Loads darknet shared library. May fail if some dependencies like OpenCV not installed + +hasGPU = os.environ.get('HAS_GPU', 'False') == 'True' +so_path = os.path.join(DIRNAME, "model_{}{}.so".format('gpu_' if hasGPU else '', platform.machine())) + +lib = CDLL(so_path, RTLD_GLOBAL) +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 hasGPU: + 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 + + + + diff --git a/ml_api/lib/backend_onnx.py b/ml_api/lib/backend_onnx.py new file mode 100644 index 000000000..47e0e060c --- /dev/null +++ b/ml_api/lib/backend_onnx.py @@ -0,0 +1,138 @@ +from typing import List, Tuple +import onnxruntime +import numpy as np +import cv2 +import os + +from lib.meta import Meta + + +class OnnxNet: + session: onnxruntime.InferenceSession + meta: Meta + + def __init__(self, onnx_path: str, meta_path: str): + if not os.path.exists(onnx_path): + raise ValueError("Invalid weight path `"+os.path.abspath(onnx_path)+"`") + if not os.path.exists(meta_path): + raise ValueError("Invalid data file path `"+os.path.abspath(meta_path)+"`") + self.session = onnxruntime.InferenceSession(onnx_path) + self.meta = Meta(meta_path) + + 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]]]: + input_h = self.session.get_inputs()[0].shape[2] + input_w = self.session.get_inputs()[0].shape[3] + width = image.shape[1] + height = image.shape[0] + + # Input + resized = cv2.resize(image, (input_w, input_h), interpolation=cv2.INTER_LINEAR) + img_in = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB) + img_in = np.transpose(img_in, (2, 0, 1)).astype(np.float32) + img_in = np.expand_dims(img_in, axis=0) + img_in /= 255.0 + + input_name = self.session.get_inputs()[0].name + outputs = self.session.run(None, {input_name: img_in}) + + detections = post_processing(outputs, width, height, thresh, nms, meta.names) + return detections[0] + + def force_cpu(self): + self.session.set_providers(['CPUExecutionProvider']) + +def nms_cpu(boxes, confs, nms_thresh=0.5, min_mode=False): + # print(boxes.shape) + x1 = boxes[:, 0] + y1 = boxes[:, 1] + x2 = boxes[:, 2] + y2 = boxes[:, 3] + + areas = (x2 - x1) * (y2 - y1) + order = confs.argsort()[::-1] + + keep = [] + while order.size > 0: + idx_self = order[0] + idx_other = order[1:] + + keep.append(idx_self) + + xx1 = np.maximum(x1[idx_self], x1[idx_other]) + yy1 = np.maximum(y1[idx_self], y1[idx_other]) + xx2 = np.minimum(x2[idx_self], x2[idx_other]) + yy2 = np.minimum(y2[idx_self], y2[idx_other]) + + w = np.maximum(0.0, xx2 - xx1) + h = np.maximum(0.0, yy2 - yy1) + inter = w * h + + if min_mode: + over = inter / np.minimum(areas[order[0]], areas[order[1:]]) + else: + over = inter / (areas[order[0]] + areas[order[1:]] - inter) + + inds = np.where(over <= nms_thresh)[0] + order = order[inds + 1] + + return np.array(keep) + +def post_processing(output, width, height, conf_thresh, nms_thresh, names): + box_array = output[0] + confs = output[1] + + if type(box_array).__name__ != 'ndarray': + box_array = box_array.cpu().detach().numpy() + confs = confs.cpu().detach().numpy() + + num_classes = confs.shape[2] + + # [batch, num, 4] + box_array = box_array[:, :, 0] + + # [batch, num, num_classes] --> [batch, num] + max_conf = np.max(confs, axis=2) + max_id = np.argmax(confs, axis=2) + + box_x1x1x2y2_to_xcycwh_scaled = lambda b: \ + ( + float(0.5 * width * (b[0] + b[2])), + float(0.5 * height * (b[1] + b[3])), + float(width * (b[2] - b[0])), + float(width * (b[3] - b[1])) + ) + dets_batch = [] + for i in range(box_array.shape[0]): + + argwhere = max_conf[i] > conf_thresh + l_box_array = box_array[i, argwhere, :] + l_max_conf = max_conf[i, argwhere] + l_max_id = max_id[i, argwhere] + + bboxes = [] + # nms for each class + for j in range(num_classes): + + cls_argwhere = l_max_id == j + ll_box_array = l_box_array[cls_argwhere, :] + ll_max_conf = l_max_conf[cls_argwhere] + ll_max_id = l_max_id[cls_argwhere] + + keep = nms_cpu(ll_box_array, ll_max_conf, nms_thresh) + + if (keep.size > 0): + ll_box_array = ll_box_array[keep, :] + ll_max_conf = ll_max_conf[keep] + ll_max_id = ll_max_id[keep] + + for k in range(ll_box_array.shape[0]): + bboxes.append([ll_box_array[k, 0], ll_box_array[k, 1], ll_box_array[k, 2], ll_box_array[k, 3], ll_max_conf[k], ll_max_conf[k], ll_max_id[k]]) + + detections = [(names[b[6]], float(b[4]), box_x1x1x2y2_to_xcycwh_scaled((b[0], b[1], b[2], b[3]))) for b in bboxes] + dets_batch.append(detections) + + + return dets_batch + + + diff --git a/ml_api/lib/detection_model.py b/ml_api/lib/detection_model.py index 7a8a99cd2..7ffb98e92 100644 --- a/ml_api/lib/detection_model.py +++ b/ml_api/lib/detection_model.py @@ -1,272 +1,64 @@ #!python3 # pylint: disable=R, W0401, W0614, W0703 -from ctypes import * -import math -import random -import os -import sys +from enum import Enum import cv2 -import platform - - -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 - - -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), - ("prob", POINTER(c_float)), - ("mask", POINTER(c_float)), - ("objectness", c_float), - ("sort_class", 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))] - - -DIRNAME = os.path.abspath( - os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "bin") -) - -hasGPU = os.environ.get('HAS_GPU', 'False') == 'True' -so_path = os.path.join(DIRNAME, "model_{}{}.so".format('gpu_' if hasGPU else '', platform.machine())) - -lib = CDLL(so_path, RTLD_GLOBAL) -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 hasGPU: - 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 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): - 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 - - -def detect(net, meta, image, thresh=.5, hier_thresh=.5, nms=.45, debug=False): - #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(net, im) - if debug: - print("did prediction") - dets = get_network_boxes(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 +import time +import argparse +from lib.meta import Meta +#GLOBALS net_main = None meta_main = None alt_names = None +# optional import for darknet +try: + from lib.backend_darknet import YoloNet + darknet_ready = True +except: + darknet_ready = False + +# optional import for onnx +try: + from lib.backend_onnx import OnnxNet + onnx_ready = True +except: + onnx_ready = False + def load_net(config_path, weight_path, meta_path): + """Loads network from config files and weights. Automatically detects the backend.""" + # nets are loaded only once and then reused global meta_main, net_main, alt_names # pylint: disable=W0603 - 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 net_main is None: - net_main = load_net_custom(config_path.encode("ascii"), weight_path.encode("ascii"), 0, 1) # batch size = 1 - if meta_main is None: - meta_main = load_meta(meta_path.encode("ascii")) + if onnx_ready and weight_path.endswith(".onnx"): + net_main = OnnxNet(weight_path, meta_path) + elif darknet_ready: + net_main = YoloNet(config_path, weight_path, meta_path) + else: + raise Exception(f"Unable to load net. Onnx_ready={onnx_ready}, Darknet_ready={darknet_ready}") + + meta_main = net_main.meta + + assert net_main is not None + assert meta_main is not None + if alt_names is None: # In Python 3, the metafile default access craps out on Windows (but not Linux) # Read the names file and create a list to feed to detect try: - with open(meta_path) as metaFH: - metaContents = metaFH.read() - import re - match = re.search("names *= *(.*)$", metaContents, re.IGNORECASE | re.MULTILINE) - if match: - result = match.group(1) - else: - result = None - try: - if os.path.exists(result): - with open(result) as namesFH: - namesList = namesFH.read().strip().split("\n") - alt_names = [x.strip() for x in namesList] - except TypeError: - pass + meta = Meta(meta_path) + alt_names = meta.names except Exception: pass return net_main, meta_main +def detect(net, meta, image, thresh=.5, hier_thresh=.5, nms=.45, debug=False): + """Runs detection on some image content""" -if __name__ == "__main__": - net_main_1, meta_main_1 = load_net("model/model.cfg", "model/model.weights", "model/model.meta") + return net.detect(meta, image, alt_names, thresh, hier_thresh, nms, debug) - import cv2 - custom_image_bgr = cv2.imread(sys.argv[1]) # use: detect(,,imagePath,) - print(detect(net_main_1, meta_main_1, custom_image_bgr, thresh=0.25)) diff --git a/ml_api/lib/geometry.py b/ml_api/lib/geometry.py new file mode 100644 index 000000000..963f4fdd8 --- /dev/null +++ b/ml_api/lib/geometry.py @@ -0,0 +1,111 @@ +from dataclasses import dataclass, asdict +from typing import Any, Dict, List, Tuple + +@dataclass +class Box: + """Detection rect""" + xc: float + yc: float + w: float + h: float + + @classmethod + def from_tuple(cls, box: Tuple[float, float, float, float]) -> 'Box': + return Box(xc=float(box[0]), yc=float(box[1]), w=float(box[2]), h=float(box[3])) + + def left(self) -> float: + return self.xc - self.w * 0.5 + + def right(self) -> float: + return self.xc + self.w * 0.5 + + def top(self) -> float: + return self.yc - self.h * 0.5 + + def bottom(self) -> float: + return self.yc + self.h * 0.5 + + def calc_iou(self, other: 'Box') -> float: + """Calculates intersection over union ration which can be used to compare boxes""" + al = self.left() + ar = self.right() + at = self.top() + ab = self.bottom() + + bl = other.left() + br = other.right() + bt = other.top() + bb = other.bottom() + + i_l = max(al, bl) + i_r = min(ar, br) + i_t = max(at, bt) + i_b = min(ab, bb) + + o_l = min(al, bl) + o_r = max(ar, br) + o_t = min(at, bt) + o_b = max(ab, bb) + + i_w = i_r - i_l + i_h = i_b - i_t + o_w = o_r - o_l + o_h = o_b - o_t + + o_a = o_w * o_h + if o_a <= 0.0: + return 0.0 + return i_w * i_h / o_a + + +@dataclass +class Detection: + """Detection result""" + name: str + confidence: float + box: Box + + @classmethod + def from_tuple_list(cls, detections: List[Tuple[str, float, Tuple[float, float, float, float]]]) -> List['Detection']: + return [Detection.from_tuple(d) for d in detections] + + @classmethod + def from_tuple(cls, detection: Tuple[str, float, Tuple[float, float, float, float]]) -> 'Detection': + box = Box.from_tuple(detection[2]) + return Detection(detection[0], float(detection[1]), box) + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> 'Detection': + return Detection(data['name'], data['confidence'], Box(**data['box'])) + + + +def compare_detections(l1: List[Detection], l2: List[Detection], threshold: float = 0.4) -> bool: + """Compares two lists of detections. Returns true if lists looks similar with some threshold""" + + # Are there all boxes from l1 matching any in l2 + for a in l1: + found = False + for b in l2: + iou = a.box.calc_iou(b.box) + if iou >= threshold: + found = True + break + if not found: + return False + + # are there all boxes in l2 matching any in l1 + # the list may differ and contain duplicates, + # that's why we need two checks + for b in l2: + found = False + for a in l1: + iou = a.box.calc_iou(b.box) + if iou >= threshold: + found = True + break + if not found: + return False + + return True + diff --git a/ml_api/lib/meta.py b/ml_api/lib/meta.py new file mode 100644 index 000000000..a9257a620 --- /dev/null +++ b/ml_api/lib/meta.py @@ -0,0 +1,27 @@ +from typing import List, Tuple +from dataclasses import dataclass, field +import os +import re + +@dataclass +class Meta: + names: List[str] = field(default_factory=list) + + def __init__(self, meta_path: str): + names = None + with open(meta_path) as f: + meta_contents = f.read() + match = re.search("names *= *(.*)$", meta_contents, re.IGNORECASE | re.MULTILINE) + if match: + names_path = match.group(1) + try: + if os.path.exists(names_path): + with open(names_path) as namesFH: + names_list = namesFH.read().strip().split("\n") + names = [x.strip() for x in names_list] + except TypeError: + pass + if names is None: + names = ['failure'] + + self.names = names diff --git a/ml_api/server.py b/ml_api/server.py index 0cdbf5107..f1899089b 100755 --- a/ml_api/server.py +++ b/ml_api/server.py @@ -26,9 +26,10 @@ # SECURITY WARNING: don't run with debug turned on in production! app.config['DEBUG'] = environ.get('DEBUG') == 'True' +MODEL_NAME = environ.get('MODEL_NAME') or 'model.weights' model_dir = path.join(path.dirname(path.realpath(__file__)), 'model') -net_main, meta_main = load_net(path.join(model_dir, 'model.cfg'), path.join(model_dir, 'model.weights'), path.join(model_dir, 'model.meta')) +net_main, meta_main = load_net(path.join(model_dir, 'model.cfg'), path.join(model_dir, MODEL_NAME), path.join(model_dir, 'model.meta')) @app.route('/p/', methods=['GET']) @token_required diff --git a/ml_api/test_data/1.jpg b/ml_api/test_data/1.jpg new file mode 100644 index 000000000..590517171 Binary files /dev/null and b/ml_api/test_data/1.jpg differ diff --git a/ml_api/test_data/1.json b/ml_api/test_data/1.json new file mode 100644 index 000000000..327a2abc7 --- /dev/null +++ b/ml_api/test_data/1.json @@ -0,0 +1 @@ +[{"name": "failure", "confidence": 0.6891447305679321, "box": {"xc": 535.3487373590469, "yc": 474.9534823894501, "w": 252.1463025212288, "h": 138.7461107969284}}, {"name": "failure", "confidence": 0.6743507385253906, "box": {"xc": 263.38869220018387, "yc": 544.919519662857, "w": 276.50628340244293, "h": 185.15091562271118}}, {"name": "failure", "confidence": 0.5817742347717285, "box": {"xc": 589.4390723705292, "yc": 481.2896337509155, "w": 121.88863205909729, "h": 94.02524149417877}}, {"name": "failure", "confidence": 0.47977277636528015, "box": {"xc": 405.2845135331154, "yc": 396.83880066871643, "w": 173.75898277759552, "h": 80.74439334869385}}, {"name": "failure", "confidence": 0.33845505118370056, "box": {"xc": 349.38735204935074, "yc": 382.78335213661194, "w": 164.63774612545967, "h": 82.647167801857}}] \ No newline at end of file diff --git a/ml_api/tests/__init__.py b/ml_api/tests/__init__.py new file mode 100644 index 000000000..37903a5c1 --- /dev/null +++ b/ml_api/tests/__init__.py @@ -0,0 +1,10 @@ +import unittest +from tests.test_detection import TestDetection +from tests.test_geometry import * + +# Test suite for all tests +def suite(): + suite = unittest.TestSuite() + suite.addTest(TestGeometry()) + suite.addTest(TestDetection()) + return suite diff --git a/ml_api/tests/test_detection.py b/ml_api/tests/test_detection.py new file mode 100644 index 000000000..ee1eb5cfb --- /dev/null +++ b/ml_api/tests/test_detection.py @@ -0,0 +1,45 @@ +from typing import List, Tuple +import unittest +from lib.detection_model import * +import os +import json +from lib.geometry import Detection, compare_detections +import cv2 + +TEST_DATA_DIR = "test_data" +DET_THRESHOLD = 0.25 +NMS_THRESHOLD = 0.4 + +def find_images() -> List[Tuple[str, List[Detection]]]: + res = [] + for name in os.listdir(TEST_DATA_DIR): + if name.endswith(".jpg"): + json_path = os.path.join(TEST_DATA_DIR, name[:-3] + "json") + if os.path.exists(json_path): + with open(json_path) as f: + items = json.load(f) + res.append((os.path.join(TEST_DATA_DIR, name), [Detection.from_dict(d) for d in items])) + return res + + +class TestDetection(unittest.TestCase): + def test_darknet(self): + net, meta = load_net("model/model.cfg", "model/model.weights", "model/model.meta") + for img_path, detections in find_images(): + custom_image_bgr = cv2.imread(img_path) + detected = detect(net, meta, custom_image_bgr, thresh=DET_THRESHOLD, nms=NMS_THRESHOLD) + detected_detections = Detection.from_tuple_list(detected) + self.assertTrue(len(detected_detections) > 0) + similar = compare_detections(detected_detections, detections) + self.assertTrue(similar) + + def test_onnx(self): + net, meta = load_net("model/model.cfg", "model/model.onnx", "model/model.meta") + for img_path, detections in find_images(): + custom_image_bgr = cv2.imread(img_path) + detected = detect(net, meta, custom_image_bgr, thresh=DET_THRESHOLD, nms=NMS_THRESHOLD) + detected_detections = Detection.from_tuple_list(detected) + self.assertTrue(len(detected_detections) > 0) + similar = compare_detections(detected_detections, detections) + self.assertTrue(similar) + diff --git a/ml_api/tests/test_geometry.py b/ml_api/tests/test_geometry.py new file mode 100644 index 000000000..943b163f7 --- /dev/null +++ b/ml_api/tests/test_geometry.py @@ -0,0 +1,15 @@ +import unittest +from lib.geometry import * + + +class TestGeometry(unittest.TestCase): + def test_box(self): + b1 = Box(1.0, 2.0, 1.0, 2.0) + b2 = Box(2.0, 3.0, 1.0, 2.0) + b3 = Box(1.0, 2.0, 2.0, 2.0) + iou_1_2 = b1.calc_iou(b2) + iou_1_3 = b1.calc_iou(b3) + + self.assertAlmostEqual(iou_1_2, 0.0) + self.assertAlmostEqual(iou_1_3, 0.5) +