Initial release — OpenMontage: the first open-source agentic video production system
11 production pipelines, 47 tools, 124 agent skills. Supports cloud APIs (fal.ai, OpenAI, ElevenLabs, Suno, HeyGen, Runway) and free local providers (diffusers, Piper TTS, WAN 2.1, Hunyuan, CogVideo). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
calesthio
@@ -0,0 +1 @@
|
||||
"""Analysis tools for content inspection, transcription, and scene detection."""
|
||||
@@ -0,0 +1,236 @@
|
||||
"""Frame sampler tool wrapping FFmpeg.
|
||||
|
||||
Extracts representative frames from video for AI analysis, thumbnails,
|
||||
or quality inspection. Supports interval-based, count-based, and
|
||||
timestamp-based extraction strategies.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class FrameSampler(BaseTool):
|
||||
name = "frame_sampler"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "analysis"
|
||||
provider = "ffmpeg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["cmd:ffmpeg"]
|
||||
install_instructions = (
|
||||
"Install FFmpeg: https://ffmpeg.org/download.html"
|
||||
)
|
||||
agent_skills = ["ffmpeg"]
|
||||
|
||||
capabilities = [
|
||||
"extract_frames_interval",
|
||||
"extract_frames_count",
|
||||
"extract_frames_timestamps",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path", "strategy"],
|
||||
"properties": {
|
||||
"input_path": {"type": "string"},
|
||||
"strategy": {
|
||||
"type": "string",
|
||||
"enum": ["interval", "count", "timestamps"],
|
||||
},
|
||||
"interval_seconds": {
|
||||
"type": "number",
|
||||
"minimum": 0.1,
|
||||
"description": "Seconds between frames (for interval strategy)",
|
||||
},
|
||||
"count": {
|
||||
"type": "integer",
|
||||
"minimum": 1,
|
||||
"description": "Total frames to extract (for count strategy)",
|
||||
},
|
||||
"timestamps": {
|
||||
"type": "array",
|
||||
"items": {"type": "number"},
|
||||
"description": "Specific timestamps in seconds (for timestamps strategy)",
|
||||
},
|
||||
"output_dir": {"type": "string"},
|
||||
"format": {"type": "string", "enum": ["png", "jpg"], "default": "jpg"},
|
||||
"quality": {"type": "integer", "minimum": 1, "maximum": 31, "default": 2},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=500)
|
||||
idempotency_key_fields = ["input_path", "strategy", "interval_seconds", "count"]
|
||||
side_effects = ["writes frame images to output_dir"]
|
||||
user_visible_verification = ["Inspect extracted frames for representative coverage"]
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
strategy = inputs["strategy"]
|
||||
fmt = inputs.get("format", "jpg")
|
||||
quality = inputs.get("quality", 2)
|
||||
output_dir = Path(inputs.get("output_dir", input_path.parent / "frames"))
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
if strategy == "interval":
|
||||
frames = self._extract_interval(input_path, output_dir, fmt, quality, inputs)
|
||||
elif strategy == "count":
|
||||
frames = self._extract_count(input_path, output_dir, fmt, quality, inputs)
|
||||
elif strategy == "timestamps":
|
||||
frames = self._extract_timestamps(input_path, output_dir, fmt, quality, inputs)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown strategy: {strategy}")
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=str(e))
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"strategy": strategy,
|
||||
"frame_count": len(frames),
|
||||
"frames": frames,
|
||||
"output_dir": str(output_dir),
|
||||
},
|
||||
artifacts=[str(output_dir)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
def _extract_interval(
|
||||
self,
|
||||
input_path: Path,
|
||||
output_dir: Path,
|
||||
fmt: str,
|
||||
quality: int,
|
||||
inputs: dict,
|
||||
) -> list[dict]:
|
||||
interval = inputs.get("interval_seconds", 5.0)
|
||||
output_pattern = str(output_dir / f"frame_%04d.{fmt}")
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
"-vf", f"fps=1/{interval}",
|
||||
]
|
||||
if fmt == "jpg":
|
||||
cmd.extend(["-qscale:v", str(quality)])
|
||||
cmd.append(output_pattern)
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
return self._collect_frames(output_dir, fmt, interval)
|
||||
|
||||
def _extract_count(
|
||||
self,
|
||||
input_path: Path,
|
||||
output_dir: Path,
|
||||
fmt: str,
|
||||
quality: int,
|
||||
inputs: dict,
|
||||
) -> list[dict]:
|
||||
count = inputs.get("count", 10)
|
||||
duration = self._get_duration(input_path)
|
||||
if duration <= 0:
|
||||
return []
|
||||
|
||||
interval = duration / count
|
||||
output_pattern = str(output_dir / f"frame_%04d.{fmt}")
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
"-vf", f"fps=1/{interval}",
|
||||
"-frames:v", str(count),
|
||||
]
|
||||
if fmt == "jpg":
|
||||
cmd.extend(["-qscale:v", str(quality)])
|
||||
cmd.append(output_pattern)
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
return self._collect_frames(output_dir, fmt, interval)
|
||||
|
||||
def _extract_timestamps(
|
||||
self,
|
||||
input_path: Path,
|
||||
output_dir: Path,
|
||||
fmt: str,
|
||||
quality: int,
|
||||
inputs: dict,
|
||||
) -> list[dict]:
|
||||
timestamps = inputs.get("timestamps", [])
|
||||
frames = []
|
||||
|
||||
for i, ts in enumerate(timestamps):
|
||||
output_file = output_dir / f"frame_{i:04d}.{fmt}"
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-ss", str(ts),
|
||||
"-i", str(input_path),
|
||||
"-frames:v", "1",
|
||||
]
|
||||
if fmt == "jpg":
|
||||
cmd.extend(["-qscale:v", str(quality)])
|
||||
cmd.append(str(output_file))
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
if output_file.exists():
|
||||
frames.append({
|
||||
"path": str(output_file),
|
||||
"timestamp_seconds": ts,
|
||||
"index": i,
|
||||
})
|
||||
|
||||
return frames
|
||||
|
||||
def _get_duration(self, input_path: Path) -> float:
|
||||
"""Get video duration in seconds via ffprobe."""
|
||||
cmd = [
|
||||
"ffprobe",
|
||||
"-v", "quiet",
|
||||
"-show_entries", "format=duration",
|
||||
"-of", "json",
|
||||
str(input_path),
|
||||
]
|
||||
result = self.run_command(cmd)
|
||||
data = json.loads(result.stdout)
|
||||
return float(data.get("format", {}).get("duration", 0))
|
||||
|
||||
def _collect_frames(
|
||||
self, output_dir: Path, fmt: str, interval: float
|
||||
) -> list[dict]:
|
||||
"""Collect extracted frame files and build metadata."""
|
||||
frames = []
|
||||
pattern = f"frame_*.{fmt}"
|
||||
for i, path in enumerate(sorted(output_dir.glob(pattern))):
|
||||
frames.append({
|
||||
"path": str(path),
|
||||
"timestamp_seconds": round(i * interval, 3),
|
||||
"index": i,
|
||||
})
|
||||
return frames
|
||||
@@ -0,0 +1,262 @@
|
||||
"""Scene detection tool wrapping PySceneDetect.
|
||||
|
||||
Detects scene boundaries and shot changes in video. Falls back to
|
||||
FFmpeg-based detection if PySceneDetect is not installed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class SceneDetect(BaseTool):
|
||||
name = "scene_detect"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "analysis"
|
||||
provider = "ffmpeg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["cmd:ffmpeg"]
|
||||
install_instructions = (
|
||||
"FFmpeg is required. For better detection install PySceneDetect:\n"
|
||||
"pip install scenedetect[opencv]"
|
||||
)
|
||||
agent_skills = ["ffmpeg"]
|
||||
|
||||
capabilities = [
|
||||
"detect_scenes",
|
||||
"detect_content_changes",
|
||||
"detect_threshold",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path"],
|
||||
"properties": {
|
||||
"input_path": {"type": "string"},
|
||||
"method": {
|
||||
"type": "string",
|
||||
"enum": ["content", "threshold", "adaptive"],
|
||||
"default": "content",
|
||||
},
|
||||
"threshold": {
|
||||
"type": "number",
|
||||
"description": "Detection threshold (method-dependent)",
|
||||
},
|
||||
"min_scene_length_seconds": {
|
||||
"type": "number",
|
||||
"minimum": 0.1,
|
||||
"default": 1.0,
|
||||
},
|
||||
"output_path": {"type": "string", "description": "Path for scene list JSON"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=2, ram_mb=1024, vram_mb=0, disk_mb=100)
|
||||
idempotency_key_fields = ["input_path", "method", "threshold"]
|
||||
side_effects = ["writes scene list JSON to output_path"]
|
||||
user_visible_verification = [
|
||||
"Spot-check detected scene boundaries against the video",
|
||||
]
|
||||
|
||||
def _has_pyscenedetect(self) -> bool:
|
||||
try:
|
||||
import scenedetect # noqa: F401
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
start = time.time()
|
||||
|
||||
if self._has_pyscenedetect():
|
||||
scenes = self._detect_pyscenedetect(inputs)
|
||||
else:
|
||||
scenes = self._detect_ffmpeg(inputs)
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
# Write scene list
|
||||
output_path = Path(
|
||||
inputs.get("output_path", str(input_path.with_suffix(".scenes.json")))
|
||||
)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text(json.dumps({"scenes": scenes}, indent=2), encoding="utf-8")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"scene_count": len(scenes),
|
||||
"scenes": scenes,
|
||||
"method": "pyscenedetect" if self._has_pyscenedetect() else "ffmpeg",
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
def _detect_pyscenedetect(self, inputs: dict[str, Any]) -> list[dict]:
|
||||
"""Use PySceneDetect for scene detection."""
|
||||
from scenedetect import open_video, SceneManager
|
||||
from scenedetect.detectors import ContentDetector, ThresholdDetector, AdaptiveDetector
|
||||
|
||||
input_path = str(inputs["input_path"])
|
||||
method = inputs.get("method", "content")
|
||||
threshold = inputs.get("threshold")
|
||||
min_scene_len = inputs.get("min_scene_length_seconds", 1.0)
|
||||
|
||||
video = open_video(input_path)
|
||||
scene_manager = SceneManager()
|
||||
|
||||
if method == "content":
|
||||
detector = ContentDetector(
|
||||
threshold=threshold or 27.0,
|
||||
min_scene_len=int(min_scene_len * video.frame_rate),
|
||||
)
|
||||
elif method == "threshold":
|
||||
detector = ThresholdDetector(
|
||||
threshold=threshold or 12.0,
|
||||
min_scene_len=int(min_scene_len * video.frame_rate),
|
||||
)
|
||||
elif method == "adaptive":
|
||||
detector = AdaptiveDetector(
|
||||
adaptive_threshold=threshold or 3.0,
|
||||
min_scene_len=int(min_scene_len * video.frame_rate),
|
||||
)
|
||||
else:
|
||||
detector = ContentDetector()
|
||||
|
||||
scene_manager.add_detector(detector)
|
||||
scene_manager.detect_scenes(video)
|
||||
scene_list = scene_manager.get_scene_list()
|
||||
|
||||
scenes = []
|
||||
for i, (scene_start, scene_end) in enumerate(scene_list):
|
||||
scenes.append({
|
||||
"index": i,
|
||||
"start_seconds": round(scene_start.get_seconds(), 3),
|
||||
"end_seconds": round(scene_end.get_seconds(), 3),
|
||||
"duration_seconds": round(
|
||||
scene_end.get_seconds() - scene_start.get_seconds(), 3
|
||||
),
|
||||
})
|
||||
|
||||
return scenes
|
||||
|
||||
def _detect_ffmpeg(self, inputs: dict[str, Any]) -> list[dict]:
|
||||
"""Fallback: use FFmpeg scene change filter."""
|
||||
input_path = str(inputs["input_path"])
|
||||
threshold = inputs.get("threshold", 0.3)
|
||||
min_scene_len = inputs.get("min_scene_length_seconds", 1.0)
|
||||
|
||||
cmd = [
|
||||
"ffprobe",
|
||||
"-v", "quiet",
|
||||
"-show_entries", "frame=pts_time",
|
||||
"-of", "json",
|
||||
"-f", "lavfi",
|
||||
f"movie='{input_path.replace(chr(92), '/').replace(':', chr(92)+':')}',select='gt(scene,{threshold})'",
|
||||
]
|
||||
|
||||
try:
|
||||
result = self.run_command(cmd, timeout=120)
|
||||
data = json.loads(result.stdout)
|
||||
except Exception:
|
||||
# If ffprobe lavfi approach fails, try a simpler method
|
||||
return self._detect_ffmpeg_simple(input_path, threshold, min_scene_len)
|
||||
|
||||
change_points = [0.0]
|
||||
for frame in data.get("frames", []):
|
||||
ts = float(frame.get("pts_time", 0))
|
||||
if ts - change_points[-1] >= min_scene_len:
|
||||
change_points.append(ts)
|
||||
|
||||
# Get total duration
|
||||
dur_cmd = [
|
||||
"ffprobe", "-v", "quiet",
|
||||
"-show_entries", "format=duration",
|
||||
"-of", "json", input_path,
|
||||
]
|
||||
dur_result = self.run_command(dur_cmd)
|
||||
total_dur = float(json.loads(dur_result.stdout)["format"]["duration"])
|
||||
change_points.append(total_dur)
|
||||
|
||||
scenes = []
|
||||
for i in range(len(change_points) - 1):
|
||||
start = change_points[i]
|
||||
end = change_points[i + 1]
|
||||
scenes.append({
|
||||
"index": i,
|
||||
"start_seconds": round(start, 3),
|
||||
"end_seconds": round(end, 3),
|
||||
"duration_seconds": round(end - start, 3),
|
||||
})
|
||||
|
||||
return scenes
|
||||
|
||||
def _detect_ffmpeg_simple(
|
||||
self, input_path: str, threshold: float, min_scene_len: float
|
||||
) -> list[dict]:
|
||||
"""Simplest fallback: split into uniform segments."""
|
||||
dur_cmd = [
|
||||
"ffprobe", "-v", "quiet",
|
||||
"-show_entries", "format=duration",
|
||||
"-of", "json", input_path,
|
||||
]
|
||||
dur_result = self.run_command(dur_cmd)
|
||||
total_dur = float(json.loads(dur_result.stdout)["format"]["duration"])
|
||||
|
||||
# Use select filter to find scene changes via stdout
|
||||
cmd = [
|
||||
"ffmpeg",
|
||||
"-i", input_path,
|
||||
"-vf", f"select='gt(scene,{threshold})',showinfo",
|
||||
"-f", "null", "-",
|
||||
]
|
||||
try:
|
||||
result = self.run_command(cmd, timeout=120)
|
||||
output = result.stderr
|
||||
except Exception:
|
||||
output = ""
|
||||
|
||||
import re
|
||||
change_points = [0.0]
|
||||
for match in re.finditer(r"pts_time:(\d+\.?\d*)", output):
|
||||
ts = float(match.group(1))
|
||||
if ts - change_points[-1] >= min_scene_len:
|
||||
change_points.append(ts)
|
||||
change_points.append(total_dur)
|
||||
|
||||
scenes = []
|
||||
for i in range(len(change_points) - 1):
|
||||
start = change_points[i]
|
||||
end = change_points[i + 1]
|
||||
scenes.append({
|
||||
"index": i,
|
||||
"start_seconds": round(start, 3),
|
||||
"end_seconds": round(end, 3),
|
||||
"duration_seconds": round(end - start, 3),
|
||||
})
|
||||
|
||||
return scenes
|
||||
@@ -0,0 +1,251 @@
|
||||
"""Transcription tool wrapping faster-whisper / WhisperX.
|
||||
|
||||
Provides speech-to-text with word-level timestamps and optional speaker
|
||||
diarization. Falls back gracefully when GPU or diarization dependencies
|
||||
are not available.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ResumeSupport,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class Transcriber(BaseTool):
|
||||
name = "transcriber"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "analysis"
|
||||
provider = "whisperx"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["python:faster_whisper"]
|
||||
install_instructions = (
|
||||
"pip install faster-whisper # CPU mode\n"
|
||||
"pip install faster-whisper[gpu] # GPU mode (requires CUDA)\n"
|
||||
"pip install whisperx # For diarization support"
|
||||
)
|
||||
agent_skills = ["speech-to-text"]
|
||||
|
||||
capabilities = [
|
||||
"transcribe",
|
||||
"word_timestamps",
|
||||
"diarization",
|
||||
"language_detection",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path"],
|
||||
"properties": {
|
||||
"input_path": {"type": "string", "description": "Path to audio or video file"},
|
||||
"model_size": {
|
||||
"type": "string",
|
||||
"enum": ["tiny", "base", "small", "medium", "large-v2", "large-v3"],
|
||||
"default": "base",
|
||||
},
|
||||
"language": {"type": "string", "description": "ISO 639-1 language code, or null for auto-detect"},
|
||||
"diarize": {"type": "boolean", "default": False},
|
||||
"output_dir": {"type": "string", "description": "Directory for output files"},
|
||||
},
|
||||
}
|
||||
|
||||
output_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"segments": {"type": "array"},
|
||||
"word_timestamps": {"type": "array"},
|
||||
"language": {"type": "string"},
|
||||
"duration_seconds": {"type": "number"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2,
|
||||
ram_mb=2048,
|
||||
vram_mb=0, # CPU by default; GPU optional
|
||||
disk_mb=500,
|
||||
network_required=False,
|
||||
)
|
||||
|
||||
retry_policy = RetryPolicy(max_retries=1, retryable_errors=["MemoryError"])
|
||||
resume_support = ResumeSupport.FROM_START
|
||||
idempotency_key_fields = ["input_path", "model_size", "language"]
|
||||
side_effects = ["writes transcript JSON to output_dir"]
|
||||
fallback = None
|
||||
user_visible_verification = [
|
||||
"Check transcript text against source audio",
|
||||
"Verify word timestamps align with speech",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
try:
|
||||
import faster_whisper # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def _has_diarization(self) -> bool:
|
||||
try:
|
||||
import whisperx # noqa: F401
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
"""Rough estimate: ~0.5x real-time on CPU for 'base' model."""
|
||||
return 60.0 # conservative default
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
model_size = inputs.get("model_size", "base")
|
||||
language = inputs.get("language")
|
||||
diarize = inputs.get("diarize", False)
|
||||
output_dir = Path(inputs.get("output_dir", input_path.parent))
|
||||
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input file not found: {input_path}")
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
from faster_whisper import WhisperModel
|
||||
except ImportError:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="faster-whisper is not installed. Run: pip install faster-whisper",
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
# Load model (CPU by default, CUDA if available)
|
||||
try:
|
||||
import torch
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
compute_type = "float16" if device == "cuda" else "int8"
|
||||
except ImportError:
|
||||
device = "cpu"
|
||||
compute_type = "int8"
|
||||
|
||||
model = WhisperModel(model_size, device=device, compute_type=compute_type)
|
||||
|
||||
# Transcribe
|
||||
segments_iter, info = model.transcribe(
|
||||
str(input_path),
|
||||
language=language,
|
||||
word_timestamps=True,
|
||||
vad_filter=True,
|
||||
)
|
||||
|
||||
segments = []
|
||||
word_timestamps = []
|
||||
|
||||
for seg in segments_iter:
|
||||
seg_data = {
|
||||
"id": seg.id,
|
||||
"start": round(seg.start, 3),
|
||||
"end": round(seg.end, 3),
|
||||
"text": seg.text.strip(),
|
||||
}
|
||||
|
||||
if seg.words:
|
||||
words = []
|
||||
for w in seg.words:
|
||||
word_entry = {
|
||||
"word": w.word,
|
||||
"start": round(w.start, 3),
|
||||
"end": round(w.end, 3),
|
||||
"probability": round(w.probability, 3),
|
||||
}
|
||||
words.append(word_entry)
|
||||
word_timestamps.append(word_entry)
|
||||
seg_data["words"] = words
|
||||
|
||||
segments.append(seg_data)
|
||||
|
||||
detected_language = language or info.language
|
||||
duration = info.duration
|
||||
|
||||
# Optional diarization pass
|
||||
if diarize and self._has_diarization():
|
||||
segments = self._apply_diarization(
|
||||
str(input_path), segments, detected_language
|
||||
)
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
result_data = {
|
||||
"segments": segments,
|
||||
"word_timestamps": word_timestamps,
|
||||
"language": detected_language,
|
||||
"duration_seconds": round(duration, 3),
|
||||
"model_size": model_size,
|
||||
"device": device,
|
||||
}
|
||||
|
||||
# Write transcript JSON
|
||||
output_path = output_dir / f"{input_path.stem}_transcript.json"
|
||||
output_path.write_text(json.dumps(result_data, indent=2), encoding="utf-8")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data=result_data,
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
def _apply_diarization(
|
||||
self,
|
||||
audio_path: str,
|
||||
segments: list[dict],
|
||||
language: str,
|
||||
) -> list[dict]:
|
||||
"""Apply WhisperX diarization to assign speaker labels."""
|
||||
try:
|
||||
import whisperx
|
||||
|
||||
# Load audio for alignment
|
||||
audio = whisperx.load_audio(audio_path)
|
||||
|
||||
# Align segments with word timestamps
|
||||
align_model, align_metadata = whisperx.load_align_model(
|
||||
language_code=language, device="cpu"
|
||||
)
|
||||
aligned = whisperx.align(
|
||||
segments, align_model, align_metadata, audio, device="cpu"
|
||||
)
|
||||
|
||||
# Diarize
|
||||
import os
|
||||
hf_token = os.environ.get("HF_TOKEN")
|
||||
if not hf_token:
|
||||
# Can't diarize without HuggingFace token for pyannote
|
||||
return segments
|
||||
|
||||
diarize_model = whisperx.DiarizationPipeline(
|
||||
use_auth_token=hf_token, device="cpu"
|
||||
)
|
||||
diarize_segments = diarize_model(audio)
|
||||
result = whisperx.assign_word_speakers(diarize_segments, aligned)
|
||||
|
||||
return result.get("segments", segments)
|
||||
except Exception:
|
||||
# Diarization is best-effort; return original segments on failure
|
||||
return segments
|
||||
@@ -0,0 +1,597 @@
|
||||
"""Video/image understanding tool using vision-language models.
|
||||
|
||||
Analyzes images or video frames using CLIP, BLIP-2, or LLaVA. Supports
|
||||
frame description, visual question answering, quality assessment, and
|
||||
scene classification. Primary use case: visual QA for automated video review.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
VIDEO_EXTENSIONS = {".mp4", ".mov", ".avi", ".webm"}
|
||||
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".webp"}
|
||||
|
||||
SCENE_CATEGORIES = [
|
||||
"indoor", "outdoor", "landscape", "cityscape", "portrait",
|
||||
"action", "close-up", "aerial", "underwater", "night",
|
||||
"studio", "nature", "urban", "abstract", "text-overlay",
|
||||
]
|
||||
|
||||
|
||||
class VideoUnderstand(BaseTool):
|
||||
name = "video_understand"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.ANALYZE
|
||||
capability = "analysis"
|
||||
provider = "transformers"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
dependencies = ["python:transformers", "python:torch"]
|
||||
install_instructions = (
|
||||
"pip install transformers torch # For CLIP/BLIP-2 visual understanding"
|
||||
)
|
||||
agent_skills = ["video-understand"]
|
||||
|
||||
capabilities = [
|
||||
"image_description",
|
||||
"visual_qa",
|
||||
"quality_assessment",
|
||||
"scene_classification",
|
||||
"object_detection",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path"],
|
||||
"properties": {
|
||||
"input_path": {
|
||||
"type": "string",
|
||||
"description": "Path to image or video file",
|
||||
},
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "Question to answer about the visual content (for VQA mode)",
|
||||
},
|
||||
"mode": {
|
||||
"type": "string",
|
||||
"enum": ["describe", "qa", "quality", "classify"],
|
||||
"default": "describe",
|
||||
"description": (
|
||||
"describe: generate caption; qa: answer query about content; "
|
||||
"quality: assess technical quality; classify: classify scene type"
|
||||
),
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["clip", "blip2", "llava"],
|
||||
"default": "clip",
|
||||
"description": "Which vision-language model to use",
|
||||
},
|
||||
"frame_indices": {
|
||||
"type": "array",
|
||||
"items": {"type": "integer"},
|
||||
"description": (
|
||||
"For video input, which frames to analyze. "
|
||||
"If not provided, samples key frames at even intervals."
|
||||
),
|
||||
},
|
||||
"max_frames": {
|
||||
"type": "integer",
|
||||
"default": 5,
|
||||
"description": "Maximum number of frames to analyze from video",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
output_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"frames": {
|
||||
"type": "array",
|
||||
"description": "Per-frame analysis results",
|
||||
},
|
||||
"summary": {"type": "string"},
|
||||
"mode": {"type": "string"},
|
||||
"model": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2,
|
||||
ram_mb=4096,
|
||||
vram_mb=2048,
|
||||
disk_mb=1000,
|
||||
network_required=False,
|
||||
)
|
||||
|
||||
idempotency_key_fields = ["input_path", "mode", "model", "query"]
|
||||
side_effects = []
|
||||
fallback = None
|
||||
user_visible_verification = [
|
||||
"Compare generated descriptions against actual visual content",
|
||||
"Verify quality scores match perceived image quality",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
"""Check if transformers and torch are both importable."""
|
||||
try:
|
||||
import transformers # noqa: F401
|
||||
import torch # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
"""Estimate runtime in seconds based on mode and frame count."""
|
||||
max_frames = inputs.get("max_frames", 5)
|
||||
mode = inputs.get("mode", "describe")
|
||||
if mode == "quality":
|
||||
return max_frames * 0.5 # quality metrics are fast
|
||||
return max_frames * 5.0 # VLM inference per frame
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
mode = inputs.get("mode", "describe")
|
||||
model_name = inputs.get("model", "clip")
|
||||
query = inputs.get("query")
|
||||
frame_indices = inputs.get("frame_indices")
|
||||
max_frames = inputs.get("max_frames", 5)
|
||||
|
||||
if not input_path.exists():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Input file not found: {input_path}",
|
||||
)
|
||||
|
||||
suffix = input_path.suffix.lower()
|
||||
is_video = suffix in VIDEO_EXTENSIONS
|
||||
is_image = suffix in IMAGE_EXTENSIONS
|
||||
|
||||
if not is_video and not is_image:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=(
|
||||
f"Unsupported file type: {suffix}. "
|
||||
f"Supported: {sorted(VIDEO_EXTENSIONS | IMAGE_EXTENSIONS)}"
|
||||
),
|
||||
)
|
||||
|
||||
if mode == "qa" and not query:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Query is required for 'qa' mode.",
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
# --- Load frames ---
|
||||
try:
|
||||
frames = self._load_frames(
|
||||
input_path, is_video, frame_indices, max_frames
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Failed to load frames: {e}")
|
||||
|
||||
if not frames:
|
||||
return ToolResult(success=False, error="No frames could be extracted.")
|
||||
|
||||
# --- Analyze each frame ---
|
||||
try:
|
||||
if mode == "quality":
|
||||
frame_results = self._analyze_quality(frames)
|
||||
elif mode == "describe":
|
||||
frame_results = self._analyze_describe(frames, model_name)
|
||||
elif mode == "qa":
|
||||
frame_results = self._analyze_qa(frames, model_name, query)
|
||||
elif mode == "classify":
|
||||
frame_results = self._analyze_classify(frames, model_name)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown mode: {mode}")
|
||||
except ImportError as e:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Missing dependency for {model_name}: {e}",
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Analysis failed: {e}")
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
# --- Build summary ---
|
||||
summary = self._build_summary(frame_results, mode)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"frames": frame_results,
|
||||
"summary": summary,
|
||||
"mode": mode,
|
||||
"model": model_name if mode != "quality" else "metrics",
|
||||
"frame_count": len(frame_results),
|
||||
},
|
||||
duration_seconds=round(elapsed, 2),
|
||||
model=model_name if mode != "quality" else None,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Frame extraction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _load_frames(
|
||||
self,
|
||||
input_path: Path,
|
||||
is_video: bool,
|
||||
frame_indices: list[int] | None,
|
||||
max_frames: int,
|
||||
) -> list:
|
||||
"""Load PIL Image objects from an image or video file."""
|
||||
from PIL import Image
|
||||
|
||||
if not is_video:
|
||||
return [Image.open(input_path).convert("RGB")]
|
||||
|
||||
return self._extract_video_frames(input_path, frame_indices, max_frames)
|
||||
|
||||
def _extract_video_frames(
|
||||
self,
|
||||
video_path: Path,
|
||||
frame_indices: list[int] | None,
|
||||
max_frames: int,
|
||||
) -> list:
|
||||
"""Extract frames from a video file using ffmpeg."""
|
||||
from PIL import Image
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp = Path(tmp_dir)
|
||||
|
||||
if frame_indices:
|
||||
# Extract specific frames using select filter
|
||||
frames_to_extract = frame_indices[:max_frames]
|
||||
select_expr = "+".join(
|
||||
f"eq(n\\,{idx})" for idx in frames_to_extract
|
||||
)
|
||||
cmd = [
|
||||
"ffmpeg", "-i", str(video_path),
|
||||
"-vf", f"select='{select_expr}'",
|
||||
"-vsync", "vfr",
|
||||
str(tmp / "frame_%04d.png"),
|
||||
"-y", "-loglevel", "error",
|
||||
]
|
||||
else:
|
||||
# Get total frame count first
|
||||
probe_cmd = [
|
||||
"ffmpeg", "-i", str(video_path),
|
||||
"-map", "0:v:0", "-c", "copy", "-f", "null", "-",
|
||||
]
|
||||
# Sample at even intervals using fps filter
|
||||
# Use a select filter that picks frames at even intervals
|
||||
cmd = [
|
||||
"ffmpeg", "-i", str(video_path),
|
||||
"-frames:v", str(max_frames),
|
||||
"-vf", f"thumbnail={max_frames}",
|
||||
str(tmp / "frame_%04d.png"),
|
||||
"-y", "-loglevel", "error",
|
||||
]
|
||||
|
||||
subprocess.run(cmd, capture_output=True, text=True, timeout=60)
|
||||
|
||||
# Load extracted frames
|
||||
frame_files = sorted(tmp.glob("frame_*.png"))
|
||||
images = []
|
||||
for f in frame_files[:max_frames]:
|
||||
images.append(Image.open(f).convert("RGB"))
|
||||
|
||||
return images
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Quality assessment (uses PIL/numpy, no VLM needed)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _analyze_quality(self, frames: list) -> list[dict[str, Any]]:
|
||||
"""Assess technical quality using simple image metrics."""
|
||||
import numpy as np
|
||||
|
||||
results = []
|
||||
for i, img in enumerate(frames):
|
||||
arr = np.array(img, dtype=np.float64)
|
||||
gray = np.mean(arr, axis=2)
|
||||
|
||||
# Blur detection: Laplacian variance (low = blurry)
|
||||
# Manual Laplacian approximation using numpy
|
||||
laplacian_kernel = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]])
|
||||
from scipy.signal import convolve2d
|
||||
laplacian = convolve2d(gray, laplacian_kernel, mode="valid")
|
||||
blur_score = float(np.var(laplacian))
|
||||
|
||||
# Brightness: mean pixel value (0-255 scale)
|
||||
brightness = float(np.mean(arr))
|
||||
|
||||
# Contrast: standard deviation of pixel values
|
||||
contrast = float(np.std(arr))
|
||||
|
||||
# Classify quality
|
||||
quality_issues = []
|
||||
if blur_score < 100:
|
||||
quality_issues.append("blurry")
|
||||
if brightness < 40:
|
||||
quality_issues.append("underexposed")
|
||||
elif brightness > 220:
|
||||
quality_issues.append("overexposed")
|
||||
if contrast < 30:
|
||||
quality_issues.append("low_contrast")
|
||||
|
||||
quality_label = "good" if not quality_issues else "issues_detected"
|
||||
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"blur_score": round(blur_score, 2),
|
||||
"brightness": round(brightness, 2),
|
||||
"contrast": round(contrast, 2),
|
||||
"quality": quality_label,
|
||||
"issues": quality_issues,
|
||||
"resolution": f"{img.width}x{img.height}",
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# VLM-based analysis modes
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _load_model(self, model_name: str):
|
||||
"""Load the requested vision-language model and processor."""
|
||||
import torch
|
||||
from transformers import (
|
||||
CLIPProcessor,
|
||||
CLIPModel,
|
||||
BlipProcessor,
|
||||
BlipForConditionalGeneration,
|
||||
Blip2Processor,
|
||||
Blip2ForConditionalGeneration,
|
||||
AutoProcessor,
|
||||
AutoModelForCausalLM,
|
||||
)
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
if model_name == "clip":
|
||||
model_id = "openai/clip-vit-base-patch32"
|
||||
processor = CLIPProcessor.from_pretrained(model_id)
|
||||
model = CLIPModel.from_pretrained(model_id).to(device)
|
||||
return model, processor, device
|
||||
|
||||
if model_name == "blip2":
|
||||
model_id = "Salesforce/blip2-opt-2.7b"
|
||||
processor = Blip2Processor.from_pretrained(model_id)
|
||||
model = Blip2ForConditionalGeneration.from_pretrained(
|
||||
model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32
|
||||
).to(device)
|
||||
return model, processor, device
|
||||
|
||||
if model_name == "llava":
|
||||
model_id = "llava-hf/llava-1.5-7b-hf"
|
||||
processor = AutoProcessor.from_pretrained(model_id)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32
|
||||
).to(device)
|
||||
return model, processor, device
|
||||
|
||||
raise ValueError(f"Unknown model: {model_name}")
|
||||
|
||||
def _analyze_describe(
|
||||
self, frames: list, model_name: str
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Generate captions for each frame."""
|
||||
import torch
|
||||
|
||||
model, processor, device = self._load_model(model_name)
|
||||
results = []
|
||||
|
||||
for i, img in enumerate(frames):
|
||||
if model_name == "clip":
|
||||
# CLIP is not a captioning model; use zero-shot classification
|
||||
# with generic scene descriptions as a caption proxy
|
||||
candidate_texts = [
|
||||
"a photo of a person", "a photo of a landscape",
|
||||
"a photo of an object", "a photo of text",
|
||||
"a photo of an animal", "a photo of a building",
|
||||
"a photo of food", "a photo of a vehicle",
|
||||
"an abstract image", "a dark scene", "a bright scene",
|
||||
]
|
||||
clip_inputs = processor(
|
||||
text=candidate_texts, images=img, return_tensors="pt", padding=True
|
||||
).to(device)
|
||||
with torch.no_grad():
|
||||
outputs = model(**clip_inputs)
|
||||
probs = outputs.logits_per_image.softmax(dim=1)[0]
|
||||
top_idx = probs.argmax().item()
|
||||
caption = candidate_texts[top_idx]
|
||||
confidence = round(probs[top_idx].item(), 3)
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"description": caption,
|
||||
"confidence": confidence,
|
||||
})
|
||||
|
||||
elif model_name in ("blip2", "llava"):
|
||||
inputs = processor(images=img, return_tensors="pt").to(device)
|
||||
with torch.no_grad():
|
||||
generated_ids = model.generate(**inputs, max_new_tokens=50)
|
||||
caption = processor.batch_decode(
|
||||
generated_ids, skip_special_tokens=True
|
||||
)[0].strip()
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"description": caption,
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
def _analyze_qa(
|
||||
self, frames: list, model_name: str, query: str
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Answer a question about each frame."""
|
||||
import torch
|
||||
|
||||
model, processor, device = self._load_model(model_name)
|
||||
results = []
|
||||
|
||||
for i, img in enumerate(frames):
|
||||
if model_name == "clip":
|
||||
# Use query as one candidate and its negation as another
|
||||
candidates = [query, f"not {query}"]
|
||||
clip_inputs = processor(
|
||||
text=candidates, images=img, return_tensors="pt", padding=True
|
||||
).to(device)
|
||||
with torch.no_grad():
|
||||
outputs = model(**clip_inputs)
|
||||
probs = outputs.logits_per_image.softmax(dim=1)[0]
|
||||
yes_prob = probs[0].item()
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"query": query,
|
||||
"answer": "yes" if yes_prob > 0.5 else "no",
|
||||
"confidence": round(max(yes_prob, 1 - yes_prob), 3),
|
||||
})
|
||||
|
||||
elif model_name in ("blip2", "llava"):
|
||||
prompt = f"Question: {query} Answer:"
|
||||
inputs = processor(
|
||||
images=img, text=prompt, return_tensors="pt"
|
||||
).to(device)
|
||||
with torch.no_grad():
|
||||
generated_ids = model.generate(**inputs, max_new_tokens=50)
|
||||
answer = processor.batch_decode(
|
||||
generated_ids, skip_special_tokens=True
|
||||
)[0].strip()
|
||||
# Remove the prompt echo if present
|
||||
if answer.startswith(prompt):
|
||||
answer = answer[len(prompt):].strip()
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"query": query,
|
||||
"answer": answer,
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
def _analyze_classify(
|
||||
self, frames: list, model_name: str
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Classify each frame into scene categories."""
|
||||
import torch
|
||||
|
||||
model, processor, device = self._load_model(model_name)
|
||||
results = []
|
||||
|
||||
for i, img in enumerate(frames):
|
||||
if model_name == "clip":
|
||||
candidate_texts = [f"a {cat} scene" for cat in SCENE_CATEGORIES]
|
||||
clip_inputs = processor(
|
||||
text=candidate_texts, images=img, return_tensors="pt", padding=True
|
||||
).to(device)
|
||||
with torch.no_grad():
|
||||
outputs = model(**clip_inputs)
|
||||
probs = outputs.logits_per_image.softmax(dim=1)[0]
|
||||
|
||||
scored = sorted(
|
||||
zip(SCENE_CATEGORIES, probs.tolist()),
|
||||
key=lambda x: x[1],
|
||||
reverse=True,
|
||||
)
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"top_category": scored[0][0],
|
||||
"confidence": round(scored[0][1], 3),
|
||||
"categories": [
|
||||
{"label": label, "score": round(score, 3)}
|
||||
for label, score in scored[:5]
|
||||
],
|
||||
})
|
||||
|
||||
elif model_name in ("blip2", "llava"):
|
||||
prompt = (
|
||||
"Classify this image into one of these categories: "
|
||||
+ ", ".join(SCENE_CATEGORIES)
|
||||
+ ". Category:"
|
||||
)
|
||||
inputs = processor(
|
||||
images=img, text=prompt, return_tensors="pt"
|
||||
).to(device)
|
||||
with torch.no_grad():
|
||||
generated_ids = model.generate(**inputs, max_new_tokens=20)
|
||||
category = processor.batch_decode(
|
||||
generated_ids, skip_special_tokens=True
|
||||
)[0].strip()
|
||||
if category.startswith(prompt):
|
||||
category = category[len(prompt):].strip()
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"top_category": category,
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Summary
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_summary(
|
||||
self, frame_results: list[dict[str, Any]], mode: str
|
||||
) -> str:
|
||||
"""Build a human-readable summary from per-frame results."""
|
||||
n = len(frame_results)
|
||||
|
||||
if mode == "describe":
|
||||
descriptions = [r.get("description", "") for r in frame_results]
|
||||
if n == 1:
|
||||
return descriptions[0]
|
||||
return f"Analyzed {n} frames. Descriptions: " + "; ".join(descriptions)
|
||||
|
||||
if mode == "qa":
|
||||
answers = [r.get("answer", "") for r in frame_results]
|
||||
if n == 1:
|
||||
return answers[0]
|
||||
return f"Analyzed {n} frames. Answers: " + "; ".join(answers)
|
||||
|
||||
if mode == "quality":
|
||||
issues_all = []
|
||||
for r in frame_results:
|
||||
issues_all.extend(r.get("issues", []))
|
||||
if not issues_all:
|
||||
return f"All {n} frame(s) passed quality checks."
|
||||
unique_issues = sorted(set(issues_all))
|
||||
return (
|
||||
f"Analyzed {n} frame(s). Issues found: {', '.join(unique_issues)}."
|
||||
)
|
||||
|
||||
if mode == "classify":
|
||||
categories = [r.get("top_category", "unknown") for r in frame_results]
|
||||
if n == 1:
|
||||
return f"Scene classified as: {categories[0]}"
|
||||
return (
|
||||
f"Analyzed {n} frames. Scene categories: "
|
||||
+ ", ".join(categories)
|
||||
)
|
||||
|
||||
return f"Analyzed {n} frame(s)."
|
||||
Reference in New Issue
Block a user