Init

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,15 @@
+---
+title: RADSeg
+emoji: 📚
+colorFrom: gray
+colorTo: blue
+sdk: gradio
+sdk_version: 6.3.0
+python_version: '3.11'
+app_file: radseg_demo.py
+pinned: true
+license: mit
+short_description: Fast & light zero-shot open vocabulary semantic segmentation
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

radseg/base.py

@@ -0,0 +1,216 @@
+import abc
+from typing_extensions import Tuple, List, override
+import torch
+
+from radseg.prompt_templates import openai_imagenet_template
+
+class ImageEncoder(abc.ABC):
+  """Interface for all image encoders"""
+
+  def __init__(self, device = None):
+    if device is None:
+      self.device = "cuda" if torch.cuda.is_available() else "cpu"
+    else:
+      self.device = device
+
+  @abc.abstractmethod
+  def is_compatible_size(self, h: int, w: int) -> bool:
+    """Checks if height and width of image is compatible with encoder"""
+    pass
+
+  @abc.abstractmethod
+  def get_nearest_size(self, h, w) -> Tuple[int, int]:
+    """Returns nearest compatible size as (h,w) tuple"""
+    pass
+
+class ImageGlobalEncoder(ImageEncoder):
+  """Interface for all image level (global) encoders"""
+
+  @abc.abstractmethod
+  def encode_image_to_vector(
+    self, rgb_image: torch.FloatTensor) -> torch.FloatTensor:
+    """Encode rgb image into an image level feature vector
+
+    Args:
+      rgb_image: Bx3xHxW float tensor representing an image with values in [0-1]
+    Returns:
+      BxC float tensor representing global features.
+    """
+
+class ImageSpatialEncoder(ImageEncoder):
+  """Interface for all pixel level / patch level encoders"""
+
+  @abc.abstractmethod
+  def encode_image_to_feat_map(
+    self, rgb_image: torch.FloatTensor) -> torch.FloatTensor:
+    """Encode rgb image into a feature map
+
+    Args:
+      rgb_image: Bx3xHxW float tensor representing an image with values in [0-1]
+    Returns:
+      BxCxH'xW' float tensor representing pixel / patch level features.
+    """
+
+class ImageSpatialGlobalEncoder(ImageSpatialEncoder, ImageGlobalEncoder):
+  "Interface for encoders that implement both global and spatial encoding APIs"
+
+  @abc.abstractmethod
+  def encode_image_to_feat_map_and_vector(self, rgb_image: torch.FloatTensor) \
+      -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+    """Encode rgb image into a feature map and an image level feature vector
+    
+    Args:
+      rgb_image: Bx3xHxW float tensor representing an image with values in [0-1]
+    Returns:
+      A tuple of 
+      - BxCxH'xW' float tensor representing pixel / patch level features.
+      - BxC float tensor representing global features.
+    """
+
+class LangImageEncoder(ImageEncoder):
+  """Interface for all image encoders that produce language aligned features"""
+
+  def __init__(self, device = None):
+    super().__init__(device)
+    self.prompt_templates = openai_imagenet_template
+
+  def insert_labels_into_templates(self, labels: List[str]) -> List[List[str]]:
+    """Inserts each labels into a set of stored templates.
+    
+    Args:
+      labels: A list of length T of class names / labels. Ex. ['cat', 'dog']
+    Returns:
+      labeled_templates as a list of length T of a list of length P of strings.
+      where T is the number of labels and P is the number of templates.
+    """
+    return [[pt(x) for pt in self.prompt_templates] for x in labels]
+
+  @abc.abstractmethod
+  def encode_labels(self, labels: List[str]) -> torch.FloatTensor:
+    """Encodes a list of labels into feature vectors using template prompts
+    
+    Args:
+      labels: A list of length T of class names / labels. Ex. ['cat', 'dog']
+    Returns:
+      A TxD float tensor where T represents the number of classes and D 
+      represents the feature space dimension.
+    """
+    pass
+
+  @abc.abstractmethod
+  def encode_prompts(self, prompts: List[str]) -> torch.FloatTensor:
+    """Encodes a list of prompts into feature vectors directly
+    
+    Args:
+      prompts: A list of strings of arbitrary length of prompts to encode.
+    Returns:
+      A TxD float tensor where T represents the number of prompts and D 
+      represents the feature space dimension.
+    """
+    pass
+
+class LangSpatialImageEncoder(LangImageEncoder, ImageSpatialEncoder):
+
+  @abc.abstractmethod
+  def align_spatial_features_with_language(
+    self, features: torch.FloatTensor) -> torch.FloatTensor:
+    """Projects image features to the language space if it is not already there.
+
+    Some encoders directly output the language aligned features in which case
+    this function will simply be a pass through.
+
+    Args:
+      features: A (BxCxHxW) float tensor representing spatial features obtained
+        from encode_image_to_feat_map.
+    Returns:
+      A (BxDxHxW) float tensor that is in the same feature space as the text 
+      embeddings obtained from encode_labels or encode_prompts.
+    """
+    pass
+
+class LangGlobalImageEncoder(LangImageEncoder, ImageGlobalEncoder):
+
+  @abc.abstractmethod
+  def align_global_features_with_language(
+    self, features: torch.FloatTensor) -> torch.FloatTensor:
+    """Projects image global features to the language space if needed
+
+    Some encoders directly output the language aligned features in which case
+    this function will simply be a pass through.
+
+    Args:
+      features: A (BxC) float tensor representing global features obtained
+        from encode_image_to_vector.
+    Returns:
+      A (BxD) float tensor that is in the same feature space as the text 
+      embeddings obtained from encode_labels or encode_prompts.
+    """
+    pass
+
+
+class LangSpatialGlobalImageEncoder(LangGlobalImageEncoder,
+                                    LangSpatialImageEncoder):
+  pass
+
+
+class ImageSemSegEncoder(LangSpatialImageEncoder):
+  """Interface for encoders that are actually semantic segmentors.
+  
+  To avoid rewriting major portions of the codebase. We consider semantic
+  segmentation as simply a onehot encoder but nevertheless still producing a 
+  feature vector that can be aggregated and compared with using cos-sim...etc.
+  Would be more efficient to deal with integers of course but we create this
+  for convenience. This way these "encoders" can be used with existing  mappers.
+  """
+
+  def __init__(self, device = None):
+    super().__init__(device)
+    self.eps = 1e-10
+
+  @property
+  @abc.abstractmethod
+  def num_classes(self) -> int:
+    pass
+
+  @property
+  @abc.abstractmethod
+  def cat_index_to_name(self):
+    """Returns a mapping from category index to category name.
+    
+    cat_index_to_name[cat_index] gives the name of that category index.
+    """
+    pass
+
+  @property
+  @abc.abstractmethod
+  def cat_name_to_index(self):
+    """Returns a mapping from category index to category name.
+    
+    cat_name_to_index[cat_name] gives the one hot index of that category name.
+    """
+    pass
+
+  @override
+  def encode_labels(self, labels: List[str]) -> torch.FloatTensor:
+    """Encodes a list of category names into onehot vectors using stored mapping
+    
+    Args:
+      labels: A list of length T of category names / labels. Ex. ['cat', 'dog']
+    Returns:
+      A TxD float tensor where T represents the number of classes and D 
+      represents the feature space dimension.
+    """
+    onehot = torch.full((len(labels), self.num_classes), self.eps,
+                         dtype=torch.float, device=self.device)
+    for i,c in enumerate(labels):
+      onehot[i, self.cat_name_to_index[c]] = 1
+    return onehot
+
+  @override
+  def encode_prompts(self, prompts: List[str]) -> torch.FloatTensor:
+    """alias for encode_labels"""
+    return self.encode_labels(prompts)
+
+  @override
+  def align_spatial_features_with_language(self, features):
+    return features

radseg/prompt_templates.py

@@ -0,0 +1,84 @@
+"""Includes common prompt templates for encoding labels"""
+
+openai_imagenet_template = [
+    lambda c: f'a bad photo of a {c}.',
+    lambda c: f'a photo of many {c}.',
+    lambda c: f'a sculpture of a {c}.',
+    lambda c: f'a photo of the hard to see {c}.',
+    lambda c: f'a low resolution photo of the {c}.',
+    lambda c: f'a rendering of a {c}.',
+    lambda c: f'graffiti of a {c}.',
+    lambda c: f'a bad photo of the {c}.',
+    lambda c: f'a cropped photo of the {c}.',
+    lambda c: f'a tattoo of a {c}.',
+    lambda c: f'the embroidered {c}.',
+    lambda c: f'a photo of a hard to see {c}.',
+    lambda c: f'a bright photo of a {c}.',
+    lambda c: f'a photo of a clean {c}.',
+    lambda c: f'a photo of a dirty {c}.',
+    lambda c: f'a dark photo of the {c}.',
+    lambda c: f'a drawing of a {c}.',
+    lambda c: f'a photo of my {c}.',
+    lambda c: f'the plastic {c}.',
+    lambda c: f'a photo of the cool {c}.',
+    lambda c: f'a close-up photo of a {c}.',
+    lambda c: f'a black and white photo of the {c}.',
+    lambda c: f'a painting of the {c}.',
+    lambda c: f'a painting of a {c}.',
+    lambda c: f'a pixelated photo of the {c}.',
+    lambda c: f'a sculpture of the {c}.',
+    lambda c: f'a bright photo of the {c}.',
+    lambda c: f'a cropped photo of a {c}.',
+    lambda c: f'a plastic {c}.',
+    lambda c: f'a photo of the dirty {c}.',
+    lambda c: f'a jpeg corrupted photo of a {c}.',
+    lambda c: f'a blurry photo of the {c}.',
+    lambda c: f'a photo of the {c}.',
+    lambda c: f'a good photo of the {c}.',
+    lambda c: f'a rendering of the {c}.',
+    lambda c: f'a {c} in a video game.',
+    lambda c: f'a photo of one {c}.',
+    lambda c: f'a doodle of a {c}.',
+    lambda c: f'a close-up photo of the {c}.',
+    lambda c: f'a photo of a {c}.',
+    lambda c: f'the origami {c}.',
+    lambda c: f'the {c} in a video game.',
+    lambda c: f'a sketch of a {c}.',
+    lambda c: f'a doodle of the {c}.',
+    lambda c: f'a origami {c}.',
+    lambda c: f'a low resolution photo of a {c}.',
+    lambda c: f'the toy {c}.',
+    lambda c: f'a rendition of the {c}.',
+    lambda c: f'a photo of the clean {c}.',
+    lambda c: f'a photo of a large {c}.',
+    lambda c: f'a rendition of a {c}.',
+    lambda c: f'a photo of a nice {c}.',
+    lambda c: f'a photo of a weird {c}.',
+    lambda c: f'a blurry photo of a {c}.',
+    lambda c: f'a cartoon {c}.',
+    lambda c: f'art of a {c}.',
+    lambda c: f'a sketch of the {c}.',
+    lambda c: f'a embroidered {c}.',
+    lambda c: f'a pixelated photo of a {c}.',
+    lambda c: f'itap of the {c}.',
+    lambda c: f'a jpeg corrupted photo of the {c}.',
+    lambda c: f'a good photo of a {c}.',
+    lambda c: f'a plushie {c}.',
+    lambda c: f'a photo of the nice {c}.',
+    lambda c: f'a photo of the small {c}.',
+    lambda c: f'a photo of the weird {c}.',
+    lambda c: f'the cartoon {c}.',
+    lambda c: f'art of the {c}.',
+    lambda c: f'a drawing of the {c}.',
+    lambda c: f'a photo of the large {c}.',
+    lambda c: f'a black and white photo of a {c}.',
+    lambda c: f'the plushie {c}.',
+    lambda c: f'a dark photo of a {c}.',
+    lambda c: f'itap of a {c}.',
+    lambda c: f'graffiti of the {c}.',
+    lambda c: f'a toy {c}.',
+    lambda c: f'itap of my {c}.',
+    lambda c: f'a photo of a cool {c}.',
+    lambda c: f'a photo of a small {c}.',
+    lambda c: f'a tattoo of the {c}.',
+]

radseg/radseg.py

@@ -0,0 +1,547 @@
+"""Includes the RADSeg Encoder.
+
+The module only relies on the base.py and prompt_templates.py files in
+image_encoders. Encoder can be copied with those files to your own project.
+
+Typical Usage:
+
+  rgb_img = torchvision.io.read_image(rgb_path)
+  rgb_img = rgb_img.float() / 255
+  rgb_img = torch.nn.functional.interpolate(
+    rgb_img.unsqueeze(0),size=(512, 512))
+
+  labels = ["car", "person"]
+
+  enc = RADSegEncoder(model_version="c-radio_v3-b", lang_model="siglip2")
+  
+  feat_map = enc.encode_image_to_feat_map(rgb_img)
+  lang_aligned_feat_map = enc.align_spatial_features_with_language(feat_map)
+
+  text_features = enc.encode_labels(labels)
+
+  from rayfronts.utils import compute_cos_sim
+  r = compute_cos_sim(text_features, lang_aligned_feat_map, softmax=True)
+"""
+
+from typing_extensions import override, List, Tuple
+
+import torch
+import numpy as np
+
+from radseg.base import ImageSemSegEncoder
+from segment_anything import sam_model_registry, SamPredictor
+from radseg.sam_utils import sam_refinement
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from timm.layers import use_fused_attn
+import math
+
+def compute_cos_sim(vec1: torch.FloatTensor,
+                    vec2: torch.FloatTensor,
+                    softmax: bool = False) -> torch.FloatTensor:
+    """Compute cosine similarity between two batches of D dim vectors.
+
+    Args:
+    vec1: NxC float tensor representing batch of vectors
+    vec2: MxC float tensor representing batch of vectors
+    softmax: If False, cosine similarity is returned. If True, softmaxed
+        probability is returned across the N dimension.
+    Returns:
+    result: MxN float tensor representing similarity/prob. where result[0,1]
+        represents the similarity of vec1[0] with vec2[1]
+    """
+    N, C1 = vec1.shape
+    M, C2 = vec2.shape
+    if C1 != C2:
+        raise ValueError(f"vec1 feature dimension '{C1}' does not match vec2"
+                        f"feature dimension '{C2}'")
+    C = C1
+
+    vec1 = vec1 / vec1.norm(dim = -1, keepdim = True)
+    vec1 = vec1.reshape(1, N, 1, C)
+
+    vec2 = vec2 / vec2.norm(dim=-1, keepdim = True)
+    vec2 = vec2.reshape(M, 1, C, 1)
+
+    sim = (vec1 @ vec2).reshape(M, N)
+    if softmax:
+        return torch.softmax(100 * sim, dim = -1)
+    else:
+        return sim
+
+class SelfCorrelatingRecursiveAttn(nn.Module):
+  def __init__(
+    self,
+    orig_attn,
+    device,
+    dim: int,
+    qk_norm: bool = False,
+    scra_scaling: int = 10
+  ) -> None:
+    super().__init__()
+    num_heads = orig_attn.num_heads
+    assert dim % num_heads == 0, "dim should be divisible by num_heads"
+    self.num_heads = num_heads
+    self.head_dim = dim // num_heads
+    self.scale = self.head_dim ** -0.5
+    self.fused_attn = use_fused_attn()
+
+    self.qkv = orig_attn.qkv
+    self.q_norm = orig_attn.q_norm if qk_norm else nn.Identity()
+    self.k_norm = orig_attn.k_norm if qk_norm else nn.Identity()
+    self.attn_drop = orig_attn.attn_drop
+    self.proj = orig_attn.proj
+    self.proj_drop = orig_attn.proj_drop
+    self.device = device
+    self.scra_scaling = scra_scaling
+
+  def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+    B, N, C = x.shape
+    x_out = self.custom_attn(x.permute(1, 0, 2))
+    x_out = x_out.permute(1, 0, 2)
+    return x_out
+
+  
+  def custom_attn(self, x):
+    num_heads = self.num_heads
+    num_tokens, bsz, embed_dim = x.size()
+    head_dim = embed_dim // num_heads
+    scale = head_dim ** -0.5
+    q, k, v = self.qkv(x).chunk(3, dim=-1)
+    q, k = self.q_norm(q), self.k_norm(k)
+
+    q = q.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+    k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+    v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+
+    attn_weights = torch.bmm(q, k.transpose(1, 2)) * scale
+    attn_weights = F.softmax(attn_weights, dim=-1)
+    attn_output = torch.bmm(attn_weights, v)
+    attn_output = attn_output.transpose(0, 1).contiguous().view(
+      -1, bsz, embed_dim)
+    attn_output = self.proj(attn_output)
+    attn_output = self.proj_drop(attn_output)
+
+    # Self Correlating Recursive Attention 
+    attn_output = attn_output.view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+    sim_tokens =  F.normalize(attn_output,dim = -1)
+    sim_matrix = torch.bmm(sim_tokens,sim_tokens.transpose(1,2))*self.scra_scaling
+    sim_matrix[sim_matrix < 0] = -torch.inf
+    sim_matrix = F.softmax(sim_matrix,dim=-1)
+    attn_output = torch.bmm(sim_matrix, v)    
+    attn_output = attn_output.transpose(0, 1).contiguous().view(
+      -1, bsz, embed_dim)
+    attn_output = self.proj(attn_output)
+    attn_output = self.proj_drop(attn_output)
+    
+    return attn_output
+
+class RADSegEncoder(ImageSemSegEncoder):
+  # TODO: Split into feature encoder and semseg encoder for clarity instead
+  # of switching behaviors based on flags.
+
+  def __init__(self,
+               device: str = None,
+               model_version: str = "radio_v3-b",
+               lang_model: str ="siglip2",
+               return_radio_features: bool = True,
+               compile: bool = False,
+               amp: bool = False,
+               predict: bool = False,
+               classes: List[str] = None,
+               text_query_mode: str = "labels",
+               prediction_thresh: float = 0.0,
+               prompt_denoising_thresh: float = 0.5,
+               scra_scaling: float = 10.0,
+               scga_scaling: float = 10.0,
+               slide_crop: int = 336,
+               slide_stride: int = 224,
+               # Sam refinement args
+               sam_refinement: bool = False,
+               sam_ckpt: str = 'sam_vit_h_4b8939.pth',
+               coarse_thresh: float = 0.10,
+               minimal_area: int = 225,
+               sam_mask_coff: float = 0.005,
+               sam_iou_thresh: float = 0.9):
+    
+    super().__init__(device)
+
+    self.compile = compile
+    self.amp = amp
+    self.model_version = model_version
+    self.return_radio_features = return_radio_features
+    adaptor_names = [lang_model, "sam"]
+    self.model = torch.hub.load("NVlabs/RADIO", "radio_model",
+                                version=model_version, progress=True,
+                                skip_validation=True,
+                                adaptor_names=adaptor_names)
+    self.model.eval()
+    self.model = self.model.to(self.device)
+    # Steal adaptors from RADIO so it does not auto compute adaptor output.
+    # We want to control when that happens.
+    self.lang_adaptor = self.model.adaptors[lang_model]
+    self.sam_adaptor = self.model.adaptors["sam"]
+    self.model.adaptors = None
+    last_block = self.model.model.blocks[-1]
+    last_block.attn = SelfCorrelatingRecursiveAttn(
+      last_block.attn,
+      dim=self.model.model.embed_dim,
+      device=self.device,
+      scra_scaling=scra_scaling)
+
+    if self.compile:
+      self.model.compile(fullgraph=True, options={"triton.cudagraphs":True})
+      self.lang_adaptor.compile(fullgraph=True, options={"triton.cudagraphs":True})
+
+    self.predict = predict
+    self.prediction_thresh = prediction_thresh
+    self.text_query_mode = text_query_mode
+    self.prompt_denoising_thresh = prompt_denoising_thresh
+
+    if self.predict:
+      if classes is None or len(classes) < 1:
+        raise Exception("Must pass list of classes when predict is True")
+      self.prompts = list(classes)
+
+      if len(self.prompts[0]) == 0: # Remove ignore class so we don't prompt
+        self.prompts.pop(0)
+      self._cat_index_to_name = {0: ''}
+      self._cat_index_to_name.update({i+1: v for i,v in enumerate(self.prompts)})
+      self._cat_name_to_index = {
+        v: k for k, v in self._cat_index_to_name.items()
+      }
+      if self.text_query_mode == "labels":
+        self.text_embeds = self.encode_labels(self.prompts, onehot=False)
+      elif self.text_query_mode == "prompts":
+        self.text_embeds = self.encode_prompts(self.prompts, onehot=False)
+      else:
+        raise ValueError("Invalid query type")
+      
+
+    self.slide_stride = slide_stride
+    self.slide_crop = slide_crop
+    self.scga_scaling = scga_scaling
+
+    # Sam refinement args
+    self.sam_refinement = sam_refinement
+    if sam_refinement:
+      self.sam_iou_thresh = sam_iou_thresh
+      self.coarse_thresh = coarse_thresh
+      self.minimal_area = minimal_area
+      self.sam_mask_coff = sam_mask_coff
+      self.sam = sam_model_registry["vit_h"](checkpoint=sam_ckpt).to(device=self.device).eval()
+      self.sam_predictor = SamPredictor(self.sam)
+      del self.sam.image_encoder.blocks
+      del self.sam.image_encoder.patch_embed
+      torch.cuda.empty_cache()
+
+  @property
+  @override
+  def num_classes(self) -> int:
+    return len(self.prompts) + 1
+
+  @property
+  @override
+  def cat_index_to_name(self):
+    return self._cat_index_to_name
+
+  @property
+  @override
+  def cat_name_to_index(self):
+    return self._cat_name_to_index
+
+  @override
+  def encode_labels(self, labels: List[str], onehot: bool = True) -> torch.FloatTensor:
+    if self.predict and onehot:
+      return super().encode_labels(labels)
+    prompts_per_label = self.insert_labels_into_templates(labels)
+    all_text_features = list()
+    for i in range(len(labels)):
+      text_features = self.encode_prompts(prompts_per_label[i], onehot=False)
+      text_features = text_features.mean(dim=0, keepdim=True)
+      all_text_features.append(text_features)
+
+    all_text_features = torch.cat(all_text_features, dim=0)
+    return all_text_features
+
+  @override
+  def encode_prompts(self, prompts: List[str], onehot: bool = True) -> torch.FloatTensor:
+    if self.predict and onehot:
+      return super().encode_labels(prompts)
+    with torch.autocast("cuda", dtype=torch.float16, enabled=self.amp):
+      text = self.lang_adaptor.tokenizer(prompts).to(self.device)
+      text_features = self.lang_adaptor.encode_text(text)
+      text_features /= text_features.norm(dim=-1, keepdim=True)
+    return text_features
+
+  @override
+  def encode_image_to_feat_map(
+    self, rgb_image: torch.FloatTensor, orig_img_size: Tuple[int] = None,
+    return_preds: bool = False, ignore_label: bool = True) -> torch.FloatTensor:
+    if orig_img_size is None:
+      orig_img_size = rgb_image.shape[-2:]
+    
+    if self.slide_crop > 0:
+      feat_map = self._sliding_inference(rgb_image, stride=self.slide_stride, crop_size=self.slide_crop)
+    else:
+      feat_map = self._single_inference(rgb_image)
+    
+    feat_map = self._self_correlating_global_aggregation(feat_map)
+    if not self.predict:
+      return feat_map
+
+    # Predict
+    seg_probs = self._get_seg_logits(feat_map)
+    seg_probs = nn.functional.interpolate(
+      seg_probs, size=orig_img_size, mode='bilinear', align_corners=False)
+
+    B, C, H, W = seg_probs.shape
+    num_cls = C
+
+    # Prompt denoising
+    seg_probs = seg_probs.permute(0, 2, 3, 1).reshape(-1, C)
+    max_sim_per_class = torch.max(seg_probs, dim=0).values
+
+    low_conf_classes = torch.argwhere(max_sim_per_class < self.prompt_denoising_thresh)
+    seg_probs[:, low_conf_classes] = 0
+
+    seg_probs = seg_probs.reshape(B, H, W, C).permute(0, 3, 1, 2) # BxCxHxW
+    max_sim_per_pixel, seg_pred = torch.max(seg_probs, dim=1, keepdim=True) # Bx1xHxW
+
+    if self.sam_refinement:
+      seg_pred_ref = list()
+      seg_probs_ref = list()
+      for b in range(B):
+        sam_image, new_h, new_w = self._preprocess_sam(rgb_image[b], target_size=1024)
+        image_features = self._single_inference(sam_image)
+        sam_features = self._get_sam_spatial_features(image_features).float()
+        sam_features = self._interpolate_to_sam_dims(sam_features)
+        sam_features = self.sam_predictor.model.image_encoder.neck(sam_features)
+        self.sam_predictor.features = sam_features
+        self.sam_predictor.is_image_set = True
+        self.sam_predictor.original_size = orig_img_size
+        self.sam_predictor.input_size = (new_h, new_w)
+        
+        refined_masks, scores, refined_logits, prompt_boxes = sam_refinement(
+          orig_img_size, seg_pred[b], seg_probs[b], num_cls, self.sam_predictor,
+          self.coarse_thresh, self.minimal_area,
+          self.sam_mask_coff, self.sam_iou_thresh)
+      
+        seg_pred_ref.append(refined_masks)
+        seg_probs_ref.append(refined_logits)
+      seg_pred = torch.stack(seg_pred_ref, dim=0)
+      seg_probs = torch.stack(seg_probs_ref, dim=0)
+
+    # Set low confidence predictions to the ignore label
+    if ignore_label:
+      seg_pred += 1
+      seg_pred[max_sim_per_pixel < self.prediction_thresh] = 0
+      seg_probs = torch.cat(
+        [torch.zeros_like(seg_probs[:, :1, :, :]), seg_probs], dim=1)
+    if return_preds:
+      return seg_probs, seg_pred
+    else:
+      return seg_probs
+
+  @override
+  def align_spatial_features_with_language(self, features: torch.FloatTensor,
+                                           onehot: bool = True):
+    if self.lang_adaptor is None:
+      raise ValueError("Cannot align to language without a lang model")
+    if not self.return_radio_features or (self.predict and onehot):
+      return features
+    B,C,H,W = features.shape
+    features = features.permute(0, 2, 3, 1).reshape(B, -1, C)
+    with torch.autocast("cuda", dtype=torch.float16, enabled=self.amp):
+      out = self.lang_adaptor.head_mlp(features)
+    return out.permute(0, 2, 1).reshape(B, -1, H, W)
+
+  def _get_sam_spatial_features(self, features: torch.FloatTensor):
+    if self.sam_adaptor is None:
+      raise ValueError("Cannot align to sam without a sam model")
+  
+    B,C,H,W = features.shape
+    features = features.permute(0, 2, 3, 1).reshape(B, -1, C)
+
+    with torch.autocast("cuda", dtype=torch.float16, enabled=self.amp):
+      out = self.sam_adaptor.feat_mlp(features)
+
+    return out.permute(0, 2, 1).reshape(B, -1, H, W)
+  
+  @override
+  def is_compatible_size(self, h: int, w: int):
+    hh, ww = self.get_nearest_size(h, w)
+    return hh == h and ww == w
+
+  @override
+  def get_nearest_size(self, h, w):
+    return self.model.get_nearest_supported_resolution(h, w)
+
+  def _interpolate_to_sam_dims(self, feature_map):
+    feat_h = feature_map.shape[2]
+    feat_w = feature_map.shape[3]
+    feature_map = F.interpolate(
+      feature_map, size=(feat_h, feat_w), mode='bilinear',
+      align_corners=False, antialias=True)
+    feature_map = F.pad(feature_map, (0, 64 - feat_w, 0, 64 - feat_h))
+    return feature_map
+
+  def _preprocess_sam(self, image: torch.Tensor, target_size: int = 1024) -> torch.Tensor:
+    """
+    Preprocess an image tensor for SAM.
+    
+    Args:
+        image (torch.Tensor): Tensor of shape (C, H, W), values in [0, 255] or [0, 1].
+        target_size (int): Size to pad/resize to (default=1024).
+    
+    Returns:
+        torch.Tensor: Preprocessed tensor of shape (1, 3, target_size, target_size).
+    """
+    if image.dim() != 3 or image.shape[0] != 3:
+      raise ValueError("Input must have shape (3, H, W)")
+
+    C, H, W = image.shape
+
+    if H >= W:
+      new_h = target_size
+      new_w = int(W * 64 /H) * 16
+    else:
+      new_w = target_size
+      new_h = int(H * 64 / W) * 16
+
+    image = image.unsqueeze(0)  # add batch dim
+    image = F.interpolate(image, size=(new_h, new_w), mode="bilinear", align_corners=False)
+
+    return image, new_h, new_w
+
+  def _self_correlating_global_aggregation(self,feat_map):
+
+    b,tokens,h,w = feat_map.shape
+    feat_map = feat_map.flatten(2,3).transpose(1,2)
+    sim_tokens = F.normalize(feat_map,dim = -1)
+    sim_matrix = torch.bmm(sim_tokens,sim_tokens.transpose(1,2))
+    sim_matrix = (sim_matrix - torch.mean(sim_matrix)) * self.scga_scaling
+    sim_matrix[sim_matrix < 0] = -torch.inf
+    sim_matrix = F.softmax(sim_matrix,dim=-1)
+    attn_output = torch.bmm(sim_matrix, feat_map)
+    attn_output = attn_output.transpose(1,2).view(b,tokens,h,w)    
+
+    return attn_output
+
+  def _get_seg_logits(self, feat_map):
+
+    image_features = self.align_spatial_features_with_language(feat_map, onehot=False)
+    feat_map = image_features.permute(0, 2, 3, 1)
+
+    B,H,W,C = feat_map.shape
+    feat_map = feat_map.reshape(-1, C)
+
+    cos_sim = compute_cos_sim(
+      self.text_embeds, feat_map,
+      softmax=True)
+
+    N = len(self.text_embeds)
+
+    cos_sim = cos_sim.reshape(B,H,W,N)
+    cos_sim = cos_sim.permute(0,3,1,2)
+    
+    return cos_sim
+
+  def _preprocess_image(self, image, stride=16, slide_crop=336):
+    longer_side = max(image.shape[2:])
+    h, w = image.shape[2:]
+    if h%stride == 0 and w%stride == 0:
+      return image
+    if longer_side % stride != 0:
+      dst_longer = (longer_side // stride + 1) * stride
+    else:
+      dst_longer = longer_side
+    new_h = int(h * dst_longer / longer_side)
+    new_w = int(w * dst_longer / longer_side)
+    if new_h % stride != 0: new_h = (new_h // stride + 1) * stride
+    if new_w % stride != 0: new_w = (new_w // stride + 1) * stride
+    new_h, new_w = max(new_h, slide_crop), max(new_w, slide_crop)
+    image = torch.nn.functional.interpolate(image, (new_h, new_w), mode='bilinear', align_corners=False)
+
+    return image
+
+  def _get_windowed_imgs(self, img, patch_size=16):
+    stride, crop_size = self.slide_stride, self.slide_crop
+    if type(img) == list:
+      img = img[0].unsqueeze(0)
+    if type(stride) == int:
+      stride = (stride, stride)
+    if type(crop_size) == int:
+      crop_size = (crop_size, crop_size)
+
+    h_stride, w_stride = stride
+    h_crop, w_crop = crop_size
+    batch_size, _, h_img, w_img = img.shape
+    h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1
+    w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1
+    crop_imgs, patch_locs = [], []
+    for h_idx in range(h_grids):
+      for w_idx in range(w_grids):
+        y1 = h_idx * h_stride
+        x1 = w_idx * w_stride
+        y2 = min(y1 + h_crop, h_img)
+        x2 = min(x1 + w_crop, w_img)
+        y1 = max(y2 - h_crop, 0)
+        x1 = max(x2 - w_crop, 0)
+        crop_img = img[:, :, y1:y2, x1:x2]
+        assert y1 % patch_size == 0 and x1 % patch_size == 0
+        assert y2 % patch_size == 0 and x2 % patch_size == 0
+        patch_locs.append(torch.tensor([y1//patch_size, x1//patch_size, y2//patch_size, x2//patch_size]))
+        # pad image when (image_size % patch_size != 0)
+        crop_imgs.append(crop_img)
+
+    batched_imgs = torch.cat(crop_imgs, dim=0) # [n_patches, 3, h, w]
+    return batched_imgs, patch_locs, (h_grids,w_grids)
+
+  def _single_inference(self, rgb_image):
+    B, C, H, W = rgb_image.shape
+    H_, W_ = H // self.model.patch_size, W // self.model.patch_size
+    with torch.autocast("cuda", dtype=torch.float16, enabled=self.amp):
+      out = self.model(rgb_image).features
+      if not self.return_radio_features:
+        out = self.lang_adaptor.head_mlp(out)
+    return out.permute(0, 2, 1).reshape(B, -1, H_, W_)
+
+  def _sliding_inference(self, img, stride=224, crop_size=336):
+    """
+    Inference by sliding-window with overlap. If h_crop > h_img or w_crop > w_img,
+    the small patch will be used to decode without padding.
+    """
+
+    if type(img) == list:
+      img = img[0].unsqueeze(0)
+    if type(stride) == int:
+      stride = (stride, stride)
+    if type(crop_size) == int:
+      crop_size = (crop_size, crop_size)
+        
+    img = self._preprocess_image(img, stride[0], crop_size[0])
+    
+    batched_imgs, patch_locs, (h_grids,w_grids) = self._get_windowed_imgs(img, patch_size=16)
+    batch_size = img.shape[0]
+    _, _, h_img, w_img = img.shape
+    
+    image_feats = self._single_inference(batched_imgs)
+    feat_dim = image_feats.shape[1]
+    dtype = image_feats.dtype
+    device = image_feats.device
+    h_feat = math.ceil(h_img / 16)
+    w_feat = math.ceil(w_img / 16)
+    feat_map = torch.zeros((batch_size, feat_dim, h_feat, w_feat),dtype = dtype, device=device)
+    count_mat = torch.zeros((batch_size, 1, h_feat, w_feat),dtype = dtype, device=device)
+    for h_idx in range(h_grids):
+      for w_idx in range(w_grids):
+        coord = patch_locs[h_idx*w_grids+w_idx]
+        img_feat = image_feats[h_idx*w_grids+w_idx]
+        feat_map[:, :,coord[0]:coord[2],coord[1]:coord[3]] += img_feat
+        count_mat[:, :,coord[0]:coord[2],coord[1]:coord[3]] += 1
+    feat_map = feat_map / count_mat  # 1, D, dst_h, dst_w
+
+    return feat_map

radseg/sam_utils.py

@@ -0,0 +1,174 @@
+import torch
+import numpy as np
+from skimage import measure
+from torchvision.ops import masks_to_boxes
+import torch.nn.functional as F
+
+
+def split_connected_regions(segmentation: np.ndarray, seg_logits, split_last=False, minimal_area=0,
+                            coarse_thresh=0.0,) -> dict:
+    """
+    Splits the segmentation mask into separate connected regions for each class.
+
+    Parameters:
+        segmentation (np.ndarray): Semantic segmentation mask of shape [C, H, W].
+        seg_logits (np.ndarray): Segmentation logits of shape [C, H, W].
+        split_last (bool): If True, the last class is not split into connected regions and regarded as the bg.
+        minimal_area (int): Minimal area of a connected region to be considered.
+    Returns:
+        dict: A dictionary where keys are class indices and values are lists of masks.
+    """
+    C, H, W = segmentation.shape
+    regions, boxes, clip_box_score, points = {}, {}, {}, {}
+    for c in range(C):
+        if not split_last and c == C - 1: continue
+        regions[c], points[c], clip_box_score[c],boxes[c] = [], [], [], []
+        class_mask = segmentation[c]
+        class_logit = seg_logits[c]
+        if class_logit.max() < coarse_thresh: continue
+        # Label connected regions
+        labeled_mask = measure.label(class_mask, connectivity=2)  # 8-connectivity
+        num_regions = labeled_mask.max()
+
+        for region_label in range(1, num_regions + 1):
+            region_mask = (labeled_mask == region_label).astype(np.uint8)
+            if region_mask.sum() < minimal_area: continue
+            # Calculate the mean logit score of the region
+            region_score = class_logit[region_mask.astype(bool)].mean()
+            #find the max logit position of the class_logit(H,W) and add it to points[c]
+            region_logit = class_logit * torch.tensor(region_mask).to(class_logit.device)
+            max_pos = torch.argmax(region_logit)
+            max_h, max_w = max_pos // class_logit.size(1), max_pos % class_logit.size(1)
+            points[c].append(torch.tensor([max_w, max_h]))
+            regions[c].append(region_mask)
+            clip_box_score[c].append(region_score)
+
+        if len(regions[c]):
+            regions[c] = torch.from_numpy(np.stack(regions[c]))
+        if regions[c] == []: boxes[c] = regions[c]
+        else: boxes[c] = masks_to_boxes(regions[c])
+    return regions, boxes, clip_box_score, points
+
+def map_refinement_coarse(refined_masks, refined_logits, coarse_boxes, seg_logits, img_shape):
+    refined_whole_logits = torch.zeros([len(coarse_boxes),img_shape[0],img_shape[1]],device=refined_masks.device) # C, H, W
+    cumsum = 0
+    for i in range(len(coarse_boxes)):
+        if coarse_boxes[i] == []: continue
+        tmp_cls_num = coarse_boxes[i].shape[0]
+        tmp_cls_mask = refined_masks[cumsum:cumsum+tmp_cls_num,0].sum(dim=0).bool() #H, W
+        tmp_cls_logit = (refined_logits[cumsum:cumsum+tmp_cls_num,0] * seg_logits[i]).max(dim=0)  #H, W
+        refined_whole_logits[i] = tmp_cls_logit.values * tmp_cls_mask
+        cumsum += tmp_cls_num
+
+    refined_whole_masks = refined_whole_logits.argmax(0, keepdim=True)
+    return refined_whole_masks, refined_whole_logits
+
+def map_failed_regions(refined_masks, refined_logits, failed_regions, seg_logits, segmentation):
+    failed_area = failed_regions.sum(dim=0, keepdim=True)!=0 #1, H, W
+    failed_logit = failed_area * seg_logits #C, H, W
+    refined_logits = torch.where(failed_logit>refined_logits, failed_logit, refined_logits)
+    refined_masks = torch.where(failed_logit.max(dim=0, keepdim=True)[0]>refined_logits.max(dim=0, keepdim=True)[0], segmentation, refined_masks)
+    return refined_masks, refined_logits
+
+def sam_refinement(img_size, segmentations, seg_logits, num_classes, predictor, coarse_thresh=0.0, minimal_area=0,
+                   sam_mask_coff=0.005,sam_iou_thresh=0.9):
+    '''
+    Args:
+        img: str or np.ndarray
+        segmentations: torch.Tensor, shape=[C, H, W]
+        seg_logits: torch.Tensor, shape=[C, H, W]
+        num_classes: int
+        predictor: Predictor
+        coarse_thresh: float
+        minimal_area: int
+        sam_mask_coff: float
+        sam_iou_thresh: float
+    '''
+    #downsample the segmentation and seg_logits to accelerate split_connected_regions
+    down_ratio = 2
+    dh, dw = seg_logits.shape[-2] // down_ratio, seg_logits.shape[-1] // down_ratio
+    if down_ratio != 1:
+        segmentations_down = F.interpolate(segmentations.unsqueeze(0).half(), [dh, dw], mode='nearest').squeeze(0).long()
+        seg_logits_down = F.interpolate(seg_logits.unsqueeze(0), [dh, dw], mode='bilinear', align_corners=False).squeeze(0)
+    else:
+        segmentations_down = segmentations
+        seg_logits_down = seg_logits
+
+    if coarse_thresh > 0:
+        segmentations_down[seg_logits_down.max(0, keepdim=True)[0] < coarse_thresh] = num_classes
+        cls_pred = F.one_hot(segmentations_down.squeeze(0).long(), num_classes=num_classes+1).permute(2, 0, 1).float()  # [C, H, W]
+        split_last = False
+    else:
+        cls_pred = F.one_hot(segmentations_down.squeeze(0).long(), num_classes=num_classes).permute(2, 0, 1).float()
+        split_last = True
+
+    coarse_regions, coarse_boxes, coarse_boxes_score, coarse_points = split_connected_regions(cls_pred.cpu().numpy(), seg_logits_down,
+                                                           split_last=split_last, minimal_area=minimal_area,coarse_thresh=coarse_thresh)
+    combined_boxes,combined_points = [], []
+    cls_wise_masks, cls_wise_scores, cls_wise_logits = [], [], []
+    failed_regions = []
+    predictor.features = predictor.features.float()
+    for i in range(num_classes):
+        tmp_r = coarse_regions.get(i, [])
+        tmp_b = coarse_boxes.get(i, [])
+        tmp_p = coarse_points.get(i, [])
+
+        if tmp_r == []: continue
+        cls_wise_r = seg_logits_down[i].unsqueeze(0)
+        #resize to longer side = 256
+        new_h, new_w = predictor.transform.get_preprocess_shape(img_size[0], img_size[1], 256)
+        cls_wise_r = F.interpolate(cls_wise_r.unsqueeze(0), [new_h, new_w], mode="bilinear", align_corners=False).squeeze(0)
+        #pad to 256x256
+        cls_wise_r = F.pad(cls_wise_r, (0, 256-new_w, 0, 256-new_h), mode='constant', value=0)
+        cls_wise_r = torch.where(cls_wise_r > coarse_thresh, sam_mask_coff*cls_wise_r, torch.tensor(0.0).type_as(segmentations)).float()
+        cls_wise_b = predictor.transform.apply_boxes_torch(tmp_b.type_as(segmentations), [dh, dw]).float()
+        cls_wise_p = torch.stack(tmp_p).type_as(segmentations)
+        cls_wise_p = predictor.transform.apply_coords_torch(cls_wise_p, [dh, dw]).unsqueeze(1)
+        cls_wise_mask, cls_wise_score, cls_wise_logit = predictor.predict_torch(
+            point_coords=cls_wise_p.float(),  # input_point,
+            point_labels=torch.ones([cls_wise_p.shape[0], 1]).type_as(segmentations).float(),  # input_label,
+            boxes=cls_wise_b,
+            mask_input=cls_wise_r,
+            multimask_output=False,
+        )
+
+        valid_masks = (cls_wise_score > sam_iou_thresh).squeeze(-1)
+        if valid_masks.sum() != tmp_r.shape[0]:
+            failed_regions.append(tmp_r[~valid_masks.cpu()])
+        if valid_masks.sum() == 0:
+            coarse_boxes[i] = []
+            continue
+        cls_wise_masks.append(cls_wise_mask[valid_masks])
+        cls_wise_scores.append(cls_wise_score[valid_masks])
+        cls_wise_logits.append(cls_wise_logit[valid_masks])
+
+        tmp_p = torch.stack(tmp_p)
+        tmp_b = tmp_b[valid_masks.cpu()]
+        tmp_p = tmp_p[valid_masks.cpu()]
+        coarse_boxes[i] = coarse_boxes[i][valid_masks.cpu()]
+
+        combined_boxes.append(tmp_b)
+        combined_points.append(tmp_p)
+
+    if len(combined_boxes) == 0:
+        return segmentations, None, seg_logits, combined_boxes
+
+    combined_boxes = torch.cat(combined_boxes, dim=0).to(predictor.device)
+    masks, scores, logits = torch.cat(cls_wise_masks, dim=0), torch.cat(cls_wise_scores, dim=0), torch.cat(cls_wise_logits, dim=0)
+    logits = predictor.model.postprocess_masks(logits, predictor.input_size, predictor.original_size)
+
+    combined_boxes = combined_boxes * down_ratio
+    if logits.shape[-1] != img_size[1] or logits.shape[-2] != img_size[0]:
+        logits = F.interpolate(logits, img_size[:2], mode='bilinear', align_corners=False)
+        masks = logits>predictor.model.mask_threshold
+
+    refined_masks, refined_logits = map_refinement_coarse(masks, logits.sigmoid(), coarse_boxes, seg_logits, img_size,)
+
+    refined_masks = torch.where(refined_logits.sum(dim=0,keepdim=True) == 0, segmentations, refined_masks)
+    if len(failed_regions):
+        failed_regions = torch.cat(failed_regions, dim=0).to(predictor.device)
+        if down_ratio != 1:
+            failed_regions = F.interpolate(failed_regions.unsqueeze(0).float(), img_size[:2], mode='nearest').squeeze(0).type_as(segmentations)
+        refined_masks, refined_logits = map_failed_regions(refined_masks, refined_logits, failed_regions, seg_logits, segmentations)
+
+    return refined_masks, scores, refined_logits,combined_boxes

radseg_demo.py

@@ -0,0 +1,285 @@
+import sys
+import os
+import logging
+import base64
+from io import BytesIO
+import colorsys
+from functools import partial
+
+import gradio as gr
+import spaces
+from PIL import Image
+import numpy as np
+import torch
+import torch.nn.functional as F
+from matplotlib import cm
+
+# Add current directory to sys.path to find radseg module
+sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
+from radseg.radseg import RADSegEncoder
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# Global state to keep track of prompts and colors (similar to inspiration app)
+prompt_list = []
+color_list = []
+
+def apply_colormap(image: np.ndarray, cmap_name='viridis') -> np.ndarray:
+    """Apply a colormap to a grayscale image and return an RGB uint8 image."""
+    if image.dtype != np.float16 and image.dtype != np.float32 and image.dtype != np.float64:
+        image = image.astype(np.float32) / 255.0
+    image = np.clip(image, 0, 1)
+    cmap = cm.get_cmap(cmap_name)
+    colored = cmap(image)[:, :, :3]  # Drop alpha channel
+    return (colored * 255).astype(np.uint8)
+
+def numpy_to_base64(img_array):
+    """Convert a NumPy array image to base64 string."""
+    if img_array.dtype != np.uint8:
+        img_array = (img_array * 255).astype(np.uint8)
+    pil_img = Image.fromarray(img_array)
+    buffered = BytesIO()
+    pil_img.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode()
+
+def make_grid_output(images, labels):
+    html = """
+    <div style='display: grid; grid-template-columns: repeat(3, 1fr); gap: 15px;'>
+    """
+    for img_array, label in zip(images, labels):
+        img_str = numpy_to_base64(img_array)
+        html += f"""
+        <div style='text-align: center;'>
+            <div style='font-weight: bold; margin-bottom: 5px;'>{label}</div>
+            <img src='data:image/png;base64,{img_str}' style='width: 100%; height: auto; border: 1px solid #ccc;' />
+        </div>
+        """
+    html += "</div>"
+    return html
+
+def generate_distinct_color(index):
+    """Generate visually distinct colors using HSV color space."""
+    hue = (index * 0.61803398875) % 1  # golden ratio
+    r, g, b = colorsys.hsv_to_rgb(hue, 0.5, 0.95)
+    return '#%02x%02x%02x' % (int(r*255), int(g*255), int(b*255))
+
+def add_prompt(prompts):
+    for prompt in prompts.split("\n"):
+        prompt = prompt.strip()
+        if not prompt or prompt in prompt_list:
+            continue
+        color = generate_distinct_color(len(prompt_list))
+        prompt_list.append(prompt)
+        color_list.append(color)
+
+    colored_prompts = [
+        f"<span style='background-color:{color}; color:#FFFFFF; padding: 2px 5px; border-radius: 3px; margin: 2px; display: inline-block;'>{p}</span>" 
+        for p, color in zip(prompt_list, color_list)
+    ]
+    return gr.update(value=""), gr.update(value=" ".join(colored_prompts))
+
+def clear_prompts():
+    prompt_list.clear()
+    color_list.clear()
+    return gr.update(value=""), gr.update(value="")
+
+def on_page_load():
+    prompt_list.clear()
+    color_list.clear()
+    return gr.update(value="")
+
+use_templates = True
+model_version = "c-radio_v3-b"
+lang_model = "siglip2"
+
+# Load model lazily
+encoder = None
+def get_encoder(scra_scaling, scga_scaling, slide_crop, slide_stride):
+    global encoder
+    if encoder is None:
+        logger.info(f"Loading encoder: {model_version} with {lang_model}")
+        try:
+            # We don't pass classes yet as we will handle that in inference
+            # or by manually encoding prompts.
+            # To avoid the "Must pass list of classes when predict is True" error,
+            # we'll use predict=False for the base encoder and manage predictions ourselves.
+            encoder = RADSegEncoder(
+                    model_version=model_version,
+                    lang_model=lang_model,
+                    scra_scaling=scra_scaling,
+                    scga_scaling=scga_scaling,
+                    slide_crop=slide_crop,
+                    slide_stride=slide_stride,
+                    sam_refinement=False,
+                    predict=False,
+                    device="cuda"
+            )
+        except Exception as e:
+            logger.error(f"Error loading encoder: {e}")
+            raise gr.Error(f"Failed to load encoder: {e}")
+    else:
+        encoder.slide_stride = slide_stride
+        encoder.slide_crop = slide_crop
+        encoder.scga_scaling = scga_scaling
+        encoder.model.model.blocks[-1].attn.scra_scaling = scra_scaling
+
+    return encoder
+
+@spaces.GPU
+@torch.inference_mode()
+def process_all(input_image, scra_scaling, scga_scaling, use_sliding_window, window_size, window_stride, softmax, resolution):
+    global encoder
+    N = len(prompt_list)
+    if N == 0:
+        raise gr.Error("You must add some prompts", duration=5)
+    
+    yield "Initializing encoder..."
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    slide_crop = window_size if use_sliding_window else 0
+    slide_stride = window_stride if use_sliding_window else 224
+    encoder = get_encoder(scra_scaling, scga_scaling, slide_crop, slide_stride)
+
+    yield "Processing image..."
+    # Convert numpy to torch
+    orig_h, orig_w = input_image.shape[:2]
+    
+    # Resize to requested resolution (maintaining aspect ratio)
+    if max(orig_h, orig_w) > resolution:
+        scale = resolution / max(orig_h, orig_w)
+        new_h, new_w = int(orig_h * scale), int(orig_w * scale)
+        # Ensure divisible by 16 for VIT/RADIO if needed, though slide crop might handle it
+        new_h = (new_h // 16) * 16
+        new_w = (new_w // 16) * 16
+        input_image_resized = np.array(Image.fromarray(input_image).resize((new_w, new_h), resample=Image.BILINEAR))
+    else:
+        input_image_resized = input_image
+
+    image_size = input_image_resized.shape[:2]
+    tensor_image = torch.from_numpy(input_image_resized).permute(2, 0, 1).to(device).float() / 255.0
+    
+    # Run RADSeg processing
+    # Note: RADSegEncoder handle sliding window if slide_crop is set.
+    # We can temporarily override resolution if needed but let's use the encoder's defaults.
+    
+    # Manual similarity calculation to match inspiration app's heatmap per prompt
+    yield "Computing features..."
+    feat_map = encoder.encode_image_to_feat_map(tensor_image.unsqueeze(0), orig_img_size=image_size)
+    
+    yield "Aligning features..."
+    aligned_feats = encoder.align_spatial_features_with_language(feat_map, onehot=False)
+    
+    yield "Encoding prompts..."
+    if use_templates:
+        prompt_embeds = encoder.encode_labels(prompt_list)
+    else:
+        prompt_embeds = encoder.encode_prompts(prompt_list)
+        
+    yield "Computing similarity..."
+    B, C, H, W = aligned_feats.shape
+    aligned_feats_flat = aligned_feats.permute(0, 2, 3, 1).reshape(-1, C)
+    
+    # Compute cosine similarity
+    # similarity: [N_prompts, H*W]
+    vec1 = prompt_embeds / prompt_embeds.norm(dim=-1, keepdim=True)
+    vec2 = aligned_feats_flat / aligned_feats_flat.norm(dim=-1, keepdim=True)
+    sim = (vec1 @ vec2.t()) # [N, H*W]
+    
+    if softmax:
+        sim = torch.softmax(100 * sim, dim=0)
+    else:
+        sim /= torch.max(sim) # For visualization
+        
+    sim = sim.reshape(N, H, W)
+    
+    # Resize heatmaps to original resolution for display
+    sim_resized = F.interpolate(sim.unsqueeze(0), size=image_size, mode='bilinear', align_corners=False).squeeze(0)
+    
+    yield "Generating outputs..."
+    heatmaps = [apply_colormap(sim_resized[i].cpu().numpy()) for i in range(N)]
+    yield make_grid_output(heatmaps, prompt_list)
+    torch.cuda.empty_cache()
+
+def main():
+    with gr.Blocks(title="RADSeg 2D Demo") as demo:
+        gr.Markdown(
+            """
+            # RADSeg: Zero-Shot Open-Vocabulary Segmentation
+            ### [Paper](https://arxiv.org/abs/2511.19704) | [Project Page](https://radseg-ovss.github.io/)
+            Test RADSeg's zero-shot open-vocabulary segmentation capabilities on any image using text prompts !
+            """
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image = gr.Image(label="Input Image", type="numpy")
+
+                with gr.Accordion("Model & Inference Settings", open=False):
+                    with gr.Group():
+                        gr.Markdown("### 1. Model Scaling")
+                        scra_scaling = gr.Slider(1.0, 20.0, 10.0, step=1.0, label="SCRA Scaling", info="Self-Correlating Recursive Attention")
+                        scga_scaling = gr.Slider(1.0, 20.0, 10.0, step=1.0, label="SCGA Scaling", info="Self-Correlating Global Aggregation")
+
+                    with gr.Group():
+                        gr.Markdown("### 2. Sliding Window")
+                        use_sliding_window = gr.Checkbox(label="Enable Sliding Window", value=True)
+                        with gr.Row():
+                            window_size = gr.Slider(224, 1024, 336, step=16, label="Window Size")
+                            window_stride = gr.Slider(112, 512, 224, step=16, label="Window Stride")
+
+                    with gr.Row():
+                        softmax = gr.Checkbox(label="Use softmax", value=True)
+                        res_slider = gr.Slider(224, 1024, 512, step=32, label="Resolution (Max Side)")
+
+                with gr.Group():
+                    gr.Markdown("### 3. Text Prompts")
+                    with gr.Row():
+                        prompt = gr.Textbox(
+                            label="Prompt", 
+                            placeholder="Type a class (e.g. 'car') and press enter",
+                            scale=4,
+                            show_label=False
+                        )
+                        add_button = gr.Button("+", scale=1)
+
+                    prompt_display = gr.HTML()
+
+                    with gr.Row():
+                        clear_button = gr.Button("Clear Prompts")
+                        run_button = gr.Button("Run Segmentation", variant="primary")
+
+            with gr.Column(scale=2):
+                output_html = gr.HTML(label="Segmentation Heatmaps")
+
+        # Event handlers
+        prompt.submit(add_prompt, inputs=prompt, outputs=[prompt, prompt_display])
+        add_button.click(add_prompt, inputs=prompt, outputs=[prompt, prompt_display])
+        clear_button.click(clear_prompts, outputs=[prompt, prompt_display])
+
+        gr.Examples(
+            examples=[
+                ["assets/example1.jpg", 10, 10, True, 336, 224, True, "Pothole\nRoad\nSky\nCar\nWater", 768],
+                ["assets/example2.jpg", 10, 10, True, 336, 224, True, "Person\nShoes\nGrey Jacket\nRed overalls\nRoad\nCrosswalk\nCar", 768],
+                ["assets/example3.jpg", 10, 10, True, 336, 224, True, "Paved ground\nFlood lights\nRed Container\nBuilding\nTanker\nSky\nClouds\nTreeline", 1024]
+            ],
+            inputs=[input_image, scra_scaling, scga_scaling, use_sliding_window, window_size, window_stride, softmax, prompt, res_slider],
+        )
+
+        run_button.click(
+            fn=process_all,
+            inputs=[
+                input_image,
+                scra_scaling, scga_scaling,
+                use_sliding_window, window_size, window_stride,
+                softmax, res_slider
+            ],
+            outputs=output_html
+        )
+        
+        demo.load(on_page_load, outputs=prompt_display)
+
+    demo.queue().launch()
+
+if __name__ == "__main__":
+    main()
+

requirements.txt

@@ -0,0 +1,23 @@
+torch==2.4.0
+torchvision==0.19.0
+
+timm==1.0.19
+segment-anything
+opencv-python-headless
+pillow
+numpy
+matplotlib
+ftfy
+regex
+pyyaml
+tqdm
+huggingface-hub
+transformers
+open-clip-torch
+scipy
+scikit-image
+addict
+yapf
+typing_extensions
+einops
+gradio