336 lines
12 KiB
Python
336 lines
12 KiB
Python
|
import base64
|
||
|
|
import copy
|
||
|
|
import json
|
||
|
|
import re
|
||
|
|
import subprocess
|
||
|
|
import tempfile
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
|
||
|
|
from app.config import settings
|
||
|
|
from app.profile_store import Profile
|
||
|
|
|
||
|
|
|
||
|
|
class VisionModelError(RuntimeError):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
VIDEO_RESULT_SCHEMA = {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"passed": {"type": ["boolean", "null"]},
|
||
|
|
"conclusion": {"type": "string"},
|
||
|
|
"summary": {"type": "string"},
|
||
|
|
"events": {
|
||
|
|
"type": "array",
|
||
|
|
"items": {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"timeRange": {"type": "string"},
|
||
|
|
"event": {"type": "string"},
|
||
|
|
"confidence": {"type": "number"},
|
||
|
|
},
|
||
|
|
"required": ["timeRange", "event", "confidence"],
|
||
|
|
},
|
||
|
|
},
|
||
|
|
"warnings": {"type": "array", "items": {"type": "string"}},
|
||
|
|
},
|
||
|
|
"required": ["passed", "conclusion", "summary", "events", "warnings"],
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _duration(path: Path) -> float:
|
||
|
|
process = subprocess.run(
|
||
|
|
[
|
||
|
|
"ffprobe",
|
||
|
|
"-v",
|
||
|
|
"error",
|
||
|
|
"-show_entries",
|
||
|
|
"format=duration",
|
||
|
|
"-of",
|
||
|
|
"default=noprint_wrappers=1:nokey=1",
|
||
|
|
str(path),
|
||
|
|
],
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
check=False,
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
value = float(process.stdout.strip())
|
||
|
|
except ValueError as exc:
|
||
|
|
raise ValueError("Unable to read video duration") from exc
|
||
|
|
if value <= 0:
|
||
|
|
raise ValueError("Video duration must be positive")
|
||
|
|
return value
|
||
|
|
|
||
|
|
|
||
|
|
def _extract_frames(
|
||
|
|
path: Path, maximum: int, maximum_width: int
|
||
|
|
) -> tuple[list[Path], float, Path]:
|
||
|
|
duration = _duration(path)
|
||
|
|
count = max(2, min(maximum, int(duration) + 1))
|
||
|
|
fps = count / duration
|
||
|
|
directory = Path(tempfile.mkdtemp(prefix="frames-", dir=settings.jobs_dir))
|
||
|
|
output = directory / "frame-%03d.jpg"
|
||
|
|
process = subprocess.run(
|
||
|
|
[
|
||
|
|
"ffmpeg",
|
||
|
|
"-v",
|
||
|
|
"error",
|
||
|
|
"-i",
|
||
|
|
str(path),
|
||
|
|
"-vf",
|
||
|
|
f"fps={fps:.8f},scale='min({maximum_width},iw)':-2",
|
||
|
|
"-frames:v",
|
||
|
|
str(count),
|
||
|
|
"-q:v",
|
||
|
|
"3",
|
||
|
|
str(output),
|
||
|
|
],
|
||
|
|
capture_output=True,
|
||
|
|
check=False,
|
||
|
|
)
|
||
|
|
frames = sorted(directory.glob("frame-*.jpg"))
|
||
|
|
if process.returncode != 0 or not frames:
|
||
|
|
error = process.stderr.decode("utf-8", errors="replace")[-500:]
|
||
|
|
raise ValueError(f"Unable to extract video frames: {error}")
|
||
|
|
return frames, duration, directory
|
||
|
|
|
||
|
|
|
||
|
|
def _select_frames(frames: list[Path], maximum: int) -> list[Path]:
|
||
|
|
if len(frames) <= maximum:
|
||
|
|
return frames
|
||
|
|
if maximum <= 1:
|
||
|
|
return [frames[len(frames) // 2]]
|
||
|
|
indices = [round(index * (len(frames) - 1) / (maximum - 1)) for index in range(maximum)]
|
||
|
|
return [frames[index] for index in indices]
|
||
|
|
|
||
|
|
|
||
|
|
def _encode_image(path: Path, maximum_width: int, extracted_width: int) -> str:
|
||
|
|
if maximum_width >= extracted_width:
|
||
|
|
image = path.read_bytes()
|
||
|
|
else:
|
||
|
|
process = subprocess.run(
|
||
|
|
[
|
||
|
|
"ffmpeg",
|
||
|
|
"-v",
|
||
|
|
"error",
|
||
|
|
"-i",
|
||
|
|
str(path),
|
||
|
|
"-vf",
|
||
|
|
f"scale='min({maximum_width},iw)':-2",
|
||
|
|
"-q:v",
|
||
|
|
"3",
|
||
|
|
"-f",
|
||
|
|
"image2pipe",
|
||
|
|
"-vcodec",
|
||
|
|
"mjpeg",
|
||
|
|
"pipe:1",
|
||
|
|
],
|
||
|
|
capture_output=True,
|
||
|
|
check=False,
|
||
|
|
)
|
||
|
|
if process.returncode != 0 or not process.stdout:
|
||
|
|
error = process.stderr.decode("utf-8", errors="replace")[-500:]
|
||
|
|
raise ValueError(f"Unable to resize video frame: {error}")
|
||
|
|
image = process.stdout
|
||
|
|
return base64.b64encode(image).decode("ascii")
|
||
|
|
|
||
|
|
|
||
|
|
def _parse_json(content: str) -> dict[str, Any]:
|
||
|
|
content = content.strip()
|
||
|
|
fenced = re.search(r"```(?:json)?\s*(.*?)\s*```", content, re.DOTALL)
|
||
|
|
if fenced:
|
||
|
|
content = fenced.group(1)
|
||
|
|
try:
|
||
|
|
parsed = json.loads(content)
|
||
|
|
return parsed if isinstance(parsed, dict) else {"result": parsed, "structured": False}
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
return {"summary": content, "structured": False}
|
||
|
|
|
||
|
|
|
||
|
|
def _fallback_reason(result: dict[str, Any], fallback_on_unknown: bool) -> str | None:
|
||
|
|
if result.get("structured") is False:
|
||
|
|
return "FAST_RESULT_NOT_STRUCTURED"
|
||
|
|
required = ("passed", "conclusion", "summary", "events", "warnings")
|
||
|
|
missing = [field for field in required if field not in result]
|
||
|
|
if missing:
|
||
|
|
return "FAST_RESULT_MISSING_FIELDS:" + ",".join(missing)
|
||
|
|
if not isinstance(result.get("conclusion"), str) or not result["conclusion"].strip():
|
||
|
|
return "FAST_RESULT_EMPTY_CONCLUSION"
|
||
|
|
if not isinstance(result.get("summary"), str) or not result["summary"].strip():
|
||
|
|
return "FAST_RESULT_EMPTY_SUMMARY"
|
||
|
|
if not isinstance(result.get("events"), list) or not isinstance(result.get("warnings"), list):
|
||
|
|
return "FAST_RESULT_INVALID_COLLECTIONS"
|
||
|
|
if fallback_on_unknown and result.get("passed") is None:
|
||
|
|
return "FAST_RESULT_INSUFFICIENT"
|
||
|
|
if result.get("passed") is not True and result.get("passed") is not False and result.get("passed") is not None:
|
||
|
|
return "FAST_RESULT_INVALID_DECISION"
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _strategy(profile: Profile, name: str) -> dict[str, Any]:
|
||
|
|
defaults = {
|
||
|
|
"FAST": {
|
||
|
|
"model": profile.config.get("model", settings.ollama_model),
|
||
|
|
"maxFrames": 6,
|
||
|
|
"maxWidth": 896,
|
||
|
|
"maxOutputTokens": 512,
|
||
|
|
"numCtx": 16384,
|
||
|
|
"emptyResponseRetries": 0,
|
||
|
|
},
|
||
|
|
"ACCURATE": {
|
||
|
|
"model": profile.config.get("model", settings.ollama_model),
|
||
|
|
"maxFrames": 12,
|
||
|
|
"maxWidth": 1280,
|
||
|
|
"maxOutputTokens": 768,
|
||
|
|
"numCtx": 32768,
|
||
|
|
"emptyResponseRetries": 1,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
configured = profile.config.get("strategies", {}).get(name.lower(), {})
|
||
|
|
return {**defaults[name], **configured}
|
||
|
|
|
||
|
|
|
||
|
|
async def _request_vision_model(
|
||
|
|
payload: dict[str, Any], empty_response_retries: int = 1
|
||
|
|
) -> tuple[str, dict[str, Any]]:
|
||
|
|
last_metadata: dict[str, Any] = {}
|
||
|
|
timeout = httpx.Timeout(settings.ollama_timeout_seconds)
|
||
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
||
|
|
for attempt in range(empty_response_retries + 1):
|
||
|
|
request_payload = copy.deepcopy(payload)
|
||
|
|
if attempt:
|
||
|
|
options = request_payload.setdefault("options", {})
|
||
|
|
options["num_predict"] = max(int(options.get("num_predict", 0)), 1536)
|
||
|
|
request_payload["messages"][0]["content"] += (
|
||
|
|
"\n\nReturn the final JSON now. Do not return reasoning without a final answer."
|
||
|
|
)
|
||
|
|
response = await client.post(
|
||
|
|
f"{settings.ollama_base_url.rstrip('/')}/api/chat", json=request_payload
|
||
|
|
)
|
||
|
|
response.raise_for_status()
|
||
|
|
body = response.json()
|
||
|
|
message = body.get("message") or {}
|
||
|
|
content = str(message.get("content") or "").strip()
|
||
|
|
metadata = {
|
||
|
|
"doneReason": body.get("done_reason"),
|
||
|
|
"promptEvalCount": body.get("prompt_eval_count"),
|
||
|
|
"evalCount": body.get("eval_count"),
|
||
|
|
"totalDurationNs": body.get("total_duration"),
|
||
|
|
"thinkingLength": len(str(message.get("thinking") or "")),
|
||
|
|
}
|
||
|
|
if content:
|
||
|
|
return content, metadata
|
||
|
|
last_metadata = metadata
|
||
|
|
raise VisionModelError(
|
||
|
|
"Vision model returned an empty final response after retry "
|
||
|
|
f"({json.dumps(last_metadata, ensure_ascii=False)})"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
async def _analyze_frames(
|
||
|
|
profile: Profile,
|
||
|
|
strategy_name: str,
|
||
|
|
frames: list[Path],
|
||
|
|
duration: float,
|
||
|
|
extracted_width: int,
|
||
|
|
prompt: str,
|
||
|
|
) -> tuple[dict[str, Any], dict[str, Any]]:
|
||
|
|
strategy = _strategy(profile, strategy_name)
|
||
|
|
selected = _select_frames(frames, max(2, int(strategy["maxFrames"])))
|
||
|
|
maximum_width = max(320, int(strategy["maxWidth"]))
|
||
|
|
images = [_encode_image(path, maximum_width, extracted_width) for path in selected]
|
||
|
|
strategy_prompt = prompt + (
|
||
|
|
f"\n\nThe video duration is {duration:.3f} seconds. "
|
||
|
|
f"The following {len(selected)} images are ordered frames sampled across the video."
|
||
|
|
)
|
||
|
|
payload = {
|
||
|
|
"model": strategy["model"],
|
||
|
|
"stream": False,
|
||
|
|
"think": bool(strategy.get("think", profile.config.get("think", False))),
|
||
|
|
"format": VIDEO_RESULT_SCHEMA,
|
||
|
|
"keep_alive": str(strategy.get("keepAlive", "30m")),
|
||
|
|
"messages": [{"role": "user", "content": strategy_prompt, "images": images}],
|
||
|
|
"options": {
|
||
|
|
"temperature": float(strategy.get("temperature", profile.config.get("temperature", 0.1))),
|
||
|
|
"num_predict": int(strategy["maxOutputTokens"]),
|
||
|
|
"num_ctx": int(strategy["numCtx"]),
|
||
|
|
},
|
||
|
|
}
|
||
|
|
content, metrics = await _request_vision_model(
|
||
|
|
payload, int(strategy.get("emptyResponseRetries", 0))
|
||
|
|
)
|
||
|
|
metadata = {
|
||
|
|
"mode": strategy_name,
|
||
|
|
"model": strategy["model"],
|
||
|
|
"sampledFrameCount": len(selected),
|
||
|
|
"maximumWidth": maximum_width,
|
||
|
|
**metrics,
|
||
|
|
}
|
||
|
|
return _parse_json(content), metadata
|
||
|
|
|
||
|
|
|
||
|
|
async def analyze_video(
|
||
|
|
profile: Profile,
|
||
|
|
media_path: Path,
|
||
|
|
extra_instruction: str = "",
|
||
|
|
analysis_mode: str = "AUTO",
|
||
|
|
) -> tuple[dict[str, Any], float, dict[str, Any]]:
|
||
|
|
requested_mode = str(analysis_mode or "AUTO").strip().upper()
|
||
|
|
if requested_mode not in {"AUTO", "FAST", "ACCURATE"}:
|
||
|
|
raise ValueError("Video analysis mode must be AUTO, FAST, or ACCURATE")
|
||
|
|
|
||
|
|
fast = _strategy(profile, "FAST")
|
||
|
|
accurate = _strategy(profile, "ACCURATE")
|
||
|
|
extraction_strategy = fast if requested_mode == "FAST" else accurate
|
||
|
|
extraction_frames = max(2, int(extraction_strategy["maxFrames"]))
|
||
|
|
extraction_width = max(320, int(extraction_strategy["maxWidth"]))
|
||
|
|
frames, duration, frame_dir = _extract_frames(
|
||
|
|
media_path, extraction_frames, extraction_width
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
prompt_path = profile.directory / str(profile.config.get("promptFile", "prompt.txt"))
|
||
|
|
prompt = prompt_path.read_text(encoding="utf-8")
|
||
|
|
if extra_instruction.strip():
|
||
|
|
prompt += "\n\nAdditional inspection requirement:\n" + extra_instruction.strip()
|
||
|
|
|
||
|
|
if requested_mode in {"FAST", "ACCURATE"}:
|
||
|
|
result, metadata = await _analyze_frames(
|
||
|
|
profile, requested_mode, frames, duration, extraction_width, prompt
|
||
|
|
)
|
||
|
|
metadata.update(
|
||
|
|
{"requestedMode": requested_mode, "fallback": False, "fallbackReason": None}
|
||
|
|
)
|
||
|
|
return result, duration, metadata
|
||
|
|
|
||
|
|
fallback_reason = None
|
||
|
|
try:
|
||
|
|
result, metadata = await _analyze_frames(
|
||
|
|
profile, "FAST", frames, duration, extraction_width, prompt
|
||
|
|
)
|
||
|
|
fallback_reason = _fallback_reason(
|
||
|
|
result, bool(profile.config.get("fallbackWhenPassedUnknown", True))
|
||
|
|
)
|
||
|
|
if fallback_reason is None:
|
||
|
|
metadata.update(
|
||
|
|
{"requestedMode": "AUTO", "fallback": False, "fallbackReason": None}
|
||
|
|
)
|
||
|
|
return result, duration, metadata
|
||
|
|
except (VisionModelError, httpx.HTTPError) as exc:
|
||
|
|
fallback_reason = f"FAST_MODEL_ERROR:{type(exc).__name__}"
|
||
|
|
|
||
|
|
result, metadata = await _analyze_frames(
|
||
|
|
profile, "ACCURATE", frames, duration, extraction_width, prompt
|
||
|
|
)
|
||
|
|
metadata.update(
|
||
|
|
{"requestedMode": "AUTO", "fallback": True, "fallbackReason": fallback_reason}
|
||
|
|
)
|
||
|
|
return result, duration, metadata
|
||
|
|
finally:
|
||
|
|
for frame in frames:
|
||
|
|
frame.unlink(missing_ok=True)
|
||
|
|
frame_dir.rmdir()
|