File size: 3,440 Bytes
c641d5f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | """Bounded YouTube download and deterministic contact-sheet extraction."""
from __future__ import annotations
import math
from pathlib import Path
from PIL import Image, ImageChops, ImageDraw, ImageStat
from .youtube import video_id
def download_youtube_video(url: str, cache_dir: Path) -> Path:
"""Download a small MP4 rendition for visual analysis."""
import yt_dlp
destination = cache_dir / "videos" / f"{video_id(url)}.mp4"
if destination.is_file() and destination.stat().st_size:
return destination
destination.parent.mkdir(parents=True, exist_ok=True)
options = {
"format": "best[ext=mp4][height<=480]/best[height<=480]/best",
"outtmpl": str(destination),
"quiet": True,
"noprogress": True,
"no_warnings": True,
"noplaylist": True,
"max_filesize": 100 * 1024 * 1024,
}
with yt_dlp.YoutubeDL(options) as downloader:
info = downloader.extract_info(url, download=True)
actual = Path(downloader.prepare_filename(info))
if actual != destination and actual.exists():
actual.replace(destination)
if not destination.exists():
raise RuntimeError("YouTube download did not produce an MP4 file")
return destination
def extract_contact_sheets(
path: str | Path,
*,
interval_seconds: float = 1.0,
max_frames: int = 128,
cells_per_sheet: int = 16,
scene_threshold: float = 28.0,
) -> list[Image.Image]:
"""Sample by interval and scene changes, returning timestamped contact sheets."""
import imageio_ffmpeg
reader = imageio_ffmpeg.read_frames(str(path), pix_fmt="rgb24")
metadata = next(reader)
width, height = metadata["size"]
fps = float(metadata.get("fps") or 25.0)
stride = max(1, round(fps * interval_seconds))
samples: list[tuple[float, Image.Image]] = []
previous: Image.Image | None = None
try:
for index, raw in enumerate(reader):
frame = Image.frombytes("RGB", (width, height), raw)
scene_change = False
thumbnail = frame.copy()
thumbnail.thumbnail((96, 96))
if previous is not None:
difference = ImageStat.Stat(
ImageChops.difference(thumbnail, previous)
).mean
scene_change = sum(difference) / len(difference) >= scene_threshold
previous = thumbnail
if index % stride and not scene_change:
continue
frame.thumbnail((240, 150))
samples.append((index / fps, frame.copy()))
if len(samples) >= max_frames:
break
finally:
reader.close()
if not samples:
raise RuntimeError("No frames could be extracted from the video")
sheets: list[Image.Image] = []
columns = int(math.sqrt(cells_per_sheet))
rows = math.ceil(cells_per_sheet / columns)
for offset in range(0, len(samples), cells_per_sheet):
batch = samples[offset : offset + cells_per_sheet]
sheet = Image.new("RGB", (columns * 240, rows * 175), "white")
draw = ImageDraw.Draw(sheet)
for cell, (timestamp, frame) in enumerate(batch):
x = (cell % columns) * 240
y = (cell // columns) * 175
sheet.paste(frame, (x, y + 20))
draw.text((x + 4, y + 3), f"{timestamp:.1f}s", fill="black")
sheets.append(sheet)
return sheets
|