miner.py

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