Spaces:
Running on Zero
Running on Zero
| """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]) | |