FastApi-ISS/utils/media_processing.py

import io
import os
import re
import subprocess
import tempfile
from typing import Tuple, List
from PIL import Image
import requests

FFMPEG = os.getenv("FFMPEG_BIN", "ffmpeg")
FFPROBE = os.getenv("FFPROBE_BIN", "ffprobe")

# ------------------------
# Shared helpers
# ------------------------

def _bool(name: str, default: bool) -> bool:
    v = os.getenv(name)
    if v is None:
        return default
    return v.strip().lower() in ("1", "true", "yes", "y")

def _int(name: str, default: int) -> int:
    try:
        return int(os.getenv(name, str(default)))
    except Exception:
        return default

def _float(name: str, default: float) -> float:
    try:
        return float(os.getenv(name, str(default)))
    except Exception:
        return default

def _parse_max_size(default: Tuple[int, int] = (1600, 1600)) -> Tuple[int, int]:
    raw = os.getenv("IMAGE_MAX_SIZE", f"{default[0]}x{default[1]}").strip().lower()
    m = re.match(r"^\s*(\d+)\s*[x,]\s*(\d+)\s*$", raw)
    if not m:
        return default
    return (int(m.group(1)), int(m.group(2)))

def download_bytes(url: str, timeout: int = 30) -> bytes:
    r = requests.get(url, stream=True, timeout=timeout)
    r.raise_for_status()
    return r.content

# ------------------------
# Image processing (WebP)
# ------------------------

def process_image_to_webp_single(image_bytes: bytes) -> Tuple[str, bytes, str]:
    """
    Resize image to:
      - scale percent from .env (default 0.85)
      - and cap to max size from .env (default 1600x1600)
    Save as WebP with quality from .env (default 85).

    Returns (variant_name, data, content_type)
    variant_name: 'webp'
    """
    scale_percent = _float("IMAGE_SCALE_PERCENT", 0.85)
    max_w, max_h = _parse_max_size((1600, 1600))
    quality = _int("IMAGE_QUALITY", 85)

    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    orig_w, orig_h = img.size

    # First apply percentage scaling
    target_w = max(1, int(orig_w * scale_percent))
    target_h = max(1, int(orig_h * scale_percent))

    # Then cap to max size if still too large
    ratio_w = max_w / target_w if target_w > max_w else 1.0
    ratio_h = max_h / target_h if target_h > max_h else 1.0
    cap_ratio = min(ratio_w, ratio_h)
    if cap_ratio < 1.0:
        target_w = max(1, int(target_w * cap_ratio))
        target_h = max(1, int(target_h * cap_ratio))

    # Resize using LANCZOS
    resized = img.resize((target_w, target_h), Image.LANCZOS)

    out = io.BytesIO()
    # method=6 is a good quality/speed balance for WebP
    resized.save(out, format="WEBP", quality=quality, method=6)
    return ("webp", out.getvalue(), "image/webp")

# ------------------------
# Video processing (H.264 @ 480p)
# ------------------------

def _probe_streams(path: str) -> dict:
    """Return ffprobe json for streams (video/audio codecs, dimensions)."""
    # Minimal probing; not strictly required but helps with SIMPLE mode decisions
    cmd = [
        FFPROBE, "-v", "error", "-select_streams", "v:0",
        "-show_entries", "stream=codec_name,width,height",
        "-of", "json", path
    ]
    try:
        out = subprocess.run(cmd, check=True, capture_output=True).stdout
        import json
        return json.loads(out.decode("utf-8"))
    except Exception:
        return {}

def transcode_video_480_h264(video_bytes: bytes) -> Tuple[str, bytes, str]:
    """
    Transcode to MP4 (H.264/AAC), height=VIDEO_HEIGHT (default 480), CRF=VIDEO_CRF (default 23).
    Adds +faststart unless disabled. Uses PRESET (default veryfast).
    If VIDEO_SIMPLE_MODE=true:
       - Use 'ultrafast' preset automatically
       - If input already H.264/AAC and <= target height, attempt stream-copy with -movflags +faststart
         (very quick) — otherwise fallback to a minimal transcode.
    Returns (variant_name, data, content_type) where variant_name='480p'
    """
    height = _int("VIDEO_HEIGHT", 480)
    crf = _int("VIDEO_CRF", 23)
    preset = os.getenv("VIDEO_PRESET", "veryfast")
    faststart = _bool("VIDEO_FASTSTART", True)
    simple = _bool("VIDEO_SIMPLE_MODE", False)

    with tempfile.TemporaryDirectory() as td:
        in_path = os.path.join(td, "in.bin")
        out_path = os.path.join(td, f"out-{height}p.mp4")
        with open(in_path, "wb") as f:
            f.write(video_bytes)

        vf = f"scale=-2:{height}"

        # SIMPLE MODE quick path
        if simple:
            # Try to detect if we can safely stream copy
            meta = _probe_streams(in_path)
            v_ok = False
            try:
                st = meta.get("streams", [{}])[0]
                vcodec = st.get("codec_name")
                v_h = int(st.get("height", 10**9))
                v_ok = (vcodec == "h264" and v_h <= height)
            except Exception:
                v_ok = False

            if v_ok:
                # Quick remux with faststart; try copy video/audio
                cmd = [FFMPEG, "-y", "-i", in_path, "-c", "copy"]
                if faststart:
                    cmd += ["-movflags", "+faststart"]
                cmd += [out_path]
                try:
                    subprocess.run(cmd, check=True, capture_output=True)
                    with open(out_path, "rb") as f:
                        return (f"{height}p", f.read(), "video/mp4")
                except Exception:
                    # fall back to quick transcode below
                    pass

            # Quick transcode with ultrafast
            cmd = [
                FFMPEG, "-y",
                "-i", in_path,
                "-vf", vf,
                "-c:v", "libx264",
                "-preset", "ultrafast",   # force quickest in simple mode
                "-crf", str(crf),
                "-c:a", "aac",
                "-b:a", "128k",
            ]
            if faststart:
                cmd += ["-movflags", "+faststart"]
            cmd += [out_path]
            subprocess.run(cmd, check=True, capture_output=True)
            with open(out_path, "rb") as f:
                return (f"{height}p", f.read(), "video/mp4")

        # Normal path (preset configurable)
        cmd = [
            FFMPEG, "-y",
            "-i", in_path,
            "-vf", vf,
            "-c:v", "libx264",
            "-preset", preset,
            "-crf", str(crf),
            "-c:a", "aac",
            "-b:a", "128k",
        ]
        if faststart:
            cmd += ["-movflags", "+faststart"]
        cmd += [out_path]

        subprocess.run(cmd, check=True, capture_output=True)
        with open(out_path, "rb") as f:
            return (f"{height}p", f.read(), "video/mp4")