Spaces:
Running on Zero
Running on Zero
| """Image loading, background removal and capture-quality checks.""" | |
| import hashlib | |
| import io | |
| from dataclasses import dataclass | |
| from functools import lru_cache | |
| import numpy as np | |
| from PIL import Image, ImageOps | |
| from returnguard import config | |
| def load_image(source) -> Image.Image: | |
| """Open a path, bytes or file-like object as an upright RGB image, downscaled.""" | |
| if isinstance(source, (bytes, bytearray)): | |
| source = io.BytesIO(source) | |
| img = ImageOps.exif_transpose(Image.open(source)).convert("RGB") | |
| img.thumbnail((config.MAX_IMAGE_SIDE, config.MAX_IMAGE_SIDE), Image.LANCZOS) | |
| return img | |
| def image_digest(img: Image.Image) -> str: | |
| return hashlib.sha1(img.tobytes() + str(img.size).encode()).hexdigest() | |
| def _rembg_session(): | |
| from rembg import new_session # heavy import, load once | |
| return new_session(config.BG_MODEL) | |
| def remove_background(img: Image.Image) -> Image.Image: | |
| """Return an RGBA cut-out of the main object, cropped to its bounding box. | |
| Results are cached on disk by image content, because evaluation re-processes | |
| the same photos many times.""" | |
| cache = config.CACHE_DIR / "cutouts" / config.BG_MODEL / f"{image_digest(img)}.png" | |
| if cache.exists(): | |
| return Image.open(cache).convert("RGBA") | |
| from rembg import remove | |
| cut = remove(img, session=_rembg_session()).convert("RGBA") | |
| alpha = np.array(cut.getchannel("A")) | |
| mask = alpha > 127 | |
| if mask.any(): | |
| ys, xs = np.where(mask) | |
| cut = cut.crop((xs.min(), ys.min(), xs.max() + 1, ys.max() + 1)) | |
| cache.parent.mkdir(parents=True, exist_ok=True) | |
| cut.save(cache) | |
| return cut | |
| class CaptureQuality: | |
| object_fraction: float # share of the original frame covered by the object | |
| sharpness: float # Laplacian variance inside the object (grayscale) | |
| ok: bool | |
| problem: str | None | |
| def _laplacian_var(gray: np.ndarray, mask: np.ndarray) -> float: | |
| lap = (np.roll(gray, -1, 0) + np.roll(gray, 1, 0) + np.roll(gray, -1, 1) + np.roll(gray, 1, 1) - 4 * gray) | |
| # erode the mask so the object outline (a hard edge against the background) doesn't count | |
| inner = mask & np.roll(mask, 3, 0) & np.roll(mask, -3, 0) & np.roll(mask, 3, 1) & np.roll(mask, -3, 1) | |
| return float(lap[inner].var()) if inner.sum() > 100 else 0.0 | |
| def assess_quality(original: Image.Image, cutout: Image.Image) -> CaptureQuality: | |
| """Is this photo usable as evidence? Measures focus on the object itself, | |
| not the silhouette edge, and checks the object is actually in frame.""" | |
| alpha = np.array(cutout.getchannel("A")) > 127 | |
| fraction = alpha.sum() / (original.width * original.height) | |
| gray = np.array(cutout.convert("L"), dtype=np.float32) | |
| sharp = _laplacian_var(gray, alpha) | |
| problem = None | |
| if fraction < config.MIN_OBJECT_FRACTION: | |
| problem = "Product not found in the photo (or far too small). Retake closer." | |
| elif sharp < config.MIN_SHARPNESS: | |
| problem = "Photo is too blurry to use as evidence. Retake with the camera steady." | |
| return CaptureQuality(round(float(fraction), 4), round(sharp, 1), problem is None, problem) | |