| """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]) |
|
|
