| """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 |
|
|
