AI-video-lifeplan-demo/backend/investigate_video_extractio...

291 lines
11 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""
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 (№1_基本情報.mp4, №2_ライフイベント.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()