Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,583 Bytes
08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 08d991a 52aa226 |
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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
import logging
import torch
from PIL import Image
import numpy as np
from torchvision import transforms
from torchvision.models.segmentation import deeplabv3_resnet50
from transformers import (
SegformerForSemanticSegmentation,
SegformerFeatureExtractor,
AutoProcessor,
CLIPSegForImageSegmentation,
)
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
class Segmenter:
"""
Generalized Semantic Segmentation Wrapper for SegFormer, DeepLabV3, and CLIPSeg.
"""
def __init__(self, model_key="nvidia/segformer-b0-finetuned-ade-512-512", device="cpu"):
"""
Args:
model_key (str): HF model identifier or 'deeplabv3_resnet50'.
device (str): 'cpu' or 'cuda'.
"""
logger.info(f"Initializing Segmenter for model '{model_key}' on {device}")
self.model_key = model_key.lower()
self.device = device
self.model = None
self.processor = None # for transformers-based models
def _load_model(self):
"""
Lazy-load the model & processor based on model_key.
"""
if self.model is not None:
return
# SegFormer
if "segformer" in self.model_key:
self.model = SegformerForSemanticSegmentation.from_pretrained(self.model_key).to(self.device).eval()
self.processor = SegformerFeatureExtractor.from_pretrained(self.model_key)
# DeepLabV3
elif self.model_key == "deeplabv3_resnet50":
self.model = deeplabv3_resnet50(pretrained=True).to(self.device).eval()
self.processor = None
# CLIPSeg
elif "clipseg" in self.model_key:
self.model = CLIPSegForImageSegmentation.from_pretrained(self.model_key).to(self.device).eval()
self.processor = AutoProcessor.from_pretrained(self.model_key)
else:
raise ValueError(f"Unsupported segmentation model key: '{self.model_key}'")
logger.info(f"Loaded segmentation model '{self.model_key}'")
def predict(self, image: Image.Image, prompt: str = "", **kwargs) -> np.ndarray:
"""
Perform segmentation.
Args:
image (PIL.Image.Image): Input.
prompt (str): Only used for CLIPSeg.
Returns:
np.ndarray: Segmentation mask (H×W).
"""
self._load_model()
# SegFormer path
if "segformer" in self.model_key:
inputs = self.processor(images=image, return_tensors="pt").to(self.device)
outputs = self.model(**inputs)
mask = outputs.logits.argmax(dim=1).squeeze().cpu().numpy()
return mask
# DeepLabV3 path
if self.model_key == "deeplabv3_resnet50":
tf = transforms.ToTensor()
inp = tf(image).unsqueeze(0).to(self.device)
with torch.no_grad():
out = self.model(inp)["out"]
mask = out.argmax(1).squeeze().cpu().numpy()
return mask
# CLIPSeg path
if "clipseg" in self.model_key:
# CLIPSeg expects both text and image
inputs = self.processor(
text=[prompt], # list of prompts
images=[image], # list of images
return_tensors="pt"
).to(self.device)
with torch.no_grad():
outputs = self.model(**inputs)
# outputs.logits shape: (batch=1, height, width)
mask = outputs.logits.squeeze(0).cpu().numpy()
# Optionally threshold to binary:
# mask = (mask > kwargs.get("threshold", 0.5)).astype(np.uint8)
return mask
raise RuntimeError("Unreachable segmentation branch")
def draw(self, image: Image.Image, mask: np.ndarray, alpha=0.5) -> Image.Image:
"""
Overlay the segmentation mask on the input image.
Args:
image (PIL.Image.Image): Original.
mask (np.ndarray): Segmentation mask.
alpha (float): Blend strength.
Returns:
PIL.Image.Image: Blended output.
"""
logger.info("Drawing segmentation overlay")
# Normalize mask to 0–255
gray = ((mask - mask.min()) / (mask.ptp()) * 255).astype(np.uint8)
mask_img = Image.fromarray(gray).convert("L").resize(image.size)
# Make it RGB
color_mask = Image.merge("RGB", (mask_img, mask_img, mask_img))
return Image.blend(image, color_mask, alpha)
|