Spaces:
Running on Zero
Running on Zero
File size: 4,157 Bytes
ea2b072 df402fb ea2b072 8dc9082 ea2b072 8dc9082 ea2b072 523a925 37fcc27 523a925 ea2b072 37fcc27 ea2b072 37fcc27 8dc9082 ea2b072 df402fb 37fcc27 8dc9082 ea2b072 df402fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | 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."""
@abstractmethod
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.
"""
@abstractmethod
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
|