scorevision: push artifact

miner.py

@@ -0,0 +1,90 @@
+import os
+import sys
+from pathlib import Path
+from typing import List, Tuple
+import sys
+import os
+
+from numpy import ndarray
+from pydantic import BaseModel
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+# from inference import predict_batch, load_model
+import importlib.util
+from pathlib import Path
+
+
+def manual_import(name, filename):
+    """
+    Manually loads a module (.so, .pyc, or .py) from a specific file path,
+    bypassing sys.meta_path import hooks.
+    """
+    # Locate the file relative to the current miner.py
+    curr_dir = Path(__file__).parent
+    file_path = curr_dir / filename
+
+    if not file_path.exists():
+        raise FileNotFoundError(f"Could not find {file_path}")
+
+    # Load the spec directly from the file path
+    spec = importlib.util.spec_from_file_location(name, file_path)
+    if spec is None:
+        raise ImportError(f"Could not load spec for {name} from {file_path}")
+
+    # Create the module and register it in sys.modules
+    module = importlib.util.module_from_spec(spec)
+
+    # Execute the module
+    spec.loader.exec_module(module)
+    return module
+
+
+# import inference
+
+
+class BoundingBox(BaseModel):
+    x1: int
+    y1: int
+    x2: int
+    y2: int
+    cls_id: int
+    conf: float
+
+
+class TVFrameResult(BaseModel):
+    frame_id: int
+    boxes: List[BoundingBox]
+    keypoints: List[Tuple[int, int]]
+
+
+class Miner:
+    def __init__(self, path_hf_repo):
+        print("model loading")
+        self.health = "Okay"
+        self.inference = None
+        self.path_hf_repo = path_hf_repo
+        self.is_start = False
+
+    def __repr__(self) -> str:
+        return self.health
+
+    def predict_batch(self, batch_images, offset, n_keypoints):
+        if self.is_start == False:
+            self.is_start = True
+            return None
+        if self.inference is None:
+            print("importing inference module")
+            # self.inference = manual_import("inference", "inference.py")
+            self.inference = manual_import(
+                "inference", "inference.cpython-312-x86_64-linux-gnu.so"
+            )
+            print("inference module loading...")
+            self.inference.load_model(self.path_hf_repo)
+            print("inference module imported")
+        results = self.inference.predict_batch(
+            batch_images,
+            offset,
+            n_keypoints,
+        )
+        return results