291 lines
11 KiB
Python
291 lines
11 KiB
Python
"""
|
||
Investigation Script 1: Video Upload and Timeline/Dialogue Extraction using Gemini
|
||
|
||
This script demonstrates:
|
||
1. Uploading video files to Google Gemini API
|
||
2. Extracting dialogue/transcription from video
|
||
3. Identifying timestamps for key moments
|
||
4. Structuring the output for video editing pipeline
|
||
|
||
Target: Process interview videos (№１_基本情報.mp4, №２_ライフイベント.mp4)
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Dict, Any, List
|
||
from dotenv import load_dotenv
|
||
from google import genai
|
||
from rich.console import Console
|
||
from rich.progress import Progress, SpinnerColumn, TextColumn
|
||
from rich.panel import Panel
|
||
from rich.json import JSON
|
||
|
||
# Load environment variables
|
||
load_dotenv()
|
||
|
||
# Initialize Rich console for beautiful output
|
||
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 VideoExtractor:
|
||
"""Extract timeline and dialogue from video using Gemini"""
|
||
|
||
def __init__(self):
|
||
# Use Gemini 3 Pro for video understanding
|
||
self.model = 'gemini-3-pro-preview' # Latest model supporting video
|
||
console.print("[green]✓[/green] Initialized Gemini 3 Pro model")
|
||
|
||
def upload_video(self, video_path: str) -> Any:
|
||
"""
|
||
Upload video file to Gemini API
|
||
|
||
Args:
|
||
video_path: Path to video file
|
||
|
||
Returns:
|
||
Uploaded file object
|
||
"""
|
||
console.print(f"[cyan]Uploading video:[/cyan] {Path(video_path).name}")
|
||
|
||
with Progress(
|
||
SpinnerColumn(),
|
||
TextColumn("[progress.description]{task.description}"),
|
||
console=console
|
||
) as progress:
|
||
task = progress.add_task("Uploading to Gemini API...", total=None)
|
||
|
||
# Upload the video file with config
|
||
from google.genai import types
|
||
with open(video_path, 'rb') as f:
|
||
upload_config = types.UploadFileConfig(
|
||
mime_type='video/mp4'
|
||
)
|
||
video_file = client.files.upload(
|
||
file=f,
|
||
config=upload_config
|
||
)
|
||
|
||
# Wait for processing to complete
|
||
while video_file.state == "PROCESSING":
|
||
time.sleep(2)
|
||
video_file = client.files.get(name=video_file.name)
|
||
|
||
progress.update(task, completed=True)
|
||
|
||
if video_file.state == "FAILED":
|
||
raise ValueError(f"Video processing failed: {video_file.state}")
|
||
|
||
console.print(f"[green]✓[/green] Video uploaded successfully: {video_file.name}")
|
||
return video_file
|
||
|
||
def extract_timeline_and_dialogue(self, video_file: Any, video_name: str) -> Dict[str, Any]:
|
||
"""
|
||
Extract structured timeline and dialogue from video
|
||
|
||
Args:
|
||
video_file: Uploaded video file object
|
||
video_name: Name of the video for context
|
||
|
||
Returns:
|
||
Dictionary containing timeline segments with timestamps and transcription
|
||
"""
|
||
console.print("\n[cyan]Analyzing video content...[/cyan]")
|
||
|
||
# Prompt for structured extraction (in Japanese for better understanding)
|
||
prompt = f"""
|
||
このビデオは保険コンサルタントと顧客のインタビュー映像です。
|
||
ビデオのタイトル: {video_name}
|
||
|
||
以下の情報を抽出してください：
|
||
|
||
1. **全体の要約**: ビデオの内容を2-3文で要約
|
||
2. **話者情報**: 誰が話しているか（例：コンサルタント、顧客）
|
||
3. **タイムライン**: 重要な場面ごとに以下の情報を含むセグメントに分割
|
||
- 開始時刻（MM:SS形式）
|
||
- 終了時刻（MM:SS形式）
|
||
- 発言内容（日本語の書き起こし）
|
||
- トピック/キーワード（例：「基本情報」「家族構成」「将来の不安」）
|
||
- 感情/トーン（例：「心配」「前向き」「真剣」）
|
||
- 映像使用の推奨度（1-5、5が最も重要）
|
||
|
||
4. **キーワード**: B-Roll素材検索用のキーワードを5-10個抽出（日本語）
|
||
|
||
**出力形式**: 必ずJSON形式で返してください。以下の構造に従ってください：
|
||
|
||
{{
|
||
"summary": "ビデオ全体の要約",
|
||
"speakers": ["話者1", "話者2"],
|
||
"duration": "MM:SS",
|
||
"segments": [
|
||
{{
|
||
"start_time": "00:00",
|
||
"end_time": "00:30",
|
||
"speaker": "コンサルタント",
|
||
"transcript": "発言内容の書き起こし",
|
||
"topic": "トピック名",
|
||
"emotion": "感情/トーン",
|
||
"importance": 4,
|
||
"keywords": ["キーワード1", "キーワード2"]
|
||
}}
|
||
],
|
||
"broll_keywords": ["素材検索用キーワード1", "キーワード2", "..."]
|
||
}}
|
||
"""
|
||
|
||
with Progress(
|
||
SpinnerColumn(),
|
||
TextColumn("[progress.description]{task.description}"),
|
||
console=console
|
||
) as progress:
|
||
task = progress.add_task("Extracting content with Gemini...", total=None)
|
||
|
||
# Generate content from video - use file object directly with prompt
|
||
response = client.models.generate_content(
|
||
model=self.model,
|
||
contents=[
|
||
video_file, # Pass the file object directly
|
||
prompt
|
||
]
|
||
)
|
||
|
||
progress.update(task, completed=True)
|
||
|
||
# Parse the response
|
||
try:
|
||
# Extract token usage information
|
||
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
|
||
}
|
||
|
||
# Display token usage
|
||
console.print(f"[cyan]📊 Token Usage:[/cyan] Input: {token_info['input_tokens']:,} | Output: {token_info['output_tokens']:,} | Total: {token_info['total_tokens']:,}")
|
||
|
||
# Extract JSON from response (handle markdown code blocks if present)
|
||
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"] = {
|
||
"token_usage": token_info,
|
||
"model": self.model,
|
||
"video_name": video_name
|
||
}
|
||
console.print("[green]✓[/green] Successfully extracted timeline and dialogue")
|
||
return result
|
||
except json.JSONDecodeError as e:
|
||
console.print(f"[yellow]⚠[/yellow] Failed to parse JSON, returning raw response")
|
||
return {
|
||
"raw_response": response.text,
|
||
"error": str(e),
|
||
"_metadata": {
|
||
"token_usage": token_info if 'token_info' in locals() else {},
|
||
"model": self.model
|
||
}
|
||
}
|
||
|
||
def save_results(self, results: Dict[str, Any], output_path: str):
|
||
"""Save extraction results to JSON file"""
|
||
output_file = Path(output_path)
|
||
output_file.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
with open(output_file, 'w', encoding='utf-8') as f:
|
||
json.dump(results, f, ensure_ascii=False, indent=2)
|
||
|
||
console.print(f"[green]✓[/green] Results saved to: {output_file}")
|
||
|
||
|
||
def main():
|
||
"""Main investigation workflow"""
|
||
console.print(Panel.fit(
|
||
"[bold cyan]Video Extraction Investigation Script[/bold cyan]\n"
|
||
"Testing Gemini 1.5 Pro for video timeline and dialogue extraction",
|
||
border_style="cyan"
|
||
))
|
||
|
||
# Initialize extractor
|
||
extractor = VideoExtractor()
|
||
|
||
# Define input and output paths
|
||
input_dir = Path("assets/0_input")
|
||
output_dir = Path("assets/1_processed/video_extractions")
|
||
|
||
# Get all MP4 files
|
||
video_files = list(input_dir.glob("*.mp4"))
|
||
|
||
if not video_files:
|
||
console.print("[red]❌ No video files found in assets/0_input/[/red]")
|
||
return
|
||
|
||
console.print(f"\n[cyan]Found {len(video_files)} video(s) to process:[/cyan]")
|
||
for video in video_files:
|
||
console.print(f" • {video.name}")
|
||
|
||
# Process each video
|
||
all_results = {}
|
||
total_tokens_used = 0
|
||
|
||
for video_path in video_files:
|
||
console.print(f"\n{'='*60}")
|
||
console.print(f"[bold]Processing: {video_path.name}[/bold]")
|
||
console.print(f"{'='*60}\n")
|
||
|
||
try:
|
||
# Upload video
|
||
video_file = extractor.upload_video(str(video_path))
|
||
|
||
# Extract timeline and dialogue
|
||
results = extractor.extract_timeline_and_dialogue(video_file, video_path.name)
|
||
|
||
# Display results
|
||
console.print("\n[bold green]Extraction Results:[/bold green]")
|
||
console.print(JSON(json.dumps(results, ensure_ascii=False)))
|
||
|
||
# Save results
|
||
output_file = output_dir / f"{video_path.stem}_extraction.json"
|
||
extractor.save_results(results, str(output_file))
|
||
|
||
all_results[video_path.name] = results
|
||
|
||
# Track total tokens
|
||
if "_metadata" in results and "token_usage" in results["_metadata"]:
|
||
total_tokens_used += results["_metadata"]["token_usage"].get("total_tokens", 0)
|
||
|
||
except Exception as e:
|
||
console.print(f"[red]❌ Error processing {video_path.name}:[/red] {e}")
|
||
continue
|
||
|
||
# Save combined results
|
||
if all_results:
|
||
combined_output = output_dir / "all_extractions.json"
|
||
extractor.save_results(all_results, str(combined_output))
|
||
|
||
console.print("\n" + "="*60)
|
||
console.print("[bold green]✓ Investigation Complete![/bold green]")
|
||
console.print(f"Processed {len(all_results)} video(s)")
|
||
console.print(f"Results saved to: {output_dir}")
|
||
console.print(f"\n[cyan]💰 Total Tokens Consumed: {total_tokens_used:,}[/cyan]")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|