generate_video_covers.py

import json
import numpy as np
import os
import shutil
import subprocess
import sys

from PIL import Image, ImageDraw, ImageFont
from mtfont import Font, SegoeFontName
from mtfs import read_file
from mtlogger import logger
from mtprompt import Prompt
from mtvalidation import to_path
from mutagen.asf import ASF, ASFByteArrayAttribute
from mutagen.mp4 import MP4, MP4Cover
from tqdm import tqdm

from _common import generate_tmp_dir, get_ext
from _constants import MKV_EXT, MP4_EXTS, VIDEO_EXTS, WMV_EXT

SCREENSHOT_SECS = 3 * 60

FONT_PATH = Font.find_by_name(SegoeFontName.BOLD)
FONT_SIZE = 24

OVERLAY_PADDING = 4
OVERLAY_LINE_GAP = 8
OVERLAY_BG_ALPHA = 192


def main():
  if len(sys.argv) > 1:
    input_path = to_path(sys.argv[1])
  else:
    input_path = Prompt.path('Enter the path to a video file or directory')

  logger.log(f'Generating video covers for "{input_path}"...')
  logger.hr()

  try:
    if os.path.isfile(input_path):
      tmp_dir = generate_tmp_dir(os.path.dirname(input_path))
      process_file(input_path, tmp_dir)
    else:
      tmp_dir = generate_tmp_dir(input_path)
      process_directory(input_path, tmp_dir)

  finally:
    shutil.rmtree(tmp_dir, ignore_errors=True)

  logger.hr()
  logger.log(f'Finished generating video covers for "{input_path}".')


def process_directory(
  dir_path: str,
  tmp_dir: str,
):
  video_files = []

  for root, dirs, file_names in os.walk(dir_path):
    if os.path.abspath(root) == os.path.abspath(tmp_dir):
      dirs[:] = []
      continue

    for file_name in file_names:
      if get_ext(file_name) in VIDEO_EXTS:
        video_files.append(os.path.join(root, file_name))

  if not video_files:
    tqdm.write('No video files found.')
    return

  for file_path in tqdm(video_files, unit='file'):
    process_file(file_path, tmp_dir)


def process_file(
  file_path: str,
  tmp_dir: str,
):
  name = os.path.basename(file_path)
  stem = os.path.splitext(name)[0]
  ext = get_ext(file_path)

  probe = get_video_info(file_path)
  if probe is None:
    tqdm.write(f'Could not read video info for: {name}')
    return

  duration_secs = float(probe['format'].get('duration', 0))
  file_size = int(probe['format'].get('size', os.path.getsize(file_path)))
  attachment_count = sum(
    stream.get('codec_type') == 'attachment'
    for stream in probe.get('streams', [])
  )

  seek_secs = min(SCREENSHOT_SECS, duration_secs * 0.25)

  frame_path = os.path.join(tmp_dir, stem + '.jpg')
  out_path = os.path.join(tmp_dir, stem + ext)

  try:
    if has_existing_cover(file_path, ext, probe):
      tqdm.write(f'  Skipping existing cover: {name}')
      return

    capture_frame(file_path, seek_secs, frame_path)
    img = Image.open(frame_path).convert('RGB')
    img = crop_black_borders(img)
    img = crop_to_16_9(img)
    img.thumbnail((256, 256), Image.LANCZOS)
    img = draw_stats_overlay(img, duration_secs, file_size)
    img.save(frame_path, 'JPEG', quality=90)

    if ext in MP4_EXTS:
      embed_mp4_cover(file_path, frame_path)
    elif ext == WMV_EXT:
      embed_asf_cover(file_path, frame_path)
    else:
      embed_cover_ffmpeg(file_path, frame_path, out_path, ext, attachment_count)
      os.replace(out_path, file_path)

    tqdm.write(f'  Embedded cover into: {name}')
  except Exception as ex:
    tqdm.write(f'  Failed: {ex}')


def embed_mp4_cover(
  file_path: str,
  cover_path: str,
):
  data = read_file(cover_path)
  tags = MP4(file_path)
  tags['covr'] = [MP4Cover(data, MP4Cover.FORMAT_JPEG)]
  tags.save()


def embed_asf_cover(
  file_path: str,
  cover_path: str,
):
  data = read_file(cover_path)
  tags = ASF(file_path)
  tags['WM/Picture'] = [ASFByteArrayAttribute(data)]
  tags.save()


def has_existing_cover(
  file_path: str,
  ext: str,
  probe: dict,
) -> bool:
  if ext in MP4_EXTS:
    media = MP4(file_path)
    return media.tags is not None and 'covr' in media.tags

  if ext == WMV_EXT:
    return 'WM/Picture' in ASF(file_path)

  for stream in probe.get('streams', []):
    stream_type = stream.get('codec_type')
    filename = stream.get('tags', {}).get('filename', '').casefold()
    is_attached_picture = stream_type == 'video' and stream.get(
      'disposition', {}
    ).get('attached_pic')
    is_cover_attachment = (
      stream_type == 'attachment' and filename == 'cover.jpg'
    )
    is_cover_video = stream_type == 'video' and filename == 'cover.jpg'
    if is_attached_picture or is_cover_attachment or is_cover_video:
      return True

  return False


def embed_cover_ffmpeg(
  file_path: str,
  cover_path: str,
  output_path: str,
  ext: str,
  attachment_count: int,
):
  if ext == MKV_EXT:
    cmd = [
      'ffmpeg',
      '-y',
      '-loglevel',
      'error',
      '-i',
      file_path,
      '-map',
      '0',
      '-attach',
      cover_path,
      f'-metadata:s:t:{attachment_count}',
      'mimetype=image/jpeg',
      f'-metadata:s:t:{attachment_count}',
      'filename=cover.jpg',
      '-c',
      'copy',
      output_path,
    ]
  else:
    cmd = [
      'ffmpeg',
      '-y',
      '-loglevel',
      'error',
      '-i',
      file_path,
      '-i',
      cover_path,
      '-map',
      '0',
      '-map',
      '1',
      '-c',
      'copy',
      '-disposition:v:1',
      'attached_pic',
      output_path,
    ]

  result = subprocess.run(cmd, check=False, capture_output=True)
  if result.returncode != 0:
    raise RuntimeError(result.stderr.decode(errors='replace'))


def get_video_info(
  file_path: str,
):
  try:
    result = subprocess.run(
      [
        'ffprobe',
        '-v',
        'quiet',
        '-analyzeduration',
        '0',
        '-probesize',
        '5000000',
        '-print_format',
        'json',
        '-show_format',
        '-show_streams',
        file_path,
      ],
      check=False,
      capture_output=True,
      encoding='utf-8',
      errors='replace',
      text=True,
    )
    return json.loads(result.stdout)
  except Exception:
    return None


def crop_to_16_9(
  img: Image.Image,
):
  w, h = img.size
  if w / h > 16 / 9:
    new_w = int(h * 16 / 9)
    x = (w - new_w) // 2
    img = img.crop((x, 0, x + new_w, h))
  return img


def crop_black_borders(
  img: Image.Image,
  threshold=10,
):
  arr = np.array(img)
  mask = arr.max(axis=2) > threshold
  rows = np.any(mask, axis=1)
  cols = np.any(mask, axis=0)
  if not rows.any() or not cols.any():
    return img
  y0, y1 = np.where(rows)[0][[0, -1]]
  x0, x1 = np.where(cols)[0][[0, -1]]
  return img.crop((x0, y0, x1 + 1, y1 + 1))


def capture_frame(
  file_path: str,
  seek_secs: float,
  output_path: str,
):
  subprocess.run(
    [
      'ffmpeg',
      '-y',
      '-loglevel',
      'error',
      '-analyzeduration',
      '0',
      '-probesize',
      '5000000',
      '-ss',
      str(seek_secs),
      '-i',
      file_path,
      '-vframes',
      '1',
      '-q:v',
      '2',
      output_path,
    ],
    check=False,
    stdout=subprocess.DEVNULL,
    stderr=subprocess.DEVNULL,
  )
  if not os.path.exists(output_path):
    raise RuntimeError(
      f'Failed to extract frame from: {os.path.basename(file_path)}'
    )


def format_duration(
  secs: float,
):
  secs = int(secs)
  h = secs // 3600
  m = (secs % 3600) // 60
  s = secs % 60
  if h > 0:
    return f'{h}:{m:02d}:{s:02d}'
  return f'{m}:{s:02d}'


def format_size(
  size_bytes: int,
):
  if size_bytes >= 1024**3:
    return f'{round(size_bytes / (1024**3), 2)} GB'
  return f'{round(size_bytes / (1024**2))} MB'


def draw_stats_overlay(
  img: Image.Image,
  duration_secs: float,
  file_size: int,
):
  draw = ImageDraw.Draw(img, 'RGBA')
  font = Font.load_by_path(FONT_PATH, FONT_SIZE)
  img_w, _ = img.size

  draw_label(
    draw, font, format_duration(duration_secs), side='left', img_w=img_w
  )
  draw_label(draw, font, format_size(file_size), side='right', img_w=img_w)

  return img.convert('RGB')


def draw_label(
  draw: ImageDraw.ImageDraw,
  font: ImageFont.FreeTypeFont,
  text: str,
  side: str,
  img_w: int,
):
  bb = draw.textbbox((0, 0), text, font=font)
  tw = bb[2] - bb[0]
  th = bb[3] - bb[1]
  top = bb[1]

  box_w = tw + OVERLAY_PADDING * 2
  box_h = th + OVERLAY_PADDING * 2

  box_x = 0 if side == 'left' else img_w - box_w
  box_x2 = box_w if side == 'left' else img_w

  draw.rectangle([box_x, 0, box_x2, box_h], fill=(0, 0, 0, OVERLAY_BG_ALPHA))

  tx = box_x + OVERLAY_PADDING
  draw.text(
    (tx, OVERLAY_PADDING - top), text, font=font, fill=(255, 255, 255, 255)
  )


if __name__ == '__main__':
  try:
    main()
  except Exception as ex:
    logger.unhandled_error(ex)

  Prompt.enter_to_exit(timeout=True)