switch to VideoMAE archit11 model

app.py

@@ -18,15 +18,16 @@ import cv2
 import gradio as gr
 import requests
 from dotenv import load_dotenv
-from transformers import AutoProcessor, AutoModel
+from transformers import VideoMAEImageProcessor, VideoMAEForVideoClassification
 
 load_dotenv()
 
-MODEL_ID = "microsoft/xclip-base-patch32"
-LABELS = ["people fighting", "people walking normally", "people talking calmly"]
-ALERT_LABEL = "people fighting"
-ALERT_THRESHOLD = 0.55
-NUM_FRAMES = 8
+MODEL_ID = os.environ.get(
+    "MODEL_ID",
+    "archit11/videomae-base-finetuned-fight-nofight-subset2",
+)
+ALERT_THRESHOLD = float(os.environ.get("ALERT_THRESHOLD", "0.60"))
+NUM_FRAMES = 16
 DB_PATH = Path("events.db")
 ALERT_COOLDOWN_SEC = 10
 
@@ -34,8 +35,23 @@ TELEGRAM_TOKEN = os.environ.get("TELEGRAM_TOKEN")
 TELEGRAM_CHAT_ID = os.environ.get("TELEGRAM_CHAT_ID")
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-processor = AutoProcessor.from_pretrained(MODEL_ID)
-model = AutoModel.from_pretrained(MODEL_ID).to(device).eval()
+processor = VideoMAEImageProcessor.from_pretrained(MODEL_ID)
+model = VideoMAEForVideoClassification.from_pretrained(MODEL_ID).to(device).eval()
+
+ID2LABEL = {int(k): v for k, v in model.config.id2label.items()}
+
+
+def _find_alert_index():
+    for idx, label in ID2LABEL.items():
+        norm = label.lower().replace(" ", "").replace("_", "").replace("-", "")
+        if norm in ("fight", "violence"):
+            return idx
+    return 0
+
+
+ALERT_IDX = _find_alert_index()
+ALERT_LABEL = ID2LABEL[ALERT_IDX]
+print(f"[config] model={MODEL_ID}  labels={ID2LABEL}  alert_on={ALERT_LABEL}")
 
 _last_alert_ts = 0.0
 _alert_lock = threading.Lock()
@@ -109,16 +125,11 @@ def should_alert(confidence):
 
 
 def run_inference(frames):
-    inputs = processor(
-        text=LABELS,
-        videos=list(frames),
-        return_tensors="pt",
-        padding=True,
-    ).to(device)
+    inputs = processor(list(frames), return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = model(**inputs)
-        probs = outputs.logits_per_video.softmax(dim=-1)[0].cpu().numpy()
-    return {label: float(p) for label, p in zip(LABELS, probs)}
+        probs = outputs.logits.softmax(dim=-1)[0].cpu().numpy()
+    return {ID2LABEL[i]: float(p) for i, p in enumerate(probs)}
 
 
 def trigger_alert(confidence, source):