from pathlib import Path from typing import Optional import hashlib, base64, binascii, math, subprocess, shutil, os, logging, uuid, tempfile from dotenv import load_dotenv from imageio_ffmpeg import get_ffmpeg_exe import boto3 load_dotenv() logger = logging.getLogger(__name__) CANONICAL_STATIC_BASE = "https://adgen.static.lookfinity.com" # ---------- COLOR FILTERS ---------- SAFE_EFFECTS = { "warm": "eq=brightness=0.01:contrast=1.02:saturation=1.08,colorbalance=rs=0.04:bs=-0.03", "fresh": "eq=brightness=0.015:contrast=1.03:saturation=1.12,colorbalance=gs=0.03:bs=0.02", "bright_clean": "eq=brightness=0.02:contrast=1.03:saturation=1.04", "natural": "eq=brightness=0.006:contrast=1.015:saturation=1.025", "product_clean": "eq=brightness=0.018:contrast=1.04:saturation=1.07,unsharp=3:3:0.2", "light_sharpen": "eq=brightness=0.005:contrast=1.02:saturation=1.04,unsharp=3:3:0.18", "soft": "eq=brightness=0.01:contrast=0.98:saturation=1.04", "muted": "eq=brightness=0.005:contrast=0.97:saturation=0.92", "cinematic": "eq=brightness=-0.005:contrast=1.06:saturation=0.98,colorbalance=rs=0.015:bs=0.025", "subtle": "eq=brightness=0.004:contrast=1.01:saturation=1.015", } def _public_base_url() -> str: base = (os.getenv("NEW_BASE") or "").strip().rstrip("/") return base or CANONICAL_STATIC_BASE # ---------- HASH UTILS ---------- def _sha256(path: Path, chunk: int = 1 << 20) -> str: h = hashlib.sha256() with path.open("rb") as f: for b in iter(lambda: f.read(chunk), b""): h.update(b) return h.hexdigest() def _hashid_short(sha256_hex: str, length: int = 12) -> str: b = binascii.unhexlify(sha256_hex) return base64.urlsafe_b64encode(b).decode().rstrip("=")[:length] # ---------- FFMPEG ---------- def _ffmpeg_bin() -> str: try: return get_ffmpeg_exe() except Exception: p = os.environ.get("FFMPEG_BIN") or shutil.which("ffmpeg") if not p: raise FileNotFoundError("ffmpeg not found. Install it or set FFMPEG_BIN") return p FFMPEG = _ffmpeg_bin() def _run_ffmpeg(cmd: list[str]) -> None: cmd = cmd[:] cmd[0] = FFMPEG proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) if proc.returncode != 0: raise RuntimeError(proc.stderr.decode(errors="ignore") or "ffmpeg failed") # ---------- THUMBNAIL ---------- def extract_thumbnail(video_path: str, time_position: str = "00:00:01") -> str: """ Extracts a thumbnail (JPEG) from the video and returns it as a base64 string. Default frame at 1 second. """ tmp_thumb = Path(tempfile.gettempdir()) / f"{Path(video_path).stem}_thumb.jpg" cmd = [ "ffmpeg", "-nostdin", "-y", "-ss", time_position, "-i", video_path, "-frames:v", "1", "-q:v", "2", str(tmp_thumb), ] _run_ffmpeg(cmd) with open(tmp_thumb, "rb") as f: encoded = base64.b64encode(f.read()).decode("utf-8") return encoded # ---------- AUGMENT VIDEO ---------- def augment_video_random( *, input_path: str, output_path: Optional[str] = None, crf: int = 20, preset: str = "medium", variation_idx: int = 0, border_width: int = 0, border_color: str = "black", effect_name: Optional[str] = None, ) -> str: if effect_name and effect_name in SAFE_EFFECTS: vf = SAFE_EFFECTS[effect_name] else: import random k_b, k_c, k_h, k_s = [random.randint(-3, 3) for _ in range(4)] brightness = max(-1.0, min(1.0, k_b * 0.2)) contrast = max(0.0, min(2.0, 1.0 + k_c * 0.03)) hue_rad = math.radians(k_h * 3.0) sat_scale = max(0.0, 1.0 + k_s * 0.05) vf = f"hue=h={hue_rad:.6f}:s={sat_scale:.4f},eq=contrast={contrast:.4f}:brightness={brightness:.4f}" if border_width > 0: ffmpeg_color = border_color.replace("#", "0x") if border_color.startswith("#") else border_color bw = int(border_width) vf += f",pad=iw+{bw*2}:ih+{bw*2}:{bw}:{bw}:{ffmpeg_color}" inp = Path(input_path) suffix = f"_augmented_{variation_idx}" if variation_idx > 0 else "_augmented" out = Path(output_path) / f"{inp.stem}{suffix}.mp4" if output_path else inp.with_name(f"{inp.stem}{suffix}.mp4") cmd = [ "ffmpeg", "-nostdin", "-y", "-i", str(inp), "-map", "0:v:0", "-map", "0:a?", "-vf", vf, "-c:v", "libx264", "-preset", preset, "-crf", str(crf), "-pix_fmt", "yuv420p", "-c:a", "copy", "-movflags", "+faststart", str(out), ] _run_ffmpeg(cmd) return str(out) # ---------- VIDEO UPLOAD TO R2 ---------- def upload_video_to_r2(video_bytes: bytes, file_name: str) -> str: s3 = boto3.client( "s3", endpoint_url=os.getenv("R2_ENDPOINT"), aws_access_key_id=os.getenv("R2_ACCESS_KEY"), aws_secret_access_key=os.getenv("R2_SECRET_KEY"), region_name="auto", ) bucket = os.getenv("R2_BUCKET_NAME") key = f"videos/{uuid.uuid4()}_{file_name}" s3.put_object(Bucket=bucket, Key=key, Body=video_bytes, ContentType="video/mp4") return f"{_public_base_url()}/{key}" # ---------- PROCESS + UPLOAD ---------- def process_video_with_hash_info( input_path: str, output_path: Optional[str] = None, variation_idx: int = 0, border_width: int = 0, border_color: str = "black", effect_name: Optional[str] = None, ) -> dict: logger.info(f"[process] Starting process_video_with_hash_info for {Path(input_path).name}") logger.info(f"[process] Step 1/5 — computing input hash") in_sha = _sha256(Path(input_path)) in_id = _hashid_short(in_sha) logger.info(f"[process] Step 2/5 — augmenting video (effect={effect_name or 'random'})") out_path = augment_video_random( input_path=input_path, output_path=output_path, variation_idx=variation_idx, border_width=border_width, border_color=border_color, effect_name=effect_name, ) logger.info(f"[process] Step 3/5 — computing output hash") out_sha = _sha256(Path(out_path)) out_id = _hashid_short(out_sha) logger.info(f"[process] Step 4/5 — extracting thumbnail") safe_tmp = Path(tempfile.gettempdir()) / Path(out_path).name shutil.copy(out_path, safe_tmp) # thumbnail try: thumbnail_b64 = extract_thumbnail(str(safe_tmp)) except Exception as e: logger.error(f"Thumbnail extraction failed: {e}") thumbnail_b64 = "" logger.info(f"[process] Step 5/5 — uploading to R2") try: with open(safe_tmp, "rb") as f: video_bytes = f.read() r2_url = upload_video_to_r2(video_bytes, Path(safe_tmp).name) except Exception as e: logger.error(f"Upload to R2 failed: {e}") r2_url = None logger.info(f"[process] Done. output_hashid={out_id}, r2_url={r2_url}") return { "input_hashid": in_id, "output_name": Path(safe_tmp).name, "output_path": str(safe_tmp), "output_hashid": out_id, "output_r2_url": r2_url, "thumbnail": thumbnail_b64, "effect": effect_name or "random", }