feat: video-edit-planner skill — transcribe, extract frames, plan edits
A conversational video-editing planning assistant. Provides tools for: - Audio transcription via bundled funasr-script (Fun-ASR-Nano / SenseVoice) - On-demand clip extraction and frame sampling (ffmpeg, hardware-accelerated) - SQLite index to track all artifacts and avoid duplicate processing - Vision analysis guidance with binary-search-style frame sampling - Iterative Markdown-table edit plan output Agent-agnostic: no platform-specific tool names, works with any agent runtime. Dependencies checked at guidance level with OS package manager install hints.
This commit is contained in:
@@ -0,0 +1,221 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Extract clips and frames from a video for visual analysis.
|
||||
|
||||
This script wraps ffmpeg to:
|
||||
1. Extract a time range as a clip (keyframe-accurate, -c copy).
|
||||
2. Extract frames from that clip (scene detection + uniform sampling).
|
||||
|
||||
It is designed to be called by the agent, not directly by the user.
|
||||
The agent decides parameters based on clip length, content type, and user needs.
|
||||
|
||||
Usage:
|
||||
# Extract a clip
|
||||
python extract_frames.py clip --video VIDEO --start 60 --end 90 --output clip.mkv
|
||||
|
||||
# Extract frames from a clip (or full video)
|
||||
python extract_frames.py frames --input clip.mkv --output-dir frames/ \\
|
||||
--width 1280 --threshold 0.3 --fps 0.5 --quality 3
|
||||
|
||||
# Check hardware acceleration support
|
||||
python extract_frames.py check-hwaccel
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def run_command(cmd: list[str], *, capture: bool = True) -> subprocess.CompletedProcess[str]:
|
||||
return subprocess.run(
|
||||
cmd,
|
||||
check=True,
|
||||
text=True,
|
||||
stdout=subprocess.PIPE if capture else None,
|
||||
stderr=subprocess.PIPE if capture else None,
|
||||
)
|
||||
|
||||
|
||||
def check_hwaccel() -> dict:
|
||||
"""Check available hardware acceleration methods."""
|
||||
if not shutil.which("ffmpeg"):
|
||||
return {"error": "ffmpeg not found"}
|
||||
proc = subprocess.run(
|
||||
["ffmpeg", "-hide_banner", "-hwaccels"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
hwaccels = [
|
||||
line.strip()
|
||||
for line in proc.stdout.splitlines()
|
||||
if line.strip() and not line.startswith("Hardware")
|
||||
]
|
||||
return {"hwaccels": hwaccels, "cuda": "cuda" in hwaccels, "vaapi": "vaapi" in hwaccels}
|
||||
|
||||
|
||||
def cmd_clip(args: argparse.Namespace) -> None:
|
||||
"""Extract a clip using -c copy (keyframe-accurate)."""
|
||||
video = Path(args.video).expanduser().resolve()
|
||||
output = Path(args.output).expanduser().resolve()
|
||||
|
||||
if not video.exists():
|
||||
sys.exit(f"Video not found: {video}")
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-hide_banner", "-y",
|
||||
"-ss", str(args.start),
|
||||
"-to", str(args.end),
|
||||
"-i", str(video),
|
||||
"-c", "copy",
|
||||
"-map", "0:v:0",
|
||||
str(output),
|
||||
]
|
||||
|
||||
try:
|
||||
run_command(cmd)
|
||||
print(json.dumps({
|
||||
"status": "ok",
|
||||
"clip_path": str(output),
|
||||
"start": args.start,
|
||||
"end": args.end,
|
||||
"duration": args.end - args.start,
|
||||
}))
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(json.dumps({"status": "error", "stderr": e.stderr}))
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def cmd_frames(args: argparse.Namespace) -> None:
|
||||
"""Extract frames from a video/clip with scene detection + uniform sampling."""
|
||||
input_path = Path(args.input).expanduser().resolve()
|
||||
output_dir = Path(args.output_dir).expanduser().resolve()
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if not input_path.exists():
|
||||
sys.exit(f"Input not found: {input_path}")
|
||||
|
||||
# Build the scale filter
|
||||
scale_filter = f"scale={args.width}:-2"
|
||||
|
||||
# Build the select filter: scene change OR uniform sampling
|
||||
# fps filter handles uniform sampling; select filter handles scene detection
|
||||
# We use fps for uniform sampling and select for scene changes, combined
|
||||
select_filter = f"select='gt(scene\\,{args.threshold})'"
|
||||
vf_parts = [scale_filter]
|
||||
|
||||
if args.fps > 0:
|
||||
# fps filter for uniform sampling
|
||||
vf_parts.append(f"fps={args.fps}")
|
||||
|
||||
# Add scene detection select filter
|
||||
vf_parts.append(select_filter)
|
||||
|
||||
# Add showinfo for scene score logging
|
||||
if args.showinfo:
|
||||
vf_parts.append("showinfo")
|
||||
|
||||
vf = ",".join(vf_parts)
|
||||
|
||||
# Build command
|
||||
cmd = ["ffmpeg", "-hide_banner", "-y"]
|
||||
|
||||
# Hardware acceleration
|
||||
if args.hwaccel:
|
||||
hwaccel_info = check_hwaccel()
|
||||
if hwaccel_info.get("cuda") and args.hwaccel == "cuda":
|
||||
cmd.extend(["-hwaccel", "cuda"])
|
||||
elif args.hwaccel != "auto":
|
||||
# Try the requested hwaccel, fall back to none
|
||||
cmd.extend(["-hwaccel", args.hwaccel])
|
||||
|
||||
cmd.extend(["-i", str(input_path)])
|
||||
cmd.extend(["-vf", vf])
|
||||
cmd.extend(["-vsync", "vfr"])
|
||||
cmd.extend(["-q:v", str(args.quality)])
|
||||
|
||||
frame_pattern = str(output_dir / "frame_%04d.jpg")
|
||||
cmd.append(frame_pattern)
|
||||
|
||||
# Capture showinfo log
|
||||
if args.showinfo:
|
||||
log_file = output_dir / "showinfo.log"
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
log_file.write_text(result.stderr)
|
||||
if result.returncode != 0:
|
||||
print(json.dumps({"status": "error", "stderr": result.stderr[-2000:]}))
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print(json.dumps({"status": "error", "error": str(e)}))
|
||||
sys.exit(1)
|
||||
else:
|
||||
try:
|
||||
run_command(cmd)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(json.dumps({"status": "error", "stderr": e.stderr[-2000:] if e.stderr else ""}))
|
||||
sys.exit(1)
|
||||
|
||||
# Count extracted frames
|
||||
frame_files = sorted(output_dir.glob("frame_*.jpg"))
|
||||
print(json.dumps({
|
||||
"status": "ok",
|
||||
"frames_dir": str(output_dir),
|
||||
"frame_count": len(frame_files),
|
||||
"width": args.width,
|
||||
"threshold": args.threshold,
|
||||
"fps": args.fps,
|
||||
"quality": args.quality,
|
||||
}))
|
||||
|
||||
|
||||
def cmd_check_hwaccel(args: argparse.Namespace) -> None:
|
||||
"""Print available hardware acceleration methods."""
|
||||
info = check_hwaccel()
|
||||
print(json.dumps(info, indent=2))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Extract clips and frames for video edit planning.")
|
||||
sub = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
# clip
|
||||
p_clip = sub.add_parser("clip", help="Extract a time range as a clip")
|
||||
p_clip.add_argument("--video", required=True, help="Source video path")
|
||||
p_clip.add_argument("--start", required=True, type=float, help="Start time in seconds")
|
||||
p_clip.add_argument("--end", required=True, type=float, help="End time in seconds")
|
||||
p_clip.add_argument("--output", required=True, help="Output clip path")
|
||||
|
||||
# frames
|
||||
p_frames = sub.add_parser("frames", help="Extract frames from a video/clip")
|
||||
p_frames.add_argument("--input", required=True, help="Input video/clip path")
|
||||
p_frames.add_argument("--output-dir", required=True, help="Output directory for frames")
|
||||
p_frames.add_argument("--width", type=int, default=1280, help="Output width in pixels (height auto-scaled)")
|
||||
p_frames.add_argument("--threshold", type=float, default=0.3, help="Scene change threshold (0-1)")
|
||||
p_frames.add_argument("--fps", type=float, default=0.5, help="Uniform sampling fps (0 = scene-only)")
|
||||
p_frames.add_argument("--quality", type=int, default=3, help="JPEG quality (2-31, lower=better)")
|
||||
p_frames.add_argument("--hwaccel", default="cuda", help="Hardware acceleration (cuda/vaapi/auto/none)")
|
||||
p_frames.add_argument("--showinfo", action="store_true", default=True, help="Log scene scores to showinfo.log")
|
||||
|
||||
# check-hwaccel
|
||||
sub.add_parser("check-hwaccel", help="Check available hardware acceleration")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
commands = {
|
||||
"clip": cmd_clip,
|
||||
"frames": cmd_frames,
|
||||
"check-hwaccel": cmd_check_hwaccel,
|
||||
}
|
||||
commands[args.command](args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,311 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Manage the video-edit-planner SQLite index.
|
||||
|
||||
The index file lives next to the video as <video_stem>.vedit.db.
|
||||
It tracks transcriptions, clips, and frames to avoid duplicate processing.
|
||||
|
||||
Usage:
|
||||
python manage_index.py init --video VIDEO_PATH
|
||||
python manage_index.py add-transcription --json-path PATH --track-index N --duration S
|
||||
python manage_index.py get-transcription
|
||||
python manage_index.py add-clip --start S --end E --path PATH --reason "..."
|
||||
python manage_index.py add-frames --clip-id N --frames-dir DIR --method scene
|
||||
python manage_index.py list
|
||||
python manage_index.py check-range --start S --end E
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sqlite3
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS transcription (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
json_path TEXT NOT NULL,
|
||||
track_index INTEGER NOT NULL,
|
||||
created_at REAL NOT NULL,
|
||||
duration_s REAL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS clips (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
start_time REAL NOT NULL,
|
||||
end_time REAL NOT NULL,
|
||||
path TEXT NOT NULL,
|
||||
created_at REAL NOT NULL,
|
||||
reason TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS frames (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
clip_id INTEGER NOT NULL REFERENCES clips(id),
|
||||
timestamp REAL NOT NULL,
|
||||
path TEXT NOT NULL,
|
||||
scene_score REAL,
|
||||
method TEXT NOT NULL,
|
||||
created_at REAL NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_frames_clip_id ON frames(clip_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_clips_time ON clips(start_time, end_time);
|
||||
"""
|
||||
|
||||
|
||||
def db_path_for_video(video_path: Path) -> Path:
|
||||
"""Return the index database path for a given video file."""
|
||||
return video_path.parent / f"{video_path.stem}.vedit.db"
|
||||
|
||||
|
||||
def connect(video_path: Path) -> sqlite3.Connection:
|
||||
db = db_path_for_video(video_path)
|
||||
if not db.exists():
|
||||
sys.exit(f"Index not found: {db}. Run 'init' first.")
|
||||
conn = sqlite3.connect(str(db))
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
return conn
|
||||
|
||||
|
||||
def cmd_init(args: argparse.Namespace) -> None:
|
||||
video = Path(args.video).expanduser().resolve()
|
||||
db = db_path_for_video(video)
|
||||
conn = sqlite3.connect(str(db))
|
||||
conn.executescript(SCHEMA)
|
||||
conn.close()
|
||||
print(json.dumps({"status": "ok", "db_path": str(db)}))
|
||||
|
||||
|
||||
def cmd_add_transcription(args: argparse.Namespace) -> None:
|
||||
video = Path(args.video).expanduser().resolve()
|
||||
conn = connect(video)
|
||||
conn.execute(
|
||||
"INSERT INTO transcription (json_path, track_index, created_at, duration_s) VALUES (?, ?, ?, ?)",
|
||||
(args.json_path, args.track_index, time.time(), args.duration),
|
||||
)
|
||||
conn.commit()
|
||||
cur = conn.execute("SELECT last_insert_rowid()")
|
||||
row_id = cur.fetchone()[0]
|
||||
conn.close()
|
||||
print(json.dumps({"status": "ok", "id": row_id}))
|
||||
|
||||
|
||||
def cmd_get_transcription(args: argparse.Namespace) -> None:
|
||||
video = Path(args.video).expanduser().resolve()
|
||||
conn = connect(video)
|
||||
rows = conn.execute("SELECT * FROM transcription ORDER BY created_at DESC").fetchall()
|
||||
conn.close()
|
||||
if not rows:
|
||||
print(json.dumps({"found": False}))
|
||||
return
|
||||
result = [dict(r) for r in rows]
|
||||
print(json.dumps({"found": True, "records": result}, ensure_ascii=False))
|
||||
|
||||
|
||||
def cmd_add_clip(args: argparse.Namespace) -> None:
|
||||
video = Path(args.video).expanduser().resolve()
|
||||
conn = connect(video)
|
||||
cur = conn.execute(
|
||||
"INSERT INTO clips (start_time, end_time, path, created_at, reason) VALUES (?, ?, ?, ?, ?)",
|
||||
(args.start, args.end, args.path, time.time(), args.reason),
|
||||
)
|
||||
conn.commit()
|
||||
clip_id = cur.lastrowid
|
||||
conn.close()
|
||||
print(json.dumps({"status": "ok", "clip_id": clip_id}))
|
||||
|
||||
|
||||
def cmd_add_frames(args: argparse.Namespace) -> None:
|
||||
"""Add frame records for all images in a directory, associating with a clip."""
|
||||
video = Path(args.video).expanduser().resolve()
|
||||
frames_dir = Path(args.frames_dir).expanduser().resolve()
|
||||
if not frames_dir.is_dir():
|
||||
sys.exit(f"Frames directory not found: {frames_dir}")
|
||||
|
||||
conn = connect(video)
|
||||
|
||||
# Read scene scores from showinfo log if available
|
||||
scene_scores: dict[str, float] = {}
|
||||
log_file = frames_dir / "showinfo.log"
|
||||
if log_file.exists():
|
||||
for line in log_file.read_text(errors="replace").splitlines():
|
||||
# Parse lavfi.scd.score and lavfi.scd.time from showinfo output
|
||||
score = None
|
||||
t = None
|
||||
for part in line.split():
|
||||
if part.startswith("score:"):
|
||||
score = float(part.split(":")[1])
|
||||
elif part.startswith("pts_time:"):
|
||||
t = float(part.split(":")[1])
|
||||
if score is not None and t is not None:
|
||||
scene_scores[f"{t:.3f}"] = score
|
||||
|
||||
# Parse timestamps from filenames: frame_XXXX.jpg
|
||||
# The agent should pass --fps and --start to compute timestamps
|
||||
fps = args.fps if args.fps else 0.5
|
||||
start_offset = args.start if args.start else 0.0
|
||||
|
||||
frame_files = sorted(frames_dir.glob("frame_*.jpg"))
|
||||
if not frame_files:
|
||||
# Try other extensions
|
||||
frame_files = sorted(frames_dir.glob("frame_*.png"))
|
||||
|
||||
if not frame_files:
|
||||
sys.exit(f"No frame files found in {frames_dir}")
|
||||
|
||||
count = 0
|
||||
for i, frame_path in enumerate(frame_files):
|
||||
# Estimate timestamp: frame number * (1/fps) + start_offset
|
||||
# This is approximate; the agent can refine with --timestamps
|
||||
timestamp = start_offset + (i / fps)
|
||||
|
||||
# Try to match with scene score
|
||||
scene_score = scene_scores.get(f"{timestamp:.3f}")
|
||||
|
||||
conn.execute(
|
||||
"INSERT INTO frames (clip_id, timestamp, path, scene_score, method, created_at) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
(args.clip_id, timestamp, str(frame_path), scene_score, args.method, time.time()),
|
||||
)
|
||||
count += 1
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
print(json.dumps({"status": "ok", "frames_added": count}))
|
||||
|
||||
|
||||
def cmd_list(args: argparse.Namespace) -> None:
|
||||
video = Path(args.video).expanduser().resolve()
|
||||
conn = connect(video)
|
||||
|
||||
transcriptions = conn.execute("SELECT * FROM transcription ORDER BY created_at DESC").fetchall()
|
||||
clips = conn.execute("SELECT * FROM clips ORDER BY start_time").fetchall()
|
||||
|
||||
result = {"transcriptions": [dict(r) for r in transcriptions], "clips": []}
|
||||
for clip in clips:
|
||||
clip_dict = dict(clip)
|
||||
frames = conn.execute(
|
||||
"SELECT * FROM frames WHERE clip_id = ? ORDER BY timestamp", (clip["id"],)
|
||||
).fetchall()
|
||||
clip_dict["frames"] = [dict(f) for f in frames]
|
||||
result["clips"].append(clip_dict)
|
||||
|
||||
conn.close()
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
def cmd_check_range(args: argparse.Namespace) -> None:
|
||||
"""Check if a time range overlaps with any existing clips."""
|
||||
video = Path(args.video).expanduser().resolve()
|
||||
conn = connect(video)
|
||||
# Find clips that overlap with the requested range
|
||||
rows = conn.execute(
|
||||
"""SELECT * FROM clips
|
||||
WHERE start_time < ? AND end_time > ?
|
||||
ORDER BY start_time""",
|
||||
(args.end, args.start),
|
||||
).fetchall()
|
||||
conn.close()
|
||||
if not rows:
|
||||
print(json.dumps({"found": False}))
|
||||
else:
|
||||
print(json.dumps({"found": True, "clips": [dict(r) for r in rows]}, ensure_ascii=False))
|
||||
|
||||
|
||||
def cmd_clean(args: argparse.Namespace) -> None:
|
||||
"""Remove frame and clip records for files that no longer exist on disk."""
|
||||
video = Path(args.video).expanduser().resolve()
|
||||
conn = connect(video)
|
||||
|
||||
# Check frames
|
||||
frame_rows = conn.execute("SELECT id, path FROM frames").fetchall()
|
||||
frames_removed = 0
|
||||
for row in frame_rows:
|
||||
if not Path(row["path"]).exists():
|
||||
conn.execute("DELETE FROM frames WHERE id = ?", (row["id"],))
|
||||
frames_removed += 1
|
||||
|
||||
# Check clips
|
||||
clip_rows = conn.execute("SELECT id, path FROM clips").fetchall()
|
||||
clips_removed = 0
|
||||
for row in clip_rows:
|
||||
if not Path(row["path"]).exists():
|
||||
conn.execute("DELETE FROM frames WHERE clip_id = ?", (row["id"],))
|
||||
conn.execute("DELETE FROM clips WHERE id = ?", (row["id"],))
|
||||
clips_removed += 1
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
print(json.dumps({"status": "ok", "frames_removed": frames_removed, "clips_removed": clips_removed}))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Manage video-edit-planner index database.")
|
||||
sub = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
# init
|
||||
p_init = sub.add_parser("init", help="Initialize index database")
|
||||
p_init.add_argument("--video", required=True, help="Path to the video file")
|
||||
|
||||
# add-transcription
|
||||
p_at = sub.add_parser("add-transcription", help="Add a transcription record")
|
||||
p_at.add_argument("--video", required=True)
|
||||
p_at.add_argument("--json-path", required=True)
|
||||
p_at.add_argument("--track-index", required=True, type=int)
|
||||
p_at.add_argument("--duration", type=float, default=None)
|
||||
|
||||
# get-transcription
|
||||
p_gt = sub.add_parser("get-transcription", help="Get latest transcription record")
|
||||
p_gt.add_argument("--video", required=True)
|
||||
|
||||
# add-clip
|
||||
p_ac = sub.add_parser("add-clip", help="Add a clip record")
|
||||
p_ac.add_argument("--video", required=True)
|
||||
p_ac.add_argument("--start", required=True, type=float, help="Start time in seconds")
|
||||
p_ac.add_argument("--end", required=True, type=float, help="End time in seconds")
|
||||
p_ac.add_argument("--path", required=True, help="Path to the clip file")
|
||||
p_ac.add_argument("--reason", default=None)
|
||||
|
||||
# add-frames
|
||||
p_af = sub.add_parser("add-frames", help="Add frame records from a directory")
|
||||
p_af.add_argument("--video", required=True)
|
||||
p_af.add_argument("--clip-id", required=True, type=int)
|
||||
p_af.add_argument("--frames-dir", required=True)
|
||||
p_af.add_argument("--method", required=True, choices=["scene", "sample", "both"])
|
||||
p_af.add_argument("--fps", type=float, default=0.5, help="Sampling fps for timestamp estimation")
|
||||
p_af.add_argument("--start", type=float, default=0.0, help="Start offset in seconds")
|
||||
|
||||
# list
|
||||
p_list = sub.add_parser("list", help="List all records")
|
||||
p_list.add_argument("--video", required=True)
|
||||
|
||||
# check-range
|
||||
p_cr = sub.add_parser("check-range", help="Check if a time range has existing clips")
|
||||
p_cr.add_argument("--video", required=True)
|
||||
p_cr.add_argument("--start", required=True, type=float)
|
||||
p_cr.add_argument("--end", required=True, type=float)
|
||||
|
||||
# clean
|
||||
p_clean = sub.add_parser("clean", help="Remove records for files that no longer exist")
|
||||
p_clean.add_argument("--video", required=True)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
commands = {
|
||||
"init": cmd_init,
|
||||
"add-transcription": cmd_add_transcription,
|
||||
"get-transcription": cmd_get_transcription,
|
||||
"add-clip": cmd_add_clip,
|
||||
"add-frames": cmd_add_frames,
|
||||
"list": cmd_list,
|
||||
"check-range": cmd_check_range,
|
||||
"clean": cmd_clean,
|
||||
}
|
||||
commands[args.command](args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1 @@
|
||||
3.13
|
||||
@@ -0,0 +1,365 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
REGULAR_MODEL_CONFIGS: dict[str, dict[str, Any]] = {
|
||||
"sensevoice": {
|
||||
"model": "iic/SenseVoiceSmall",
|
||||
"vad_model": "fsmn-vad",
|
||||
"vad_kwargs": {"max_single_segment_time": 30000},
|
||||
},
|
||||
"paraformer": {
|
||||
"model": "paraformer-zh",
|
||||
"vad_model": "fsmn-vad",
|
||||
"punc_model": "ct-punc",
|
||||
},
|
||||
"paraformer-en": {
|
||||
"model": "paraformer-en",
|
||||
"vad_model": "fsmn-vad",
|
||||
},
|
||||
}
|
||||
|
||||
NANO_MODEL_CONFIG: dict[str, Any] = {
|
||||
"model": "FunAudioLLM/Fun-ASR-Nano-2512",
|
||||
"vad_model": "fsmn-vad",
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class AudioTrack:
|
||||
index: int
|
||||
codec: str | None
|
||||
channels: int | None
|
||||
channel_layout: str | None
|
||||
language: str | None
|
||||
title: str | None
|
||||
|
||||
|
||||
def require_command(name: str) -> None:
|
||||
if shutil.which(name) is None:
|
||||
raise SystemExit(f"缺少命令: {name}")
|
||||
|
||||
|
||||
def run_command(cmd: list[str], *, capture: bool = True) -> subprocess.CompletedProcess[str]:
|
||||
return subprocess.run(
|
||||
cmd,
|
||||
check=True,
|
||||
text=True,
|
||||
stdout=subprocess.PIPE if capture else None,
|
||||
stderr=subprocess.PIPE if capture else None,
|
||||
)
|
||||
|
||||
|
||||
def list_audio_tracks(media_path: Path) -> list[AudioTrack]:
|
||||
require_command("ffprobe")
|
||||
proc = run_command(
|
||||
[
|
||||
"ffprobe",
|
||||
"-v",
|
||||
"error",
|
||||
"-select_streams",
|
||||
"a",
|
||||
"-show_entries",
|
||||
"stream=index,codec_name,channels,channel_layout:stream_tags=language,title",
|
||||
"-of",
|
||||
"json",
|
||||
str(media_path),
|
||||
]
|
||||
)
|
||||
data = json.loads(proc.stdout)
|
||||
tracks: list[AudioTrack] = []
|
||||
for stream in data.get("streams", []):
|
||||
tags = stream.get("tags") or {}
|
||||
tracks.append(
|
||||
AudioTrack(
|
||||
index=int(stream["index"]),
|
||||
codec=stream.get("codec_name"),
|
||||
channels=stream.get("channels"),
|
||||
channel_layout=stream.get("channel_layout"),
|
||||
language=tags.get("language"),
|
||||
title=tags.get("title"),
|
||||
)
|
||||
)
|
||||
return tracks
|
||||
|
||||
|
||||
def print_audio_tracks(media_path: Path) -> None:
|
||||
tracks = list_audio_tracks(media_path)
|
||||
if not tracks:
|
||||
print("未找到音频轨")
|
||||
return
|
||||
for i, track in enumerate(tracks, 1):
|
||||
parts = [
|
||||
f"#{i}",
|
||||
f"stream_index={track.index}",
|
||||
f"codec={track.codec or '?'}",
|
||||
f"channels={track.channels or '?'}",
|
||||
]
|
||||
if track.channel_layout:
|
||||
parts.append(f"layout={track.channel_layout}")
|
||||
if track.language:
|
||||
parts.append(f"lang={track.language}")
|
||||
if track.title:
|
||||
parts.append(f"title={track.title}")
|
||||
print(" ".join(parts))
|
||||
|
||||
|
||||
def resolve_track(media_path: Path, requested: int | None) -> int:
|
||||
tracks = list_audio_tracks(media_path)
|
||||
if not tracks:
|
||||
raise SystemExit(f"未找到音频轨: {media_path}")
|
||||
if requested is None:
|
||||
return tracks[0].index
|
||||
valid_indexes = {track.index for track in tracks}
|
||||
if requested in valid_indexes:
|
||||
return requested
|
||||
raise SystemExit(
|
||||
f"音轨 stream index 不存在: {requested}\n"
|
||||
f"可用音轨: {', '.join(str(t.index) for t in tracks)}\n"
|
||||
"提示: --track 使用 ffprobe/ffmpeg 的 stream index，不是第几条音轨的序号。"
|
||||
)
|
||||
|
||||
|
||||
def extract_audio(media_path: Path, wav_path: Path, track_index: int) -> None:
|
||||
require_command("ffmpeg")
|
||||
wav_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
cmd = [
|
||||
"ffmpeg",
|
||||
"-hide_banner",
|
||||
"-y",
|
||||
"-i",
|
||||
str(media_path),
|
||||
"-map",
|
||||
f"0:{track_index}",
|
||||
"-vn",
|
||||
"-ar",
|
||||
"16000",
|
||||
"-ac",
|
||||
"1",
|
||||
"-c:a",
|
||||
"pcm_s16le",
|
||||
str(wav_path),
|
||||
]
|
||||
subprocess.run(cmd, check=True)
|
||||
|
||||
|
||||
def clean_text(text: str, *, sensevoice: bool) -> str:
|
||||
if sensevoice:
|
||||
try:
|
||||
from funasr.utils.postprocess_utils import rich_transcription_postprocess
|
||||
|
||||
text = rich_transcription_postprocess(text)
|
||||
except Exception:
|
||||
# Keep transcription usable even if FunASR changes this helper.
|
||||
pass
|
||||
return re.sub(r"<\|[^|]*\|>", "", text).strip()
|
||||
|
||||
|
||||
def get_audio_duration(audio_path: Path) -> float | None:
|
||||
try:
|
||||
import soundfile as sf
|
||||
|
||||
return round(float(sf.info(str(audio_path)).duration), 3)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def normalize_result(
|
||||
raw_result: list[dict[str, Any]],
|
||||
*,
|
||||
audio_path: Path,
|
||||
source_media: Path,
|
||||
track_index: int,
|
||||
model_name: str,
|
||||
model_id: str,
|
||||
language: str | None,
|
||||
diarize: bool,
|
||||
elapsed: float,
|
||||
sensevoice: bool,
|
||||
) -> dict[str, Any]:
|
||||
first = raw_result[0] if raw_result else {}
|
||||
text = clean_text(str(first.get("text", "")), sensevoice=sensevoice)
|
||||
segments: list[dict[str, Any]] = []
|
||||
for seg in first.get("sentence_info") or []:
|
||||
item: dict[str, Any] = {
|
||||
"start": seg.get("start", 0),
|
||||
"end": seg.get("end", 0),
|
||||
"text": clean_text(str(seg.get("sentence") or seg.get("text") or ""), sensevoice=sensevoice),
|
||||
}
|
||||
if diarize and "spk" in seg:
|
||||
item["speaker"] = seg["spk"]
|
||||
segments.append(item)
|
||||
|
||||
output: dict[str, Any] = {
|
||||
"text": text,
|
||||
"segments": segments,
|
||||
"file": source_media.name,
|
||||
"audio_file": audio_path.name,
|
||||
"track": track_index,
|
||||
"model": model_name,
|
||||
"model_id": model_id,
|
||||
"language": language or "auto",
|
||||
"audio_duration_s": get_audio_duration(audio_path),
|
||||
"processing_s": round(elapsed, 3),
|
||||
}
|
||||
for key in ("timestamps", "timestamp"):
|
||||
if key in first:
|
||||
output[key] = first[key]
|
||||
return output
|
||||
|
||||
|
||||
def write_json(path: Path, data: dict[str, Any]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
|
||||
|
||||
def make_output_path(output_dir: Path, media_path: Path, track_index: int, suffix: str) -> Path:
|
||||
return output_dir / f"{media_path.stem}_track{track_index}_{suffix}.json"
|
||||
|
||||
|
||||
def resolve_output_dir(output_dir: Path | None, media_path: Path) -> Path:
|
||||
if output_dir is None:
|
||||
return media_path.parent
|
||||
return output_dir.expanduser().resolve()
|
||||
|
||||
|
||||
def add_common_args(parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument("media", type=Path, help="视频/音频文件路径")
|
||||
parser.add_argument("--track", type=int, default=None, help="ffmpeg stream index；默认第一条音频轨")
|
||||
parser.add_argument("--list-tracks", action="store_true", help="列出音轨后退出")
|
||||
parser.add_argument("--output-dir", "-o", type=Path, default=None, help="输出目录；默认源视频/音频文件所在目录")
|
||||
parser.add_argument("--language", "-l", default="auto", help="语言，如 auto/zh/en/ja/ko/yue；默认 auto")
|
||||
parser.add_argument("--device", default=None, help="设备，如 cuda:0/cpu；默认自动")
|
||||
parser.add_argument("--no-diarize", action="store_true", help="禁用 cam++ 说话人分离")
|
||||
parser.add_argument("--keep-wav", action="store_true", help="保留提取出的 16k wav 文件")
|
||||
parser.add_argument("--wav-dir", type=Path, default=None, help="临时 wav 输出目录；默认系统临时目录")
|
||||
parser.add_argument("--verbose", "-v", action="store_true", help="输出更多进度信息")
|
||||
|
||||
|
||||
def prepare_audio(args: argparse.Namespace) -> tuple[Path, int, tempfile.TemporaryDirectory[str] | None]:
|
||||
media_path: Path = args.media.expanduser().resolve()
|
||||
if not media_path.exists():
|
||||
raise SystemExit(f"文件不存在: {media_path}")
|
||||
|
||||
if args.list_tracks:
|
||||
print_audio_tracks(media_path)
|
||||
raise SystemExit(0)
|
||||
|
||||
track_index = resolve_track(media_path, args.track)
|
||||
temp_dir: tempfile.TemporaryDirectory[str] | None = None
|
||||
if args.keep_wav:
|
||||
wav_dir = args.wav_dir.expanduser().resolve() if args.wav_dir else resolve_output_dir(args.output_dir, media_path)
|
||||
wav_dir.mkdir(parents=True, exist_ok=True)
|
||||
elif args.wav_dir is not None:
|
||||
wav_dir = args.wav_dir.expanduser().resolve()
|
||||
wav_dir.mkdir(parents=True, exist_ok=True)
|
||||
else:
|
||||
temp_dir = tempfile.TemporaryDirectory(prefix="funasr-script-")
|
||||
wav_dir = Path(temp_dir.name)
|
||||
|
||||
wav_path = wav_dir / f"{media_path.stem}_track{track_index}.wav"
|
||||
print(f"[1/2] 提取音轨 stream_index={track_index} -> {wav_path}", file=sys.stderr)
|
||||
extract_audio(media_path, wav_path, track_index)
|
||||
return wav_path, track_index, temp_dir
|
||||
|
||||
|
||||
def auto_device(device: str | None) -> str:
|
||||
if device:
|
||||
return device
|
||||
import torch
|
||||
|
||||
return "cuda:0" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
|
||||
def transcribe_with_config(
|
||||
*,
|
||||
wav_path: Path,
|
||||
media_path: Path,
|
||||
track_index: int,
|
||||
output_path: Path,
|
||||
model_name: str,
|
||||
config: dict[str, Any],
|
||||
language: str | None,
|
||||
device: str | None,
|
||||
diarize: bool,
|
||||
use_itn: bool,
|
||||
sensevoice: bool,
|
||||
verbose: bool,
|
||||
) -> dict[str, Any]:
|
||||
from funasr import AutoModel
|
||||
|
||||
config = config.copy()
|
||||
if diarize and "spk_model" not in config:
|
||||
config["spk_model"] = "cam++"
|
||||
resolved_device = auto_device(device)
|
||||
|
||||
print(f"[2/2] 加载模型 {model_name} ({config['model']}) on {resolved_device}", file=sys.stderr)
|
||||
load_start = time.time()
|
||||
model = AutoModel(device=resolved_device, disable_update=True, **config)
|
||||
if verbose:
|
||||
print(f"模型加载耗时: {time.time() - load_start:.1f}s", file=sys.stderr)
|
||||
|
||||
gen_kw: dict[str, Any] = {"input": str(wav_path), "batch_size": 1}
|
||||
if language and language != "auto":
|
||||
gen_kw["language"] = language
|
||||
if use_itn:
|
||||
gen_kw["use_itn"] = True
|
||||
|
||||
start = time.time()
|
||||
raw_result = model.generate(**gen_kw)
|
||||
elapsed = time.time() - start
|
||||
output = normalize_result(
|
||||
raw_result,
|
||||
audio_path=wav_path,
|
||||
source_media=media_path,
|
||||
track_index=track_index,
|
||||
model_name=model_name,
|
||||
model_id=str(config["model"]),
|
||||
language=language,
|
||||
diarize=diarize,
|
||||
elapsed=elapsed,
|
||||
sensevoice=sensevoice,
|
||||
)
|
||||
write_json(output_path, output)
|
||||
print(f"完成: {output_path}", file=sys.stderr)
|
||||
print(f"转录耗时: {elapsed:.2f}s", file=sys.stderr)
|
||||
return output
|
||||
|
||||
|
||||
def run_pipeline(args: argparse.Namespace, *, model_name: str, config: dict[str, Any], suffix: str, use_itn: bool, sensevoice: bool) -> Path:
|
||||
media_path = args.media.expanduser().resolve()
|
||||
temp_dir: tempfile.TemporaryDirectory[str] | None = None
|
||||
try:
|
||||
wav_path, track_index, temp_dir = prepare_audio(args)
|
||||
output_path = make_output_path(resolve_output_dir(args.output_dir, media_path), media_path, track_index, suffix)
|
||||
transcribe_with_config(
|
||||
wav_path=wav_path,
|
||||
media_path=media_path,
|
||||
track_index=track_index,
|
||||
output_path=output_path,
|
||||
model_name=model_name,
|
||||
config=config,
|
||||
language=args.language,
|
||||
device=args.device,
|
||||
diarize=not args.no_diarize,
|
||||
use_itn=use_itn,
|
||||
sensevoice=sensevoice,
|
||||
verbose=args.verbose,
|
||||
)
|
||||
print(output_path)
|
||||
return output_path
|
||||
finally:
|
||||
if temp_dir is not None:
|
||||
temp_dir.cleanup()
|
||||
@@ -0,0 +1,30 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
from funasr_common import REGULAR_MODEL_CONFIGS, add_common_args, run_pipeline
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="funasr-fast",
|
||||
description="提取指定音轨并用 SenseVoice 快速转录。",
|
||||
)
|
||||
add_common_args(parser)
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = build_parser().parse_args()
|
||||
run_pipeline(
|
||||
args,
|
||||
model_name="sensevoice",
|
||||
config=REGULAR_MODEL_CONFIGS["sensevoice"],
|
||||
suffix="sensevoice",
|
||||
use_itn=True,
|
||||
sensevoice=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,30 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
from funasr_common import NANO_MODEL_CONFIG, add_common_args, run_pipeline
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="funasr-nano",
|
||||
description="提取指定音轨并用 Fun-ASR-Nano-2512 转录。",
|
||||
)
|
||||
add_common_args(parser)
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = build_parser().parse_args()
|
||||
run_pipeline(
|
||||
args,
|
||||
model_name="fun-asr-nano",
|
||||
config=NANO_MODEL_CONFIG,
|
||||
suffix="fun-asr-nano",
|
||||
use_itn=True,
|
||||
sensevoice=False,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,37 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
from funasr_common import REGULAR_MODEL_CONFIGS, add_common_args, run_pipeline
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="funasr-regular",
|
||||
description="提取指定音轨并用常规 FunASR 管线转录（SenseVoice/Paraformer）。",
|
||||
)
|
||||
add_common_args(parser)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
choices=sorted(REGULAR_MODEL_CONFIGS),
|
||||
default="sensevoice",
|
||||
help="常规 FunASR 模型，默认 sensevoice",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = build_parser().parse_args()
|
||||
config = REGULAR_MODEL_CONFIGS[args.model]
|
||||
run_pipeline(
|
||||
args,
|
||||
model_name=args.model,
|
||||
config=config,
|
||||
suffix=args.model,
|
||||
use_itn=args.model == "sensevoice",
|
||||
sensevoice=args.model == "sensevoice",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from funasr_nano import main
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,25 @@
|
||||
[project]
|
||||
name = "funasr-script"
|
||||
version = "0.1.0"
|
||||
description = "Small local FunASR transcription scripts for extracting a video audio track and writing JSON transcripts."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.13"
|
||||
dependencies = [
|
||||
"funasr>=1.3.14",
|
||||
"soundfile>=0.13.1",
|
||||
"torch>=2.13.0",
|
||||
"torchaudio>=2.11.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
funasr-script = "funasr_nano:main"
|
||||
funasr-nano = "funasr_nano:main"
|
||||
funasr-fast = "funasr_fast:main"
|
||||
funasr-regular = "funasr_regular:main"
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=69"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[tool.setuptools]
|
||||
py-modules = ["funasr_common", "funasr_regular", "funasr_fast", "funasr_nano", "main"]
|
||||
Generated
+1906
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user