Create vision.py

vision.py

@@ -0,0 +1,37 @@
+import json
+from pathlib import Path
+from typing import Dict
+from PIL import Image
+import torch
+from transformers import BlipProcessor, BlipForConditionalGeneration
+from config import VISION_MODEL
+
+_processor = None
+_model = None
+_device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def _load_blip():
+    global _processor, _model
+    if _processor is None or _model is None:
+        _processor = BlipProcessor.from_pretrained(VISION_MODEL)
+        _model = BlipForConditionalGeneration.from_pretrained(VISION_MODEL).to(_device)
+        _model.eval()
+    return _processor, _model
+
+def caption_image(img_path: Path) -> str:
+    processor, model = _load_blip()
+    img = Image.open(str(img_path)).convert("RGB")
+    inputs = processor(img, return_tensors="pt").to(_device)
+    with torch.inference_mode():
+        out_ids = model.generate(**inputs, max_new_tokens=40)
+    return processor.decode(out_ids[0], skip_special_tokens=True)
+
+def caption_folder(frames_dir: Path) -> Dict[str, str]:
+    results = {}
+    for p in sorted(frames_dir.glob("*.jpg")):
+        results[p.name] = caption_image(p)
+    return results
+
+def dump_json(data, out_path: Path):
+    with open(out_path, "w", encoding="utf-8") as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)