Initial commit

2026-05-03 23:47:54 +03:00
commit 1dac1749bf
74 changed files with 6697 additions and 0 deletions
@@ -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)