miner.py

import os
import sys
from pathlib import Path
from typing import List, Tuple, Dict
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)
    sys.modules[name] = module
    
    # Execute the module
    spec.loader.exec_module(module)
    return module

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: Path) -> None:
        print("model laoding")
        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: List[ndarray], offset: int, n_keypoints: int) -> List[TVFrameResult]:
        if self.is_start == False:
            self.is_start = True
            return None
        if self.inference is None:
            self.inference = manual_import("inference", "inference.cpython-312-x86_64-linux-gnu.so")
            self.inference.load_model(self.path_hf_repo)
            
        results = self.inference.predict_batch(
            batch_images,
            offset,
            n_keypoints,
        )
        return results