Spaces:
Paused
Paused
File size: 638 Bytes
cf54850 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | from __future__ import annotations
import os
from src.engines.sstgnn.engine import SSTGNNEngine
from src.services.media_utils import extract_video_frames
class SSTGNNModule:
def __init__(self, cache_dir: str = "/data/model_cache"):
os.environ.setdefault("MODEL_CACHE_DIR", cache_dir)
self.engine = SSTGNNEngine()
def score(self, video_path: str) -> dict:
frames = extract_video_frames(video_path, max_frames=60)
if not frames:
return {"s3": 0.5, "vram_mb": 0}
result = self.engine.run_video(frames)
return {"s3": round(float(result.confidence), 4), "vram_mb": 0}
|