generated from westfarn/web_django_template
westfarn
b2029b6c0f
Replaces the client_site template branding, adds shop/shipping, and points beta CI at master for easy deploy. Closes #1
122 lines
3.4 KiB
Python
122 lines
3.4 KiB
Python
"""Turn a raw product photo into a centered cutout on a square canvas."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from io import BytesIO
|
|
from pathlib import PurePosixPath
|
|
|
|
from PIL import Image, ImageOps
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
CANVAS_SIZE = 1200
|
|
CUTOUT_MAX_SIDE = 1600
|
|
PADDING = 0.1
|
|
ALPHA_THRESHOLD = 128
|
|
|
|
_session = None
|
|
|
|
|
|
def _rembg_session():
|
|
global _session
|
|
if _session is None:
|
|
from rembg import new_session
|
|
|
|
_session = new_session("u2net")
|
|
return _session
|
|
|
|
|
|
def open_image(data: bytes) -> Image.Image:
|
|
image = Image.open(BytesIO(data))
|
|
image.load()
|
|
image = ImageOps.exif_transpose(image) or image
|
|
if getattr(image, "n_frames", 1) > 1:
|
|
image.seek(0)
|
|
image = image.copy()
|
|
return image
|
|
|
|
|
|
def downscale(image: Image.Image, max_side: int) -> Image.Image:
|
|
width, height = image.size
|
|
longest = max(width, height)
|
|
if longest <= max_side:
|
|
return image
|
|
scale = max_side / longest
|
|
return image.resize(
|
|
(max(1, int(width * scale)), max(1, int(height * scale))),
|
|
Image.Resampling.LANCZOS,
|
|
)
|
|
|
|
|
|
def harden_cutout(image: Image.Image, threshold: int = ALPHA_THRESHOLD) -> Image.Image:
|
|
"""Drop faint shadow/halo pixels rembg leaves around the subject."""
|
|
rgba = image.convert("RGBA")
|
|
red, green, blue, alpha = rgba.split()
|
|
alpha = alpha.point(lambda value: value if value >= threshold else 0)
|
|
return Image.merge("RGBA", (red, green, blue, alpha))
|
|
|
|
|
|
def cutout_subject(image: Image.Image) -> Image.Image:
|
|
"""Knock out the background. Falls back to the original on failure."""
|
|
rgba = downscale(image.convert("RGBA"), CUTOUT_MAX_SIDE)
|
|
try:
|
|
from rembg import remove
|
|
|
|
buf = BytesIO()
|
|
rgba.save(buf, format="PNG")
|
|
result = remove(
|
|
buf.getvalue(),
|
|
session=_rembg_session(),
|
|
post_process_mask=True,
|
|
)
|
|
return harden_cutout(Image.open(BytesIO(result)).convert("RGBA"))
|
|
except Exception:
|
|
logger.exception("product photo cutout failed; using original")
|
|
return rgba
|
|
|
|
|
|
def center_on_square(
|
|
image: Image.Image,
|
|
*,
|
|
size: int = CANVAS_SIZE,
|
|
padding: float = PADDING,
|
|
) -> Image.Image:
|
|
rgba = image.convert("RGBA")
|
|
alpha = rgba.getchannel("A")
|
|
mask = alpha.point(lambda value: 255 if value > ALPHA_THRESHOLD else 0)
|
|
bbox = mask.getbbox()
|
|
canvas = Image.new("RGBA", (size, size), (0, 0, 0, 0))
|
|
if bbox is None:
|
|
return canvas
|
|
cropped = rgba.crop(bbox)
|
|
max_inner = max(1, int(size * (1 - 2 * padding)))
|
|
width, height = cropped.size
|
|
scale = min(max_inner / width, max_inner / height)
|
|
new_w = max(1, int(width * scale))
|
|
new_h = max(1, int(height * scale))
|
|
resized = cropped.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
|
x = (size - new_w) // 2
|
|
y = (size - new_h) // 2
|
|
canvas.paste(resized, (x, y), resized)
|
|
return canvas
|
|
|
|
|
|
def encode_png(image: Image.Image) -> bytes:
|
|
buf = BytesIO()
|
|
image.save(buf, format="PNG", optimize=True)
|
|
return buf.getvalue()
|
|
|
|
|
|
def prepare_product_photo(data: bytes) -> tuple[bytes, str]:
|
|
"""Return a PNG cutout on a square canvas, plus content type."""
|
|
image = open_image(data)
|
|
cut = cutout_subject(image)
|
|
framed = center_on_square(cut)
|
|
return encode_png(framed), "image/png"
|
|
|
|
|
|
def product_image_filename(original: str) -> str:
|
|
stem = PurePosixPath(original or "product").stem.strip() or "product"
|
|
return f"{stem[:200]}.png"
|