Spaces:
Running on Zero
Running on Zero
| from abc import ABC, abstractmethod | |
| from PIL import Image, ImageFilter | |
| def resize_for_upload(image: Image.Image, max_dimension: int = 1536) -> Image.Image: | |
| """Downscale so the longest edge is at most max_dimension, keeping aspect ratio.""" | |
| width, height = image.size | |
| longest_edge = max(width, height) | |
| if longest_edge <= max_dimension: | |
| return image | |
| scale = max_dimension / longest_edge | |
| new_size = (round(width * scale), round(height * scale)) | |
| return image.resize(new_size, Image.LANCZOS) | |
| class Segmenter(ABC): | |
| """Produces object masks from an image and point prompts.""" | |
| def encode(self, image: Image.Image) -> dict: | |
| """Validate/warm up the model for this image; return small JSON-safe metadata. | |
| Must not return tensors or other GPU-resident objects: under ZeroGPU, | |
| encode() and decode() may run in different ephemeral worker | |
| processes, and only the host (which has no real GPU) can persist | |
| data between them - so decode() re-derives the embedding from the | |
| original image instead of reusing anything encode() returns. | |
| """ | |
| def decode(self, image: Image.Image, points: list[tuple[int, int, bool]]) -> Image.Image: | |
| """Encode the image and decode it with the given point prompts; return an 'L' mode mask.""" | |
| import numpy as np | |
| class Sam2ModelSegmenter(Segmenter): | |
| """Wraps SAM2 hiera-base-plus via the `sam2` package for server-side inference.""" | |
| def __init__(self, hf_model_id: str = "facebook/sam2.1-hiera-base-plus", device: str = "cuda"): | |
| # Do NOT build the model here: __init__ runs at Space startup, which | |
| # is outside any @spaces.GPU-decorated call. ZeroGPU's CUDA | |
| # emulation mode active at that point cannot correctly satisfy | |
| # SAM2's Hydra-based instantiation of MemoryAttention (it raises | |
| # "Expected all tensors to be on the same device ... cuda:0 | |
| # (ZeroGPU) and cpu" regardless of whether "cpu" or "cuda" is | |
| # requested). Build lazily on first encode(), which - once this | |
| # method is wrapped with @spaces.GPU - runs inside a real, | |
| # physically-attached GPU worker instead of the emulation. | |
| self._hf_model_id = hf_model_id | |
| self._device = device | |
| self._predictor = None | |
| def _ensure_loaded(self) -> None: | |
| if self._predictor is not None: | |
| return | |
| from sam2.build_sam import build_sam2_hf | |
| from sam2.sam2_image_predictor import SAM2ImagePredictor | |
| sam2_model = build_sam2_hf(self._hf_model_id, device=self._device) | |
| self._predictor = SAM2ImagePredictor(sam2_model) | |
| def encode(self, image: Image.Image) -> dict: | |
| self._ensure_loaded() | |
| return {"image_size": image.size} | |
| def decode( | |
| self, image: Image.Image, points: list[tuple[int, int, bool]], edge_padding: int = 6 | |
| ) -> Image.Image: | |
| self._ensure_loaded() | |
| image_array = np.array(image.convert("RGB")) | |
| self._predictor.set_image(image_array) | |
| point_coords = np.array([[x, y] for x, y, _ in points]) | |
| point_labels = np.array([1 if positive else 0 for _, _, positive in points]) | |
| masks, scores, _ = self._predictor.predict( | |
| point_coords=point_coords, | |
| point_labels=point_labels, | |
| multimask_output=True, | |
| ) | |
| best_mask = masks[int(np.argmax(scores))] | |
| mask_image = Image.fromarray((best_mask * 255).astype("uint8"), mode="L") | |
| mask_image = mask_image.resize(image.size) | |
| # SAM2's boundary tends to sit a few pixels inside the object's true | |
| # edge rather than exactly on it, leaving a visible sliver of the | |
| # object outside the selection - padding it out here means the mask | |
| # the user sees (and that gets sent for removal) already comfortably | |
| # covers the object, instead of relying on them to notice and patch | |
| # small gaps with extra taps/brush strokes. | |
| if edge_padding > 0: | |
| mask_image = mask_image.filter(ImageFilter.MaxFilter(edge_padding * 2 + 1)) | |
| return mask_image | |