miner.py

import os
import site
import tempfile

# 1. Get the path to the current environment's site-packages
site_packages_path = site.getsitepackages()[0]
target_dir = os.path.join(site_packages_path, "torch_pip_tmp")

# 2. Create the target directory if it doesn't exist
os.makedirs(target_dir, exist_ok=True)

# 3. Redirect Pip & System Temp Folders
# TMPDIR is used by Linux/macOS; TEMP/TMP are used by Windows
os.environ['TMPDIR'] = target_dir
os.environ['TEMP'] = target_dir
os.environ['TMP'] = target_dir
os.environ['PIP_CACHE_DIR'] = target_dir # Specifically for pip cache

# 4. Redirect Torch and Torchvision Caches
# TORCH_HOME controls where models and hub resources are saved
os.environ['TORCH_HOME'] = target_dir

# 5. Tell the Python tempfile module to use this directory for the current session
tempfile.tempdir = target_dir #

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 turbovision import Turbo

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:
        self.turbo = Turbo(path_hf_repo)

    def __repr__(self) -> str:
        return f"Miner(turbo={self.turbo})"

    def predict_batch(self, batch_images: List[ndarray], offset: int, n_keypoints: int) -> List[TVFrameResult]: 
        return self.turbo.predict_batch(batch_images, offset, n_keypoints)