image_resizer/image_ops.py

"""Image resize, format conversion, and save operations."""

from __future__ import annotations

import os

from PIL import Image, ImageFile, ImageSequence

from image_resizer.config import (
  SAVE_QUALITY,
  WHITE_BG_BORDER_PX,
  WHITE_BG_MIN_ALPHA,
  WHITE_BG_MIN_RATIO,
  WHITE_BG_THRESHOLD,
)

ImageFile.LOAD_TRUNCATED_IMAGES = True


def resize_with_padding(
  img: Image.Image,
  target_w: int,
  target_h: int,
  fill: tuple[int, int, int] = (255, 255, 255),
) -> Image.Image:
  """Resize image to fit within target dimensions with letterboxing."""
  img.thumbnail((target_w, target_h), Image.LANCZOS)
  bg = Image.new("RGB", (target_w, target_h), fill)
  x = (target_w - img.width) // 2
  y = (target_h - img.height) // 2
  bg.paste(img, (x, y))
  return bg


def pad_without_resize(
  img: Image.Image,
  target_w: int,
  target_h: int,
  fill: tuple[int, int, int] = (255, 255, 255),
) -> Image.Image:
  """Expand canvas to at least target size; never scale or stretch the image."""
  canvas_w = max(img.width, target_w)
  canvas_h = max(img.height, target_h)
  if canvas_w == img.width and canvas_h == img.height:
    return img

  mode = "RGBA" if "A" in img.getbands() else "RGB"
  fill_color: tuple = (*fill, 255) if mode == "RGBA" else fill
  bg = Image.new(mode, (canvas_w, canvas_h), fill_color)
  x = (canvas_w - img.width) // 2
  y = (canvas_h - img.height) // 2
  if img.mode != mode:
    img = img.convert(mode)
  bg.paste(img, (x, y), img if mode == "RGBA" and "A" in img.getbands() else None)
  return bg


def _format_and_extension(fmt: str) -> tuple[str, str]:
  fmt_upper = fmt.upper()
  save_fmt = "JPEG" if fmt_upper in ("JPEG", "JPG") else fmt_upper
  ext = "jpg" if fmt_upper in ("JPEG", "JPG") else fmt.lower()
  return save_fmt, ext


def has_white_background(img: Image.Image) -> bool:
  """Return True if the border is mostly white or transparent (no background)."""
  img = img.convert("RGBA")
  width, height = img.size
  if width == 0 or height == 0:
    return False

  border = min(WHITE_BG_BORDER_PX, width // 2, height // 2)
  pixels: list[tuple[int, int, int, int]] = []

  if border == 0:
    pixels = list(img.getdata())
  else:
    for y in range(height):
      for x in range(width):
        if x < border or x >= width - border or y < border or y >= height - border:
          pixels.append(img.getpixel((x, y)))

  if not pixels:
    return False

  ok_count = 0
  for r, g, b, a in pixels:
    # Transparent / cutout ("no background") counts as OK for MAIN checks.
    if a < WHITE_BG_MIN_ALPHA:
      ok_count += 1
      continue
    if (
      r >= WHITE_BG_THRESHOLD
      and g >= WHITE_BG_THRESHOLD
      and b >= WHITE_BG_THRESHOLD
    ):
      ok_count += 1

  return ok_count / len(pixels) >= WHITE_BG_MIN_RATIO


def open_image_frames(path: str) -> list[Image.Image]:
  """Open an image file and return a list of RGBA frames (one per GIF frame)."""
  img = Image.open(path)
  if getattr(img, "is_animated", False):
    img.seek(0)
    return [frame.convert("RGBA").copy() for frame in ImageSequence.Iterator(img)]
  return [img.convert("RGBA")]


def save_processed_image(
  img: Image.Image,
  name: str,
  fmt: str,
  out_dir: str,
  width: int = None,
  height: int = None,
  frame_index: int = None,
  total_frames: int = 1,
  pad_only: bool = False,
) -> str:
  """Composite, optionally resize or pad, format-convert, and save a frame."""
  img = img.convert("RGBA")
  bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
  img = Image.alpha_composite(bg, img)

  save_fmt, ext = _format_and_extension(fmt)
  if save_fmt == "JPEG":
    img = img.convert("RGB")

  if width is not None and height is not None:
    if pad_only:
      img = pad_without_resize(img, width, height)
    else:
      img = resize_with_padding(img, width, height)

  if total_frames > 1 and frame_index is not None:
    fname = f"{name}_frame{frame_index + 1}.{ext}"
  else:
    fname = f"{name}.{ext}"

  out_path = os.path.join(out_dir, fname)
  save_kwargs = {"format": save_fmt}
  if save_fmt == "JPEG":
    save_kwargs.update(quality=SAVE_QUALITY, subsampling=0, optimize=True)
  elif save_fmt in ("WEBP", "AVIF"):
    save_kwargs["quality"] = SAVE_QUALITY

  img.save(out_path, **save_kwargs)
  return out_path


# Backward-compatible alias
def save_resized_image(
  img: Image.Image,
  name: str,
  fmt: str,
  width: int,
  height: int,
  out_dir: str,
  frame_index: int = None,
  total_frames: int = 1,
) -> str:
  return save_processed_image(
    img, name, fmt, out_dir,
    width=width, height=height,
    frame_index=frame_index, total_frames=total_frames,
  )