Spaces:
Running on Zero
Running on Zero
| """Local visual features: object-only DINOv2 embeddings and a colour signature. | |
| Why object-only pooling: the cut-out is padded to a square before the ViT sees | |
| it. For tall items (bottles, shakers) 65-77% of the patches are padding, and a | |
| plain patch-mean then scores *any* two tall objects as similar because they | |
| share the padding. We average only the patches the background-removal mask | |
| marks as object. | |
| Why a colour signature: DINOv2 is trained with colour-jitter augmentation, so it | |
| is deliberately colour-invariant. A blue speaker and a black speaker of the same | |
| model look almost identical to it. The median chroma of the object pixels | |
| catches exactly that substitution, for free. | |
| """ | |
| from dataclasses import dataclass | |
| from functools import lru_cache | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from returnguard import config | |
| _MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) | |
| _STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) | |
| _PAD = tuple(int(v * 255) for v in _MEAN) # pad with the normalisation mean -> ~zero input | |
| def _model(): | |
| from transformers import AutoModel | |
| model = AutoModel.from_pretrained(config.EMBED_MODEL) | |
| model.eval() | |
| return model | |
| def _square(img: Image.Image, fill) -> Image.Image: | |
| side = max(img.size) | |
| canvas = Image.new(img.mode, (side, side), fill) | |
| canvas.paste(img, ((side - img.width) // 2, (side - img.height) // 2)) | |
| return canvas | |
| class ViewFeatures: | |
| embedding: np.ndarray # L2-normalised object-only patch mean, float32 [D] | |
| lab: np.ndarray # median object colour (L*, a*, b*), float32 [3] | |
| object_patches: int | |
| def to_dict(self): | |
| return {"embedding": self.embedding.tolist(), "lab": self.lab.tolist(), | |
| "object_patches": self.object_patches} | |
| def from_dict(cls, d): | |
| return cls(np.asarray(d["embedding"], dtype=np.float32), np.asarray(d["lab"], dtype=np.float32), | |
| int(d["object_patches"])) | |
| def _rgb_to_lab(rgb: np.ndarray) -> np.ndarray: | |
| """sRGB [0,255] -> CIE Lab (D65).""" | |
| c = rgb.astype(np.float32) / 255.0 | |
| c = np.where(c > 0.04045, ((c + 0.055) / 1.055) ** 2.4, c / 12.92) | |
| xyz = c @ np.array([[0.4124, 0.3576, 0.1805], [0.2126, 0.7152, 0.0722], [0.0193, 0.1192, 0.9505]]).T | |
| xyz /= np.array([0.95047, 1.0, 1.08883]) | |
| f = np.where(xyz > 0.008856, np.cbrt(xyz), 7.787 * xyz + 16 / 116) | |
| L = 116 * f[:, 1] - 16 | |
| a = 500 * (f[:, 0] - f[:, 1]) | |
| b = 200 * (f[:, 1] - f[:, 2]) | |
| return np.stack([L, a, b], axis=1) | |
| def color_signature(cutout: Image.Image) -> np.ndarray: | |
| """Median CIE Lab colour of the object pixels (L*, a*, b*). | |
| A median over thousands of object pixels is robust to highlights, shadows | |
| and small background-removal errors. A histogram was tried first and was too | |
| sensitive to white-balance shifts between two photos of the same item.""" | |
| arr = np.array(cutout.convert("RGBA")) | |
| pix = arr[arr[..., 3] > 127][:, :3] | |
| if len(pix) < 50: | |
| return np.full(3, np.nan, dtype=np.float32) | |
| if len(pix) > 40000: | |
| pix = pix[np.random.default_rng(0).choice(len(pix), 40000, replace=False)] | |
| return np.median(_rgb_to_lab(pix), axis=0).astype(np.float32) | |
| def view_features(cutout: Image.Image) -> ViewFeatures: | |
| size = config.EMBED_SIZE | |
| rgba = cutout.convert("RGBA") | |
| rgb = _square(rgba.convert("RGB"), _PAD).resize((size, size), Image.BICUBIC) | |
| # put the original pixels back only where the object is; background -> pad colour | |
| alpha_sq = _square(rgba.getchannel("A"), 0).resize((size, size), Image.BILINEAR) | |
| a = np.array(alpha_sq, dtype=np.float32)[..., None] / 255.0 | |
| pixels = np.array(rgb, dtype=np.float32) * a + np.array(_PAD, dtype=np.float32) * (1 - a) | |
| x = (pixels / 255.0 - _MEAN) / _STD | |
| x = torch.from_numpy(x.transpose(2, 0, 1)).unsqueeze(0) | |
| tokens = _model()(pixel_values=x).last_hidden_state[0] # [1 + N (+ registers), D] | |
| grid = size // 14 | |
| patches = tokens[-grid * grid:] # patch tokens are the last N | |
| patch_alpha = np.array(alpha_sq.resize((grid, grid), Image.BOX), dtype=np.float32).flatten() / 255.0 | |
| weights = torch.from_numpy((patch_alpha > 0.5).astype(np.float32)) | |
| if weights.sum() < 1: | |
| weights = torch.ones_like(weights) | |
| emb = (patches * weights[:, None]).sum(0) / weights.sum() | |
| emb = torch.nn.functional.normalize(emb, dim=0) | |
| return ViewFeatures(emb.numpy().astype(np.float32), color_signature(cutout), int(weights.sum())) | |
| def shape_similarity(a: ViewFeatures, b: ViewFeatures) -> float: | |
| return float(np.dot(a.embedding, b.embedding)) | |
| def color_shift(a: ViewFeatures, b: ViewFeatures) -> float: | |
| """Chroma difference (Delta E in the a*b* plane) between the two objects' median | |
| colours. Lightness is left out because exposure changes it between photos. | |
| Re-photographing the same item stays below ~2.6 on our data; the blue vs black | |
| speaker of the same model is ~4.""" | |
| if np.isnan(a.lab).any() or np.isnan(b.lab).any(): | |
| return 0.0 | |
| return float(np.linalg.norm(a.lab[1:] - b.lab[1:])) | |