screenshot-remove-black-edge.py

#!/usr/bin/env -S uv --quiet run --script
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "numpy>=1.24",
#     "pillow>=9.0"
# ]
# ///
"""
Five-Iteration Bezel Trimmer
============================
Removes dark top/bottom bezels from mobile screenshots and flattens transparency to white.
Now includes --inplace-tilde to rename originals to filename~ before overwriting.
"""

import argparse
import sys
from pathlib import Path
import shutil
import numpy as np
from PIL import Image, ImageDraw, ImageOps


def load_with_exif(path: Path) -> Image.Image:
    im = Image.open(path)
    try:
        im = ImageOps.exif_transpose(im)
    except Exception:
        pass
    return im


def flatten_to_white(im: Image.Image) -> Image.Image:
    if im.mode in ("RGBA", "LA") or (im.mode == "P" and "transparency" in im.info):
        bg = Image.new("RGBA", im.size, (255, 255, 255, 255))
        bg.alpha_composite(im.convert("RGBA"))
        return bg.convert("RGB")
    return im.convert("RGB")


def compute_luminance(arr: np.ndarray) -> np.ndarray:
    r = arr[..., 0].astype(np.float32)
    g = arr[..., 1].astype(np.float32)
    b = arr[..., 2].astype(np.float32)
    return 0.2126 * r + 0.7152 * g + 0.0722 * b


def moving_max(v: np.ndarray, win: int) -> np.ndarray:
    if win <= 1 or v.size < win:
        return v
    pad = win - 1
    padded = np.pad(v, (pad // 2, pad - pad // 2), mode="edge")
    out = np.empty_like(v)
    for i in range(v.shape[0]):
        out[i] = np.max(padded[i : i + win])
    return out


def adaptive_threshold(y: np.ndarray, lo_pct=5.0, hi_pct=25.0) -> float:
    lo = np.percentile(y, lo_pct)
    hi = np.percentile(y, hi_pct)
    return float(lo * 0.8 + hi * 0.2)


def find_crop_rows(
    im: Image.Image,
    dark_thresh: int = 60,
    dark_frac: float = 0.9,
    pad: int = 4,
    auto_thresh: bool = False,
    top_only: bool = False,
    bottom_only: bool = False,
    min_keep: int = 400,
    max_chop: int = None,
    brightness_span_min: float = 10.0,
):
    arr = np.asarray(im.convert("RGB"))
    y = compute_luminance(arr)
    h, _ = y.shape

    if auto_thresh:
        est = adaptive_threshold(y)
        dark_thresh = int(max(10, min(100, est)))

    dark = y < dark_thresh
    row_dark_fraction = moving_max(dark.mean(axis=1), 5)

    top, bottom = 0, h - 1
    while top < h and row_dark_fraction[top] >= dark_frac:
        top += 1
    while bottom >= 0 and row_dark_fraction[bottom] >= dark_frac:
        bottom -= 1

    if top_only:
        bottom = h - 1
    if bottom_only:
        top = 0
    if max_chop is not None:
        top = min(top, max_chop)
        bottom = max(bottom, h - 1 - max_chop)

    if bottom <= top or (bottom - top + 1) < min_keep:
        return 0, h

    kept = y[top : bottom + 1, :]
    if kept.size == 0 or (kept.max() - kept.min()) < brightness_span_min:
        return 0, h

    top = max(0, top - pad if not bottom_only else 0)
    bottom = min(h - 1, bottom + pad if not top_only else h - 1)

    return top, bottom + 1


def draw_debug_overlay(im: Image.Image, top: int, bottom: int) -> Image.Image:
    out = im.convert("RGB").copy()
    dr = ImageDraw.Draw(out)
    w, _ = out.size
    for y in (top, bottom - 1):
        dr.line([(0, y), (w, y)], fill=(255, 0, 0), width=3)
    return out


def backup_original(src: Path):
    backup_path = src.with_name(src.name + "~")
    shutil.copy2(src, backup_path)
    return backup_path


def process_file(src: Path, dst: Path, args):
    im = load_with_exif(src)
    im = flatten_to_white(im)

    if args.fixed_chop > 0:
        h = im.height
        t = min(args.fixed_chop, h // 2 - 1)
        b = h - t
        if not args.dry_run:
            im.crop((0, t, im.width, b)).save(dst, quality=95)
        return True, f"fixed-chop top={t}, bottom={b}"

    top, bottom = find_crop_rows(
        im,
        args.dark_thresh,
        args.dark_frac,
        args.pad,
        args.auto_thresh,
        args.top_only,
        args.bottom_only,
        args.min_keep,
        args.max_chop,
        args.brightness_span_min,
    )

    if args.debug_overlay:
        draw_debug_overlay(im, top, bottom).save(
            dst.with_name(dst.stem + "_debug" + dst.suffix)
        )
    if not args.dry_run:
        im.crop((0, top, im.width, bottom)).save(dst, quality=95)
    return True, f"crop rows [{top}:{bottom}) height_out={bottom - top}"


def main():
    ap = argparse.ArgumentParser(
        description="Trim dark top/bottom bezels and flatten transparency to white."
    )
    ap.add_argument("inputs", nargs="+", help="Input files or directories.")
    ap.add_argument("-o", "--outdir", default="trimmed", help="Output directory.")
    ap.add_argument("--suffix", default="_trim", help="Suffix for output filenames.")
    ap.add_argument("--inplace", action="store_true", help="Overwrite originals.")
    ap.add_argument(
        "--inplace-tilde",
        action="store_true",
        help="Rename original to filename~ before overwriting.",
    )
    ap.add_argument(
        "--dark-thresh",
        type=int,
        default=60,
        help="Luminance threshold for 'dark' (0-255).",
    )
    ap.add_argument(
        "--dark-frac",
        type=float,
        default=0.9,
        help="Row is bezel if >= this fraction is dark.",
    )
    ap.add_argument(
        "--pad", type=int, default=4, help="Safety pixels kept inside crop."
    )
    ap.add_argument(
        "--auto-thresh", action="store_true", help="Auto-estimate dark threshold."
    )
    ap.add_argument("--top-only", action="store_true", help="Only trim top bezel.")
    ap.add_argument(
        "--bottom-only", action="store_true", help="Only trim bottom bezel."
    )
    ap.add_argument(
        "--fixed-chop",
        type=int,
        default=0,
        help="Remove exactly N pixels from top/bottom.",
    )
    ap.add_argument("--min-keep", type=int, default=400, help="Minimum height to keep.")
    ap.add_argument(
        "--max-chop",
        type=int,
        default=None,
        help="Max pixels to remove from each side.",
    )
    ap.add_argument(
        "--brightness-span-min",
        type=float,
        default=10.0,
        help="Min luminance span required.",
    )
    ap.add_argument(
        "--debug-overlay", action="store_true", help="Save debug image with crop lines."
    )
    ap.add_argument(
        "--dry-run", action="store_true", help="Compute crop only, don't write."
    )
    args = ap.parse_args()

    files = []
    for item in args.inputs:
        p = Path(item)
        if p.is_dir():
            for ext in ("*.png", "*.jpg", "*.jpeg", "*.webp"):
                files.extend(p.glob(ext))
        elif p.exists():
            files.append(p)

    if not files:
        print("No input files.", file=sys.stderr)
        sys.exit(2)

    outdir = Path(args.outdir)
    if not args.inplace and not args.inplace_tilde:
        outdir.mkdir(parents=True, exist_ok=True)

    for src in files:
        if args.inplace_tilde:
            backup_original(src)
            dst = src
        elif args.inplace:
            dst = src
        else:
            dst = outdir / (src.stem + args.suffix + src.suffix)

        ok, msg = process_file(src, dst, args)
        print(f"{'OK' if ok else 'ERR'} {src.name} -> {dst.name}: {msg}")


if __name__ == "__main__":
    main()