forge/backend/app/services/frame_extractor.py

118 lines
4 KiB
Python

import os
import subprocess
import logging
from uuid import uuid4
from datetime import datetime
from app.database import SessionLocal
from app.models.asset import Asset
from app.config import settings
logger = logging.getLogger(__name__)
def extract_frame(asset_id: str, timestamp: float):
"""
Extract a frame from a video asset at a specific timestamp.
"""
print(f"DEBUG: Extracting frame for asset {asset_id} at {timestamp}")
db = SessionLocal()
try:
# Get input asset
from uuid import UUID
try:
uuid_id = UUID(str(asset_id))
except ValueError:
print(f"DEBUG: Invalid UUID string: {asset_id}")
raise ValueError("Invalid asset ID format")
asset = db.query(Asset).filter(Asset.id == uuid_id).first()
if not asset:
print(f"DEBUG: Asset {asset_id} not found in DB")
raise ValueError("Asset not found")
if not asset.file_path or not os.path.exists(asset.file_path):
print(f"DEBUG: File path not found: {asset.file_path}")
raise ValueError(f"Video file not found on disk: {asset.file_path}")
# Generate output filename
# Format: {original_name}_frame_{timestamp}.png
base_name = os.path.splitext(asset.original_filename)[0]
# Clean timestamp format to be safe (replace . with -)
time_str = f"{timestamp:.3f}".replace('.', '-')
filename = f"{base_name}_frame_{time_str}_{uuid4().hex[:6]}.png"
storage_path = os.path.join(settings.storage_path, "images")
os.makedirs(storage_path, exist_ok=True)
output_path = os.path.join(storage_path, filename)
print(f"DEBUG: Output path: {output_path}")
# Build ffmpeg command
# -ss before -i for faster seeking
# -vframes 1 to get one frame
# -q:v 2 for high quality jpg, but we want png so usually just default or compression level
# PNG is lossless by default in ffmpeg usually.
cmd = [
'ffmpeg',
'-y', # Overwrite
'-ss', str(timestamp),
'-i', asset.file_path,
'-vframes', '1',
output_path
]
logger.info(f"Extracting frame with command: {' '.join(cmd)}")
print(f"DEBUG: Running command: {cmd}")
result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode != 0:
logger.error(f"FFmpeg failed: {result.stderr}")
print(f"DEBUG: FFmpeg failed stderr: {result.stderr}")
print(f"DEBUG: FFmpeg failed stdout: {result.stdout}")
raise ValueError(f"Frame extraction failed: {result.stderr}")
if not os.path.exists(output_path):
print(f"DEBUG: Output file missing after success return code")
raise ValueError("Output file was not created")
# Get file size
file_size = os.path.getsize(output_path)
# Get dimensions if possible (assume same as video or read it)
# We can use Pillow if installed, or just use input video dims
width = asset.width
height = asset.height
# Determine mime type
mime_type = "image/png"
# Create new asset
new_asset = Asset(
user_id=asset.user_id,
project_id=asset.project_id,
original_filename=filename,
stored_filename=filename,
file_path=output_path,
file_type="image",
mime_type=mime_type,
file_size_bytes=file_size,
width=width,
height=height,
source_module="frame_extractor",
parent_asset_id=asset.id,
asset_metadata={
"source_video_id": str(asset.id),
"timestamp": timestamp
}
)
db.add(new_asset)
db.commit()
db.refresh(new_asset)
return new_asset
finally:
db.close()