"""Draw the red outline the LoRA expects around the object to remove.

    python draw_outline.py photo.jpg mask.png photo_outlined.png

The mask is white (> 127) on the object, at any size (it is resized to the photo with nearest
neighbour). The outline is a pure red ring (255, 0, 0) starting 2 px outside the mask, 3 px wide,
or 4 px when the photo's short side is over 900 px: the style the model expects. Save the
result losslessly (png) if you keep it, so the red stays pure.

Needs numpy, opencv-python(-headless) and pillow.
"""

from __future__ import annotations

import sys

import cv2
import numpy as np
from PIL import Image

RED = (255, 0, 0)
OUTLINE_SUFFIX = "The area to edit is outlined in red. Remove the red outline as well."


def _within(mask: np.ndarray, r: float) -> np.ndarray:
    """Pixels within distance r of the mask (Euclidean distance transform)."""
    dist = cv2.distanceTransform((~mask).astype(np.uint8), cv2.DIST_L2, 5)
    return dist <= r


def outline_ring(mask: np.ndarray, width: int | None = None, gap: int = 2) -> np.ndarray:
    """Boolean ring from `gap` to `gap + width` px outside the boolean mask."""
    h, w = mask.shape
    width = width or (3 if min(h, w) <= 900 else 4)
    return _within(mask, gap + width) & ~_within(mask, gap)


def draw_outline(photo: Image.Image, mask: Image.Image) -> Image.Image:
    """The photo with the mask's outline drawn in pure red. An empty mask returns the photo."""
    photo = photo.convert("RGB")
    m = np.array(mask.convert("L").resize(photo.size, Image.NEAREST)) > 127
    out = np.array(photo)
    if m.any():
        out[outline_ring(m)] = RED
    return Image.fromarray(out)


def outline_prompt(instruction: str) -> str:
    """The instruction followed by the sentence the model expects with outlined images."""
    return f"{instruction.rstrip()} {OUTLINE_SUFFIX}"


if __name__ == "__main__":
    if len(sys.argv) != 4:
        sys.exit(__doc__)
    draw_outline(Image.open(sys.argv[1]), Image.open(sys.argv[2])).save(sys.argv[3])