espresso_frame/server/app/image_pipeline.py

"""Resize, quantize, and pack a photo into the panel's raw 4bpp format."""

from __future__ import annotations

from PIL import Image, ImageOps

EPD_WIDTH = 800
EPD_HEIGHT = 480

# Approximate sRGB for each of the panel's 6 ink colors. These are
# reasonable placeholders, not measured values -- Waveshare doesn't publish
# exact color primaries for this panel. Tune them once you can compare a
# rendered test image against the real panel.
PALETTE_RGB = [
    (0, 0, 0),  # BLACK
    (255, 255, 255),  # WHITE
    (255, 219, 0),  # YELLOW
    (207, 0, 15),  # RED
    (0, 39, 133),  # BLUE
    (0, 133, 55),  # GREEN
]

# The panel's actual 4-bit color codes (see firmware/components/epd7in3e),
# in the same order as PALETTE_RGB. 0x4 is intentionally unused upstream.
PANEL_CODES = [0x0, 0x1, 0x2, 0x3, 0x5, 0x6]


def _build_palette_image() -> Image.Image:
    pal_img = Image.new("P", (1, 1))
    pal_img.putpalette([channel for rgb in PALETTE_RGB for channel in rgb])
    return pal_img


_PALETTE_IMAGE = _build_palette_image()


def _plain_center_crop_box(
    img_width: int, img_height: int, target_width: int, target_height: int
) -> tuple[float, float, int, int]:
    """The largest target_width:target_height window centered in the
    source image -- the same box ImageOps.fit() computes internally when
    there's no face-aware shift to apply. Returns (left, top, crop_w,
    crop_h); left/top are floats (not yet rounded) since callers that go
    on to face-shift this box need the unrounded center point."""
    target_ratio = target_width / target_height
    if img_width / img_height > target_ratio:
        crop_h = img_height
        crop_w = int(crop_h * target_ratio)
    else:
        crop_w = img_width
        crop_h = int(crop_w / target_ratio)

    left = (img_width - crop_w) / 2
    top = (img_height - crop_h) / 2
    return left, top, crop_w, crop_h


def _face_aware_crop_box(
    img_width: int, img_height: int, target_width: int, target_height: int, faces: list[dict]
) -> tuple[int, int, int, int]:
    """Largest crop window matching target_width:target_height that fits
    inside the source image. Starts from the plain center crop and only
    shifts it the minimum amount needed to bring any faces that would
    otherwise be cut off back on screen -- an already-fine composition
    (faces already fully inside the center crop) is left untouched rather
    than re-centered on the faces. If the faces themselves span wider than
    the crop window allows, centers on their midpoint as best-effort,
    since there's no shift that fits them all regardless.

    Each face's box is given relative to its own imageWidth/imageHeight
    (the resolution Immich ran detection on), which may differ from the
    downloaded preview's resolution passed in here, so each box is scaled
    into img_width/img_height space before use.
    """
    min_x = min_y = float("inf")
    max_x = max_y = float("-inf")
    for face in faces:
        face_w = face.get("imageWidth") or img_width
        face_h = face.get("imageHeight") or img_height
        scale_x = img_width / face_w
        scale_y = img_height / face_h
        min_x = min(min_x, face["boundingBoxX1"] * scale_x)
        max_x = max(max_x, face["boundingBoxX2"] * scale_x)
        min_y = min(min_y, face["boundingBoxY1"] * scale_y)
        max_y = max(max_y, face["boundingBoxY2"] * scale_y)

    left, top, crop_w, crop_h = _plain_center_crop_box(img_width, img_height, target_width, target_height)

    if max_x - min_x <= crop_w:
        if min_x < left:
            left = min_x
        elif max_x > left + crop_w:
            left = max_x - crop_w
    else:
        left = (min_x + max_x) / 2 - crop_w / 2

    if max_y - min_y <= crop_h:
        if min_y < top:
            top = min_y
        elif max_y > top + crop_h:
            top = max_y - crop_h
    else:
        top = (min_y + max_y) / 2 - crop_h / 2

    left = max(0, min(left, img_width - crop_w))
    top = max(0, min(top, img_height - crop_h))

    return (int(left), int(top), int(left) + crop_w, int(top) + crop_h)


def render_frame(source: Image.Image, faces: list[dict] | None = None) -> bytes:
    """Fits `source` to the panel's resolution, quantizes it to the 6-color
    palette with Floyd-Steinberg dithering, and packs 2 pixels/byte the way
    epd7in3e.c expects. Always returns exactly EPD_WIDTH*EPD_HEIGHT/2 bytes.

    If `faces` (from ImmichClient.get_asset_faces) is non-empty, crops
    toward keeping them on screen instead of a plain center-crop.
    """
    fitted = ImageOps.exif_transpose(source.convert("RGB"))

    if faces:
        box = _face_aware_crop_box(fitted.width, fitted.height, EPD_WIDTH, EPD_HEIGHT, faces)
        fitted = fitted.crop(box).resize((EPD_WIDTH, EPD_HEIGHT), Image.LANCZOS)
    else:
        fitted = ImageOps.fit(fitted, (EPD_WIDTH, EPD_HEIGHT), method=Image.LANCZOS)

    quantized = fitted.quantize(palette=_PALETTE_IMAGE, dither=Image.Dither.FLOYDSTEINBERG)
    pixels = quantized.load()

    out = bytearray(EPD_WIDTH * EPD_HEIGHT // 2)
    i = 0
    for y in range(EPD_HEIGHT):
        for x in range(0, EPD_WIDTH, 2):
            left = PANEL_CODES[pixels[x, y]]
            right = PANEL_CODES[pixels[x + 1, y]]
            out[i] = (left << 4) | right
            i += 1

    return bytes(out)