96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
#!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)
|