Update models/segmentation/segmenter.py

models/segmentation/segmenter.py

@@ -4,86 +4,121 @@ from PIL import Image
 import numpy as np
 from torchvision import transforms
 from torchvision.models.segmentation import deeplabv3_resnet50
-from transformers import SegformerForSemanticSegmentation, SegformerFeatureExtractor
+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 and DeepLabV3.
+    Generalized Semantic Segmentation Wrapper for SegFormer, DeepLabV3, and CLIPSeg.
     """
 
     def __init__(self, model_key="nvidia/segformer-b0-finetuned-ade-512-512", device="cpu"):
         """
-        Initialize the segmentation model.
-
         Args:
-            model_key (str): Model identifier, e.g., Hugging Face model id or 'deeplabv3_resnet50'.
-            device (str): Inference device ("cpu" or "cuda").
+            model_key (str): HF model identifier or 'deeplabv3_resnet50'.
+            device (str): 'cpu' or 'cuda'.
         """
-        logger.info(f"Initializing segmenter with model: {model_key}")
+        logger.info(f"Initializing Segmenter for model '{model_key}' on {device}")
+        self.model_key = model_key.lower()
         self.device = device
-        self.model_key = model_key
-        self.model, self.processor = self._load_model()
+        self.model = None
+        self.processor = None  # for transformers-based models
 
     def _load_model(self):
         """
-        Load the segmentation model and processor.
-
-        Returns:
-            Tuple[torch.nn.Module, Optional[Processor]]
+        Lazy-load the model & processor based on model_key.
         """
+        if self.model is not None:
+            return
+
+        # SegFormer
         if "segformer" in self.model_key:
-            model = SegformerForSemanticSegmentation.from_pretrained(self.model_key).to(self.device)
-            processor = SegformerFeatureExtractor.from_pretrained(self.model_key)
-            return model, processor
+            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":
-            model = deeplabv3_resnet50(pretrained=True).to(self.device).eval()
-            return model, None
+            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 model key: {self.model_key}")
+            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, **kwargs):
+    def predict(self, image: Image.Image, prompt: str = "", **kwargs) -> np.ndarray:
         """
-        Perform segmentation on the input image.
+        Perform segmentation.
 
         Args:
-            image (PIL.Image.Image): Input image.
-
+            image (PIL.Image.Image): Input.
+            prompt (str): Only used for CLIPSeg.
         Returns:
-            np.ndarray: Segmentation mask.
+            np.ndarray: Segmentation mask (H×W).
         """
-        logger.info("Running segmentation")
+        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
 
-        elif self.model_key == "deeplabv3_resnet50":
-            transform = transforms.Compose([
-                transforms.ToTensor(),
-            ])
-            inputs = transform(image).unsqueeze(0).to(self.device)
+        # 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)["out"]
-                mask = outputs.argmax(1).squeeze().cpu().numpy()
-                return mask
+                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
 
-    def draw(self, image: Image.Image, mask: np.ndarray, alpha=0.5):
+        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 image.
+            image (PIL.Image.Image): Original.
             mask (np.ndarray): Segmentation mask.
             alpha (float): Blend strength.
-
         Returns:
-            PIL.Image.Image: Image with mask overlay.
+            PIL.Image.Image: Blended output.
         """
         logger.info("Drawing segmentation overlay")
-        mask_img = Image.fromarray((mask * 255 / mask.max()).astype(np.uint8)).convert("L").resize(image.size)
-        mask_colored = Image.merge("RGB", (mask_img, mask_img, mask_img))
-        return Image.blend(image, mask_colored, alpha)
+        # 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)