Spaces:
Sleeping
Sleeping
File size: 2,409 Bytes
cfbbb2f 28470ee cfbbb2f 28470ee cfbbb2f 28470ee cfbbb2f 28470ee cfbbb2f 28470ee cfbbb2f 28470ee cfbbb2f 28470ee cfbbb2f 28470ee cfbbb2f 28470ee cfbbb2f 28470ee | 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 | import os
import sys
import torch
import numpy as np
from PIL import Image
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'sam2'))
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
MODEL_CFG = "configs/sam2.1/sam2.1_hiera_s.yaml"
CHECKPOINT = os.path.join(os.path.dirname(__file__), "models", "sam2.1_hiera_small.pt")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
SAM2_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'sam2')
print(f"[SAM2] Using device: {DEVICE}")
def load_sam2_model():
os.chdir(SAM2_DIR)
model = build_sam2(MODEL_CFG, CHECKPOINT, device=DEVICE)
predictor = SAM2ImagePredictor(model)
print("[SAM2] Model loaded successfully")
return predictor
def get_mask_for_box(predictor, image_rgb: np.ndarray, box: list) -> np.ndarray:
"""
Get precise pixel mask for a single bounding box using SAM2.
box = [x1, y1, x2, y2] in absolute pixels.
Returns boolean mask same size as image.
"""
predictor.set_image(image_rgb)
box_array = np.array(box, dtype=np.float32)
with torch.inference_mode():
with torch.autocast(device_type=DEVICE, dtype=torch.float16):
masks, scores, _ = predictor.predict(
point_coords = None,
point_labels = None,
box = box_array[None, :], # SAM2 expects (1, 4)
multimask_output = True,
)
# pick highest confidence mask
best_idx = scores.argmax()
best_mask = masks[best_idx].astype(bool)
return best_mask
def mask_to_transparent_png(image_rgb: np.ndarray, mask: np.ndarray) -> Image.Image:
"""
Apply mask to image — keep masked pixels, make everything else transparent.
Returns RGBA PIL image.
"""
rgba = np.zeros((*image_rgb.shape[:2], 4), dtype=np.uint8)
rgba[..., :3] = image_rgb
rgba[..., 3] = (mask * 255).astype(np.uint8) # alpha = 255 where object is
# crop to bounding box of mask to minimise image size
rows = np.where(mask.any(axis=1))[0]
cols = np.where(mask.any(axis=0))[0]
if len(rows) == 0 or len(cols) == 0:
return Image.fromarray(rgba, 'RGBA')
y1, y2 = int(rows.min()), int(rows.max())
x1, x2 = int(cols.min()), int(cols.max())
cropped = rgba[y1:y2+1, x1:x2+1]
return Image.fromarray(cropped, 'RGBA')
|