Initial commit
This commit is contained in:
Mustafa KURU
@@ -0,0 +1,25 @@
|
||||
import os
|
||||
from functools import wraps
|
||||
|
||||
from flask import Response, request
|
||||
|
||||
ML_API_TOKEN = os.environ.get("ML_API_TOKEN")
|
||||
|
||||
|
||||
def token_required(f):
|
||||
@wraps(f)
|
||||
def check_authorization(*args, **kwargs):
|
||||
if (
|
||||
request.headers.get("Authorization") == f"Bearer {ML_API_TOKEN}"
|
||||
or request.args.get("token") == ML_API_TOKEN
|
||||
):
|
||||
return f(*args, **kwargs)
|
||||
return Response(status=401)
|
||||
|
||||
@wraps(f)
|
||||
def passthru(*args, **kwargs):
|
||||
return f(*args, **kwargs)
|
||||
|
||||
if ML_API_TOKEN:
|
||||
return check_authorization
|
||||
return passthru
|
||||
@@ -0,0 +1,254 @@
|
||||
# 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
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
#!python3
|
||||
|
||||
# pylint: disable=R, W0401, W0614, W0703
|
||||
from lib.meta import Meta
|
||||
from os import environ, path
|
||||
|
||||
alt_names = None
|
||||
|
||||
darknet_ready = True
|
||||
try:
|
||||
from lib.darknet import YoloNet
|
||||
except Exception as e:
|
||||
print(f'Error during importing YoloNet! - {e}')
|
||||
darknet_ready = False
|
||||
|
||||
onnx_ready = True
|
||||
try:
|
||||
from lib.onnx import OnnxNet
|
||||
except Exception as e:
|
||||
print(f'Error during importing OnnxNet! - {e}')
|
||||
onnx_ready = False
|
||||
|
||||
|
||||
def load_net(config_path, meta_path, weights_path=None):
|
||||
|
||||
def try_loading_net(net_config_priority):
|
||||
for net_config in net_config_priority:
|
||||
weights = net_config['weights_path']
|
||||
use_gpu = net_config['use_gpu']
|
||||
|
||||
net_main = None
|
||||
try:
|
||||
print(f'----- Trying to load weights: {weights} - use_gpu = {use_gpu} -----')
|
||||
if weights.endswith(".onnx"):
|
||||
if not onnx_ready:
|
||||
raise Exception('Not loading ONNX net due to previous import failure. Check earlier log for errors.')
|
||||
net_main = OnnxNet(weights, meta_path, use_gpu)
|
||||
|
||||
elif weights.endswith(".darknet"):
|
||||
if not darknet_ready:
|
||||
raise Exception('Not loading darknet net due to previous import failure. Check earlier log for errors.')
|
||||
net_main = YoloNet(weights, meta_path, config_path, use_gpu)
|
||||
|
||||
else:
|
||||
raise Exception(f'Can not recognize net from weights file surfix: {weights}')
|
||||
|
||||
print('Succeeded!')
|
||||
return net_main
|
||||
except Exception as e:
|
||||
print(f'Failed! - {e}')
|
||||
|
||||
raise Exception(f'Failed to load any net after trying: {net_config_priority}')
|
||||
|
||||
global alt_names # pylint: disable=W0603
|
||||
|
||||
model_dir = path.join(path.dirname(path.realpath(__file__)), '..', 'model')
|
||||
use_gpu = environ.get('ML_USE_GPU', 'false').lower() in ('1', 'true', 'yes', 'on')
|
||||
preferred_backend = environ.get('ML_MODEL_BACKEND', 'onnx').lower()
|
||||
|
||||
cpu_priority = [
|
||||
dict(weights_path=path.join(model_dir, 'model-weights.onnx'), use_gpu=False),
|
||||
dict(weights_path=path.join(model_dir, 'model-weights.darknet'), use_gpu=False),
|
||||
]
|
||||
gpu_priority = [
|
||||
dict(weights_path=path.join(model_dir, 'model-weights.onnx'), use_gpu=True),
|
||||
dict(weights_path=path.join(model_dir, 'model-weights.darknet'), use_gpu=True),
|
||||
]
|
||||
|
||||
if preferred_backend == 'darknet':
|
||||
cpu_priority.reverse()
|
||||
gpu_priority.reverse()
|
||||
|
||||
net_config_priority = gpu_priority + cpu_priority if use_gpu else cpu_priority
|
||||
if weights_path is not None:
|
||||
net_config_priority = (
|
||||
[dict(weights_path=weights_path, use_gpu=True), dict(weights_path=weights_path, use_gpu=False)]
|
||||
if use_gpu
|
||||
else [dict(weights_path=weights_path, use_gpu=False)]
|
||||
)
|
||||
|
||||
net_main = try_loading_net(net_config_priority)
|
||||
|
||||
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:
|
||||
meta = Meta(meta_path)
|
||||
alt_names = meta.names
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return net_main
|
||||
|
||||
def detect(net, image, thresh=.5, hier_thresh=.5, nms=.45, debug=False):
|
||||
return net.detect(net.meta, image, alt_names, thresh, hier_thresh, nms, debug)
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -0,0 +1,132 @@
|
||||
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, use_gpu: bool):
|
||||
providers = ['CUDAExecutionProvider'] if use_gpu else ['CPUExecutionProvider']
|
||||
self.session = onnxruntime.InferenceSession(onnx_path, providers=providers)
|
||||
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 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
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
https://tsd-pub-static.s3.amazonaws.com/ml-models/model-weights-8be06cde4e.darknet
|
||||
@@ -0,0 +1 @@
|
||||
https://tsd-pub-static.s3.amazonaws.com/ml-models/model-weights-5a6b1be1fa.onnx
|
||||
@@ -0,0 +1,258 @@
|
||||
[net]
|
||||
# Testing
|
||||
batch=64
|
||||
subdivisions=8
|
||||
# Training
|
||||
# batch=64
|
||||
# subdivisions=8
|
||||
height=416
|
||||
width=416
|
||||
channels=3
|
||||
momentum=0.9
|
||||
decay=0.0005
|
||||
angle=0
|
||||
saturation = 1.5
|
||||
exposure = 1.5
|
||||
hue=.1
|
||||
|
||||
learning_rate=0.001
|
||||
burn_in=1000
|
||||
max_batches = 50000
|
||||
policy=steps
|
||||
steps=40000,60000
|
||||
scales=.1,.1
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=32
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[maxpool]
|
||||
size=2
|
||||
stride=2
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=64
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[maxpool]
|
||||
size=2
|
||||
stride=2
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=128
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=64
|
||||
size=1
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=128
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[maxpool]
|
||||
size=2
|
||||
stride=2
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=256
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=128
|
||||
size=1
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=256
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[maxpool]
|
||||
size=2
|
||||
stride=2
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=512
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=256
|
||||
size=1
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=512
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=256
|
||||
size=1
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=512
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[maxpool]
|
||||
size=2
|
||||
stride=2
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=1024
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=512
|
||||
size=1
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=1024
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=512
|
||||
size=1
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
filters=1024
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
activation=leaky
|
||||
|
||||
|
||||
#######
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
filters=1024
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
filters=1024
|
||||
activation=leaky
|
||||
|
||||
[route]
|
||||
layers=-9
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
size=1
|
||||
stride=1
|
||||
pad=1
|
||||
filters=64
|
||||
activation=leaky
|
||||
|
||||
[reorg3d]
|
||||
stride=2
|
||||
|
||||
[route]
|
||||
layers=-1,-4
|
||||
|
||||
[convolutional]
|
||||
batch_normalize=1
|
||||
size=3
|
||||
stride=1
|
||||
pad=1
|
||||
filters=1024
|
||||
activation=leaky
|
||||
|
||||
[convolutional]
|
||||
size=1
|
||||
stride=1
|
||||
pad=1
|
||||
filters=30
|
||||
activation=linear
|
||||
|
||||
|
||||
[region]
|
||||
anchors = 1.3221, 1.73145, 3.19275, 4.00944, 5.05587, 8.09892, 9.47112, 4.84053, 11.2364, 10.0071
|
||||
bias_match=1
|
||||
classes=1
|
||||
coords=4
|
||||
num=5
|
||||
softmax=1
|
||||
jitter=.3
|
||||
rescore=1
|
||||
|
||||
object_scale=5
|
||||
noobject_scale=1
|
||||
class_scale=1
|
||||
coord_scale=1
|
||||
|
||||
absolute=1
|
||||
thresh = .6
|
||||
random=1
|
||||
@@ -0,0 +1,2 @@
|
||||
classes= 1
|
||||
names = /app/model/names
|
||||
@@ -0,0 +1 @@
|
||||
failure
|
||||
@@ -0,0 +1,6 @@
|
||||
ipdb
|
||||
flask>=1.0
|
||||
redis==3.0.1
|
||||
newrelic==4.12.0.113
|
||||
requests==2.21.0
|
||||
gunicorn==19.9.0
|
||||
@@ -0,0 +1,275 @@
|
||||
#!/usr/bin/env python
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from datetime import datetime, timezone
|
||||
import logging
|
||||
from os import environ, path
|
||||
from time import perf_counter
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
import cv2
|
||||
import flask
|
||||
from flask import Response, jsonify, request
|
||||
import numpy as np
|
||||
import requests
|
||||
|
||||
from auth import token_required
|
||||
from lib.detection_model import detect, load_net
|
||||
|
||||
THRESH = float(environ.get("ML_DETECTION_BOX_THRESHOLD", "0.08"))
|
||||
REQUEST_TIMEOUT = (
|
||||
float(environ.get("ML_IMAGE_CONNECT_TIMEOUT", "2")),
|
||||
float(environ.get("ML_IMAGE_READ_TIMEOUT", "10")),
|
||||
)
|
||||
MAX_RECENT_REQUESTS = int(environ.get("ML_RECENT_REQUESTS", "100"))
|
||||
|
||||
app = flask.Flask(__name__)
|
||||
app.config["DEBUG"] = environ.get("DEBUG") == "True"
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(levelname)s %(name)s %(message)s",
|
||||
)
|
||||
app.logger.setLevel(logging.INFO)
|
||||
|
||||
STARTED_AT = datetime.now(timezone.utc)
|
||||
RECENT_REQUESTS: deque[dict] = deque(maxlen=MAX_RECENT_REQUESTS)
|
||||
|
||||
model_dir = path.join(path.dirname(path.realpath(__file__)), "model")
|
||||
net_main = load_net(path.join(model_dir, "model.cfg"), path.join(model_dir, "model.meta"))
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _redact_url(raw_url: str | None) -> str | None:
|
||||
if not raw_url:
|
||||
return None
|
||||
parsed = urlsplit(raw_url)
|
||||
return urlunsplit((parsed.scheme, parsed.netloc, parsed.path, "", ""))
|
||||
|
||||
|
||||
def _record_request(entry: dict) -> None:
|
||||
stored_entry = dict(entry)
|
||||
if "image_url" in stored_entry:
|
||||
stored_entry["image_url"] = _redact_url(stored_entry.get("image_url"))
|
||||
RECENT_REQUESTS.appendleft({"time": _now_iso(), **stored_entry})
|
||||
app.logger.info(
|
||||
"prediction status=%s detections=%s duration_ms=%s image_host=%s error=%s",
|
||||
entry.get("status"),
|
||||
entry.get("detections", 0),
|
||||
entry.get("duration_ms"),
|
||||
urlsplit(entry.get("image_url") or "").netloc,
|
||||
entry.get("error"),
|
||||
)
|
||||
|
||||
|
||||
def _fetch_image(image_url: str) -> np.ndarray:
|
||||
response = requests.get(image_url, stream=True, timeout=REQUEST_TIMEOUT)
|
||||
response.raise_for_status()
|
||||
img_array = np.array(bytearray(response.content), dtype=np.uint8)
|
||||
image = cv2.imdecode(img_array, -1)
|
||||
if image is None:
|
||||
raise ValueError("image_decode_failed")
|
||||
return image
|
||||
|
||||
|
||||
def _status_payload() -> dict:
|
||||
return {
|
||||
"ok": net_main is not None,
|
||||
"started_at": STARTED_AT.isoformat(),
|
||||
"model": {
|
||||
"classes": ["failure"],
|
||||
"box_threshold": THRESH,
|
||||
"backend": type(net_main).__name__ if net_main is not None else None,
|
||||
"use_gpu": environ.get("ML_USE_GPU", "false"),
|
||||
"model_backend_preference": environ.get("ML_MODEL_BACKEND", "onnx"),
|
||||
},
|
||||
"requests": {
|
||||
"recent_count": len(RECENT_REQUESTS),
|
||||
"max_recent": MAX_RECENT_REQUESTS,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@app.route("/", methods=["GET"])
|
||||
def dashboard():
|
||||
"""Render a small operational status page."""
|
||||
status = _status_payload()
|
||||
rows = "\n".join(
|
||||
"<tr>"
|
||||
f"<td>{entry['time']}</td>"
|
||||
f"<td>{entry.get('status', '')}</td>"
|
||||
f"<td>{entry.get('detections', 0)}</td>"
|
||||
f"<td>{entry.get('duration_ms', '')}</td>"
|
||||
f"<td>{entry.get('error') or ''}</td>"
|
||||
f"<td>{_redact_url(entry.get('image_url')) or ''}</td>"
|
||||
"</tr>"
|
||||
for entry in list(RECENT_REQUESTS)[:20]
|
||||
)
|
||||
body = f"""<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>Elegoo Spaghetti Detection ML Server</title>
|
||||
<style>
|
||||
body {{ font-family: system-ui, sans-serif; margin: 24px; color: #1f2933; }}
|
||||
code {{ background: #eef2f7; padding: 2px 5px; border-radius: 4px; }}
|
||||
table {{ border-collapse: collapse; width: 100%; margin-top: 16px; }}
|
||||
th, td {{ border-bottom: 1px solid #d9e2ec; padding: 8px; text-align: left; font-size: 14px; }}
|
||||
.ok {{ color: #137333; font-weight: 700; }}
|
||||
.bad {{ color: #b3261e; font-weight: 700; }}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>Elegoo Spaghetti Detection ML Server</h1>
|
||||
<p>Status: <span class="{'ok' if status['ok'] else 'bad'}">{'ok' if status['ok'] else 'error'}</span></p>
|
||||
<p>Backend: <code>{status['model']['backend']}</code> | GPU opt-in: <code>{status['model']['use_gpu']}</code> | Box threshold: <code>{status['model']['box_threshold']}</code></p>
|
||||
<p>Health: <code>/hc/</code> | JSON status: <code>/api/status</code> | Token-protected logs: <code>/api/logs?token=<token></code></p>
|
||||
<h2>Recent Requests</h2>
|
||||
<table>
|
||||
<thead><tr><th>Time</th><th>Status</th><th>Detections</th><th>ms</th><th>Error</th><th>Image URL without token</th></tr></thead>
|
||||
<tbody>{rows}</tbody>
|
||||
</table>
|
||||
</body>
|
||||
</html>"""
|
||||
return Response(body, mimetype="text/html")
|
||||
|
||||
|
||||
@app.route("/api/status", methods=["GET"])
|
||||
def api_status():
|
||||
"""Return JSON server status."""
|
||||
return jsonify(_status_payload())
|
||||
|
||||
|
||||
@app.route("/api/logs", methods=["GET"])
|
||||
@token_required
|
||||
def api_logs():
|
||||
"""Return recent request logs."""
|
||||
return jsonify({"requests": list(RECENT_REQUESTS)})
|
||||
|
||||
|
||||
@app.route("/debug/image", methods=["GET"])
|
||||
@token_required
|
||||
def debug_image():
|
||||
"""Check whether the server can fetch and decode an image URL."""
|
||||
image_url = request.args.get("img")
|
||||
if not image_url:
|
||||
return jsonify({"ok": False, "error": "missing_image_url"}), 400
|
||||
started = perf_counter()
|
||||
try:
|
||||
image = _fetch_image(image_url)
|
||||
return jsonify(
|
||||
{
|
||||
"ok": True,
|
||||
"duration_ms": round((perf_counter() - started) * 1000),
|
||||
"shape": list(image.shape),
|
||||
"image_url": _redact_url(image_url),
|
||||
}
|
||||
)
|
||||
except requests.RequestException as err:
|
||||
return jsonify({"ok": False, "error": "image_fetch_failed", "message": str(err)}), 502
|
||||
except ValueError as err:
|
||||
return jsonify({"ok": False, "error": str(err)}), 422
|
||||
|
||||
|
||||
@app.route("/p/", methods=["GET"])
|
||||
@token_required
|
||||
def get_p():
|
||||
"""Run prediction for the image URL in the img query parameter."""
|
||||
image_url = request.args.get("img")
|
||||
if not image_url:
|
||||
_record_request({"status": 400, "error": "missing_image_url", "detections": 0})
|
||||
return jsonify(
|
||||
{
|
||||
"detections": [],
|
||||
"error": "missing_image_url",
|
||||
"message": "Missing img query parameter.",
|
||||
}
|
||||
), 400
|
||||
|
||||
started = perf_counter()
|
||||
try:
|
||||
image = _fetch_image(image_url)
|
||||
detections = detect(net_main, image, thresh=THRESH)
|
||||
duration_ms = round((perf_counter() - started) * 1000)
|
||||
_record_request(
|
||||
{
|
||||
"status": 200,
|
||||
"detections": len(detections),
|
||||
"duration_ms": duration_ms,
|
||||
"image_url": image_url,
|
||||
}
|
||||
)
|
||||
return jsonify({"detections": detections, "duration_ms": duration_ms})
|
||||
except requests.RequestException as err:
|
||||
duration_ms = round((perf_counter() - started) * 1000)
|
||||
_record_request(
|
||||
{
|
||||
"status": 502,
|
||||
"error": "image_fetch_failed",
|
||||
"message": str(err),
|
||||
"duration_ms": duration_ms,
|
||||
"image_url": image_url,
|
||||
"detections": 0,
|
||||
}
|
||||
)
|
||||
return jsonify(
|
||||
{
|
||||
"detections": [],
|
||||
"error": "image_fetch_failed",
|
||||
"message": str(err),
|
||||
}
|
||||
), 502
|
||||
except ValueError as err:
|
||||
duration_ms = round((perf_counter() - started) * 1000)
|
||||
_record_request(
|
||||
{
|
||||
"status": 422,
|
||||
"error": str(err),
|
||||
"duration_ms": duration_ms,
|
||||
"image_url": image_url,
|
||||
"detections": 0,
|
||||
}
|
||||
)
|
||||
return jsonify(
|
||||
{
|
||||
"detections": [],
|
||||
"error": str(err),
|
||||
"message": "The image URL did not return a decodable image.",
|
||||
}
|
||||
), 422
|
||||
except Exception as err:
|
||||
duration_ms = round((perf_counter() - started) * 1000)
|
||||
app.logger.exception("Unable to process image")
|
||||
_record_request(
|
||||
{
|
||||
"status": 500,
|
||||
"error": "prediction_failed",
|
||||
"message": str(err),
|
||||
"duration_ms": duration_ms,
|
||||
"image_url": image_url,
|
||||
"detections": 0,
|
||||
}
|
||||
)
|
||||
return jsonify(
|
||||
{
|
||||
"detections": [],
|
||||
"error": "prediction_failed",
|
||||
"message": str(err),
|
||||
}
|
||||
), 500
|
||||
|
||||
|
||||
@app.route("/hc/", methods=["GET"])
|
||||
def health_check():
|
||||
"""Health check for Home Assistant and Docker."""
|
||||
if net_main is not None:
|
||||
return "ok", 200
|
||||
return "error", 503
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app.run(host="0.0.0.0", port=3333, threaded=False)
|
||||
@@ -0,0 +1,6 @@
|
||||
import server
|
||||
|
||||
application = server.app
|
||||
|
||||
if __name__ == "__main__":
|
||||
application.run()
|
||||
@@ -0,0 +1,17 @@
|
||||
#!/usr/bin/with-contenv bashio
|
||||
|
||||
declare ML_API_TOKEN
|
||||
|
||||
ML_API_TOKEN=$(bashio::config 'obico_api_secret')
|
||||
export ML_API_TOKEN
|
||||
export ML_USE_GPU=$(bashio::config 'use_gpu')
|
||||
export GUNICORN_TIMEOUT=$(bashio::config 'gunicorn_timeout')
|
||||
|
||||
cd /app
|
||||
FLASK_APP=server.py venv/bin/gunicorn \
|
||||
--bind "0.0.0.0:3333" \
|
||||
--workers "${GUNICORN_WORKERS:-1}" \
|
||||
--timeout "${GUNICORN_TIMEOUT:-120}" \
|
||||
--error-logfile - \
|
||||
--log-level info \
|
||||
wsgi:application
|
||||
Reference in New Issue
Block a user