413 lines
14 KiB
Python
413 lines
14 KiB
Python
"""
|
||
Investigation Script 3: Subtitle/Transcript Extraction - Granularity Comparison
|
||
|
||
This script tests different approaches for subtitle extraction:
|
||
1. Movie-level (entire video summary)
|
||
2. Second-level (time-stamped subtitles per second/phrase)
|
||
3. Word-level (individual word timestamps)
|
||
|
||
Compares two methods:
|
||
- Method A: Upload video directly to Gemini
|
||
- Method B: Extract audio with ffmpeg, upload audio only
|
||
|
||
Tests which is faster, more accurate, and uses fewer tokens.
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import time
|
||
import subprocess
|
||
from pathlib import Path
|
||
from typing import Dict, Any, List, Optional, Literal
|
||
from dotenv import load_dotenv
|
||
from google import genai
|
||
from google.genai import types
|
||
from rich.console import Console
|
||
from rich.progress import Progress, SpinnerColumn, TextColumn
|
||
from rich.panel import Panel
|
||
from rich.table import Table
|
||
from rich.json import JSON
|
||
|
||
# Load environment variables
|
||
load_dotenv()
|
||
|
||
# Initialize Rich console
|
||
console = Console()
|
||
|
||
# Configure Gemini API
|
||
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
|
||
if not GEMINI_API_KEY:
|
||
console.print("[red]❌ GEMINI_API_KEY not found in .env file[/red]")
|
||
exit(1)
|
||
|
||
# Initialize Gemini client
|
||
client = genai.Client(api_key=GEMINI_API_KEY)
|
||
|
||
|
||
class AudioExtractor:
|
||
"""Extract audio from video using ffmpeg"""
|
||
|
||
@staticmethod
|
||
def extract_audio(video_path: str, output_path: str) -> Optional[str]:
|
||
"""
|
||
Extract audio from video file using ffmpeg
|
||
|
||
Args:
|
||
video_path: Path to video file
|
||
output_path: Path for output audio file
|
||
|
||
Returns:
|
||
Path to extracted audio file, or None if failed
|
||
"""
|
||
try:
|
||
console.print(f"[cyan]Extracting audio from video...[/cyan]")
|
||
|
||
# Use ffmpeg to extract audio as MP3
|
||
cmd = [
|
||
'ffmpeg',
|
||
'-i', video_path,
|
||
'-vn', # No video
|
||
'-acodec', 'libmp3lame', # MP3 codec
|
||
'-q:a', '2', # High quality
|
||
'-y', # Overwrite output
|
||
output_path
|
||
]
|
||
|
||
result = subprocess.run(
|
||
cmd,
|
||
capture_output=True,
|
||
text=True,
|
||
check=True
|
||
)
|
||
|
||
console.print(f"[green]✓[/green] Audio extracted: {Path(output_path).name}")
|
||
return output_path
|
||
|
||
except subprocess.CalledProcessError as e:
|
||
console.print(f"[red]❌ FFmpeg error:[/red] {e.stderr}")
|
||
return None
|
||
except FileNotFoundError:
|
||
console.print("[red]❌ FFmpeg not found. Please install ffmpeg.[/red]")
|
||
return None
|
||
|
||
|
||
class SubtitleExtractor:
|
||
"""Extract subtitles at different granularity levels"""
|
||
|
||
def __init__(self):
|
||
self.model = 'gemini-3-pro-preview'
|
||
console.print("[green]✓[/green] Initialized Subtitle Extractor")
|
||
|
||
def upload_media(self, file_path: str, mime_type: str) -> Any:
|
||
"""Upload video or audio file to Gemini"""
|
||
console.print(f"[cyan]Uploading {mime_type.split('/')[0]}:[/cyan] {Path(file_path).name}")
|
||
|
||
with Progress(
|
||
SpinnerColumn(),
|
||
TextColumn("[progress.description]{task.description}"),
|
||
console=console
|
||
) as progress:
|
||
task = progress.add_task("Uploading to Gemini API...", total=None)
|
||
|
||
with open(file_path, 'rb') as f:
|
||
upload_config = types.UploadFileConfig(mime_type=mime_type)
|
||
media_file = client.files.upload(file=f, config=upload_config)
|
||
|
||
# Wait for processing
|
||
while media_file.state == "PROCESSING":
|
||
time.sleep(2)
|
||
media_file = client.files.get(name=media_file.name)
|
||
|
||
progress.update(task, completed=True)
|
||
|
||
if media_file.state == "FAILED":
|
||
raise ValueError(f"Media processing failed: {media_file.state}")
|
||
|
||
console.print(f"[green]✓[/green] Media uploaded successfully")
|
||
return media_file
|
||
|
||
def extract_movie_level(self, media_file: Any) -> Dict[str, Any]:
|
||
"""
|
||
Extract movie-level subtitle (entire video summary)
|
||
|
||
Returns:
|
||
Dictionary with summary and full transcript
|
||
"""
|
||
prompt = """
|
||
このビデオ/音声の内容を全体として書き起こしてください。
|
||
|
||
以下の情報を抽出してください：
|
||
1. **全体の要約**: ビデオ全体の内容を2-3文で要約
|
||
2. **完全な書き起こし**: すべての発言を時系列で書き起こし
|
||
3. **話者**: 誰が話しているか識別
|
||
|
||
**出力形式**: 必ずJSON形式で返してください：
|
||
{
|
||
"summary": "全体要約",
|
||
"full_transcript": "完全な書き起こしテキスト",
|
||
"speakers": ["話者1", "話者2"],
|
||
"duration": "MM:SS"
|
||
}
|
||
"""
|
||
return self._generate_content(media_file, prompt, "Movie-level")
|
||
|
||
def extract_second_level(self, media_file: Any) -> Dict[str, Any]:
|
||
"""
|
||
Extract second-level subtitles (time-stamped by phrase/sentence)
|
||
|
||
Returns:
|
||
Dictionary with time-stamped subtitle segments
|
||
"""
|
||
prompt = """
|
||
この音声/ビデオの内容を時間軸に沿って正確に書き起こし、字幕形式で出力してください。
|
||
|
||
**要件**:
|
||
- 各発言を開始時刻と終了時刻付きでセグメント化
|
||
- 話者を識別
|
||
- 時刻はMM:SS.mmm形式（秒単位で正確に）
|
||
|
||
**出力形式**: 必ずJSON形式で返してください：
|
||
{
|
||
"subtitles": [
|
||
{
|
||
"index": 1,
|
||
"start_time": "00:00.000",
|
||
"end_time": "00:03.500",
|
||
"speaker": "話者名",
|
||
"text": "発言内容"
|
||
}
|
||
],
|
||
"total_segments": 0
|
||
}
|
||
"""
|
||
return self._generate_content(media_file, prompt, "Second-level")
|
||
|
||
def extract_word_level(self, media_file: Any) -> Dict[str, Any]:
|
||
"""
|
||
Extract word-level subtitles (timestamp for each word)
|
||
|
||
Returns:
|
||
Dictionary with word-by-word timestamps
|
||
"""
|
||
prompt = """
|
||
この音声/ビデオの内容を単語レベルで書き起こし、各単語のタイムスタンプを付けてください。
|
||
|
||
**要件**:
|
||
- 各単語に開始時刻を付与
|
||
- 話者を識別
|
||
- 時刻はMM:SS.mmm形式
|
||
|
||
**出力形式**: 必ずJSON形式で返してください：
|
||
{
|
||
"word_timestamps": [
|
||
{
|
||
"word": "単語",
|
||
"start_time": "00:00.000",
|
||
"speaker": "話者名"
|
||
}
|
||
],
|
||
"total_words": 0
|
||
}
|
||
"""
|
||
return self._generate_content(media_file, prompt, "Word-level")
|
||
|
||
def _generate_content(self, media_file: Any, prompt: str, level: str) -> Dict[str, Any]:
|
||
"""Generate content with token tracking"""
|
||
console.print(f"[cyan]Extracting {level} subtitles...[/cyan]")
|
||
|
||
start_time = time.time()
|
||
|
||
with Progress(
|
||
SpinnerColumn(),
|
||
TextColumn("[progress.description]{task.description}"),
|
||
console=console
|
||
) as progress:
|
||
task = progress.add_task(f"Generating {level} content...", total=None)
|
||
|
||
response = client.models.generate_content(
|
||
model=self.model,
|
||
contents=[media_file, prompt]
|
||
)
|
||
|
||
progress.update(task, completed=True)
|
||
|
||
elapsed_time = time.time() - start_time
|
||
|
||
# Extract token usage
|
||
token_info = {
|
||
"input_tokens": response.usage_metadata.prompt_token_count if hasattr(response, 'usage_metadata') else 0,
|
||
"output_tokens": response.usage_metadata.candidates_token_count if hasattr(response, 'usage_metadata') else 0,
|
||
"total_tokens": response.usage_metadata.total_token_count if hasattr(response, 'usage_metadata') else 0
|
||
}
|
||
|
||
console.print(f"[cyan]📊 Tokens:[/cyan] {token_info['total_tokens']:,} | [cyan]⏱ Time:[/cyan] {elapsed_time:.2f}s")
|
||
|
||
# Parse response
|
||
try:
|
||
response_text = response.text
|
||
if "```json" in response_text:
|
||
json_start = response_text.find("```json") + 7
|
||
json_end = response_text.find("```", json_start)
|
||
response_text = response_text[json_start:json_end].strip()
|
||
elif "```" in response_text:
|
||
json_start = response_text.find("```") + 3
|
||
json_end = response_text.find("```", json_start)
|
||
response_text = response_text[json_start:json_end].strip()
|
||
|
||
result = json.loads(response_text)
|
||
result["_metadata"] = {
|
||
"level": level,
|
||
"token_usage": token_info,
|
||
"processing_time_seconds": elapsed_time,
|
||
"model": self.model
|
||
}
|
||
|
||
console.print(f"[green]✓[/green] {level} extraction complete")
|
||
return result
|
||
|
||
except json.JSONDecodeError as e:
|
||
console.print(f"[yellow]⚠[/yellow] JSON parsing failed, returning raw response")
|
||
return {
|
||
"raw_response": response.text,
|
||
"error": str(e),
|
||
"_metadata": {
|
||
"level": level,
|
||
"token_usage": token_info,
|
||
"processing_time_seconds": elapsed_time
|
||
}
|
||
}
|
||
|
||
|
||
def compare_methods(video_path: str, output_dir: Path, granularity: Literal["movie", "second", "word"]):
|
||
"""
|
||
Compare video vs audio extraction for a specific granularity level
|
||
|
||
Args:
|
||
video_path: Path to video file
|
||
output_dir: Output directory for results
|
||
granularity: Level of granularity to test
|
||
"""
|
||
extractor = SubtitleExtractor()
|
||
audio_extractor = AudioExtractor()
|
||
results = {}
|
||
|
||
console.print(f"\n[bold cyan]Testing {granularity.upper()}-level extraction[/bold cyan]")
|
||
console.print("="*60)
|
||
|
||
# Method A: Direct video upload
|
||
console.print("\n[bold]Method A: Video Upload[/bold]")
|
||
try:
|
||
video_file = extractor.upload_media(video_path, 'video/mp4')
|
||
|
||
if granularity == "movie":
|
||
result_a = extractor.extract_movie_level(video_file)
|
||
elif granularity == "second":
|
||
result_a = extractor.extract_second_level(video_file)
|
||
else: # word
|
||
result_a = extractor.extract_word_level(video_file)
|
||
|
||
results["method_a_video"] = result_a
|
||
|
||
except Exception as e:
|
||
console.print(f"[red]❌ Method A failed:[/red] {e}")
|
||
results["method_a_video"] = {"error": str(e)}
|
||
|
||
# Method B: Audio extraction + upload
|
||
console.print("\n[bold]Method B: Audio Extraction + Upload[/bold]")
|
||
try:
|
||
# Extract audio
|
||
audio_path = str(output_dir / f"{Path(video_path).stem}_audio.mp3")
|
||
extracted_audio = audio_extractor.extract_audio(video_path, audio_path)
|
||
|
||
if extracted_audio:
|
||
audio_file = extractor.upload_media(extracted_audio, 'audio/mpeg')
|
||
|
||
if granularity == "movie":
|
||
result_b = extractor.extract_movie_level(audio_file)
|
||
elif granularity == "second":
|
||
result_b = extractor.extract_second_level(audio_file)
|
||
else: # word
|
||
result_b = extractor.extract_word_level(audio_file)
|
||
|
||
results["method_b_audio"] = result_b
|
||
else:
|
||
results["method_b_audio"] = {"error": "Audio extraction failed"}
|
||
|
||
except Exception as e:
|
||
console.print(f"[red]❌ Method B failed:[/red] {e}")
|
||
results["method_b_audio"] = {"error": str(e)}
|
||
|
||
return results
|
||
|
||
|
||
def main():
|
||
"""Main investigation workflow"""
|
||
console.print(Panel.fit(
|
||
"[bold cyan]Subtitle Extraction Investigation[/bold cyan]\n"
|
||
"Comparing Video vs Audio | Movie/Second/Word Level Granularity",
|
||
border_style="cyan"
|
||
))
|
||
|
||
# Setup paths
|
||
input_dir = Path("assets/0_input")
|
||
output_dir = Path("assets/1_processed/subtitle_investigation")
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# Get first video for testing
|
||
video_files = list(input_dir.glob("*.mp4"))
|
||
|
||
if not video_files:
|
||
console.print("[red]❌ No video files found[/red]")
|
||
return
|
||
|
||
test_video = video_files[0] # Use first video for comparison
|
||
console.print(f"\n[cyan]Test Video:[/cyan] {test_video.name}")
|
||
|
||
# Test all granularity levels
|
||
all_results = {}
|
||
|
||
for granularity in ["movie", "second", "word"]:
|
||
results = compare_methods(str(test_video), output_dir, granularity)
|
||
all_results[granularity] = results
|
||
|
||
# Save individual results
|
||
result_file = output_dir / f"{test_video.stem}_{granularity}_comparison.json"
|
||
with open(result_file, 'w', encoding='utf-8') as f:
|
||
json.dump(results, f, ensure_ascii=False, indent=2)
|
||
|
||
console.print(f"[green]✓[/green] Results saved: {result_file.name}")
|
||
|
||
# Generate comparison summary
|
||
console.print("\n" + "="*60)
|
||
console.print("[bold green]Comparison Summary[/bold green]\n")
|
||
|
||
table = Table(title="Performance Comparison", show_header=True)
|
||
table.add_column("Granularity", style="cyan")
|
||
table.add_column("Method", style="yellow")
|
||
table.add_column("Tokens", justify="right", style="magenta")
|
||
table.add_column("Time (s)", justify="right", style="blue")
|
||
|
||
for granularity, results in all_results.items():
|
||
for method_name, data in results.items():
|
||
if "_metadata" in data:
|
||
meta = data["_metadata"]
|
||
table.add_row(
|
||
granularity.upper(),
|
||
method_name.replace("_", " ").title(),
|
||
f"{meta['token_usage']['total_tokens']:,}",
|
||
f"{meta['processing_time_seconds']:.2f}"
|
||
)
|
||
|
||
console.print(table)
|
||
|
||
# Save complete comparison
|
||
summary_file = output_dir / "complete_comparison.json"
|
||
with open(summary_file, 'w', encoding='utf-8') as f:
|
||
json.dump(all_results, f, ensure_ascii=False, indent=2)
|
||
|
||
console.print(f"\n[green]✓[/green] Complete comparison saved: {summary_file}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|