update sam2

app.py

@@ -472,13 +472,29 @@ def make_output_images_section():
 
 def make_parameters_section():
     gr.Markdown('### Parameters')
-    model_dropdown = gr.Dropdown(["SAM(sam_vit_b)", "MobileSAM", "DiNO(dinov2_vitb14_reg)", "CLIP(openai/clip-vit-base-patch16)", "MAE(vit_base)"], label="Backbone", value="SAM(sam_vit_b)", elem_id="model_name")
+    model_names = [
+        "SAM(sam_vit_b)",
+        "MobileSAM",
+        "DiNO(dinov2_vitb14_reg)",
+        "CLIP(openai/clip-vit-base-patch16)",
+        "MAE(vit_base)",
+        "SAM2(sam2_hiera_b+)",
+        "SAM2(sam2_hiera_t)",
+    ]
+    model_dropdown = gr.Dropdown(model_names, label="Backbone", value="SAM(sam_vit_b)", elem_id="model_name")
     layer_slider = gr.Slider(0, 11, step=1, label="Backbone: Layer index", value=11, elem_id="layer")
     node_type_dropdown = gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?")
     num_eig_slider = gr.Slider(1, 1000, step=1, label="NCUT: Number of eigenvectors", value=100, elem_id="num_eig", info='increase for more clusters')
-    affinity_focal_gamma_slider = gr.Slider(0.01, 1, step=0.01, label="NCUT: Affinity focal gamma", value=0.5, elem_id="affinity_focal_gamma", info="decrease for shaper segmentation")
 
+    def change_layer_slider(model_name):
+        if model_name == "SAM2(sam2_hiera_b+)":
+            return gr.Slider(0, 23, step=1, label="Backbone: Layer index", value=23, elem_id="layer", visible=True)
+        else:
+            return gr.Slider(0, 11, step=1, label="Backbone: Layer index", value=11, elem_id="layer", visible=True)
+    model_dropdown.change(fn=change_layer_slider, inputs=model_dropdown, outputs=layer_slider)
+    
     with gr.Accordion("➡️ Click to expand: more parameters", open=False):
+        affinity_focal_gamma_slider = gr.Slider(0.01, 1, step=0.01, label="NCUT: Affinity focal gamma", value=0.5, elem_id="affinity_focal_gamma", info="decrease for shaper segmentation")
         num_sample_ncut_slider = gr.Slider(100, 50000, step=100, label="NCUT: num_sample", value=10000, elem_id="num_sample_ncut", info="Nyström approximation")
         sampling_method_dropdown = gr.Dropdown(["fps", "random"], label="NCUT: Sampling method", value="fps", elem_id="sampling_method", info="Nyström approximation")
         knn_ncut_slider = gr.Slider(1, 100, step=1, label="NCUT: KNN", value=10, elem_id="knn_ncut", info="Nyström approximation")

backbone.py

@@ -1,5 +1,6 @@
 from typing import Optional, Tuple
 from einops import rearrange
+import requests
 import torch
 import torch.nn.functional as F
 import timm
@@ -228,6 +229,98 @@ class SAM(torch.nn.Module):
 MODEL_DICT["SAM(sam_vit_b)"] = SAM()
 
 
+class SAM2(nn.Module):
+
+    def __init__(self, model_cfg='sam2_hiera_b+',):
+        super().__init__()
+
+        try:
+            from sam2.build_sam import build_sam2
+        except ImportError:
+            print("Please install segment_anything_2 from https://github.com/facebookresearch/segment-anything-2.git")
+            return
+        
+        config_dict = {
+            'sam2_hiera_large': ("sam2_hiera_large.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt"),
+            'sam2_hiera_b+': ("sam2_hiera_base_plus.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt"),
+            'sam2_hiera_s': ("sam2_hiera_small.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt"),
+            'sam2_hiera_t': ("sam2_hiera_tiny.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt"),
+        }
+        filename, url = config_dict[model_cfg]
+        if not os.path.exists(filename):
+            print(f"Downloading {url}")
+            r = requests.get(url)
+            with open(filename, 'wb') as f:
+                f.write(r.content)
+        sam2_checkpoint = filename
+        
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        sam2_model = build_sam2(model_cfg, sam2_checkpoint, device=device)
+
+        image_encoder = sam2_model.image_encoder
+        image_encoder.eval()
+        
+        from sam2.modeling.backbones.hieradet import do_pool
+        from sam2.modeling.backbones.utils import window_partition, window_unpartition
+        def new_forward(self, x: torch.Tensor) -> torch.Tensor:
+            shortcut = x  # B, H, W, C
+            x = self.norm1(x)
+
+            # Skip connection
+            if self.dim != self.dim_out:
+                shortcut = do_pool(self.proj(x), self.pool)
+
+            # Window partition
+            window_size = self.window_size
+            if window_size > 0:
+                H, W = x.shape[1], x.shape[2]
+                x, pad_hw = window_partition(x, window_size)
+
+            # Window Attention + Q Pooling (if stage change)
+            x = self.attn(x)
+            if self.q_stride:
+                # Shapes have changed due to Q pooling
+                window_size = self.window_size // self.q_stride[0]
+                H, W = shortcut.shape[1:3]
+
+                pad_h = (window_size - H % window_size) % window_size
+                pad_w = (window_size - W % window_size) % window_size
+                pad_hw = (H + pad_h, W + pad_w)
+
+            # Reverse window partition
+            if self.window_size > 0:
+                x = window_unpartition(x, window_size, pad_hw, (H, W))
+
+            self.attn_output = x.clone()
+            
+            x = shortcut + self.drop_path(x)
+            # MLP
+            mlp_out = self.mlp(self.norm2(x))
+            self.mlp_output = mlp_out.clone()
+            x = x + self.drop_path(mlp_out)
+            self.block_output = x.clone()
+            return x        
+        
+        setattr(image_encoder.trunk.blocks[0].__class__, 'forward', new_forward)
+        
+        self.image_encoder = image_encoder
+        
+        
+        
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        output = self.image_encoder(x)
+        attn_outputs, mlp_outputs, block_outputs = [], [], []
+        for block in self.image_encoder.trunk.blocks:
+            attn_outputs.append(block.attn_output)
+            mlp_outputs.append(block.mlp_output)
+            block_outputs.append(block.block_output)
+        return attn_outputs, mlp_outputs, block_outputs
+
+
+MODEL_DICT["SAM2(sam2_hiera_b+)"] = SAM2(model_cfg='sam2_hiera_b+')
+MODEL_DICT["SAM2(sam2_hiera_t)"] = SAM2(model_cfg='sam2_hiera_t')
+
 class DiNOv2(torch.nn.Module):
     def __init__(self, ver="dinov2_vitb14_reg"):
         super().__init__()
@@ -282,6 +375,16 @@ class DiNOv2(torch.nn.Module):
 
 MODEL_DICT["DiNO(dinov2_vitb14_reg)"] = DiNOv2()
 
+def resample_position_embeddings(embeddings, h, w):
+    cls_embeddings = embeddings[0]
+    patch_embeddings = embeddings[1:]  # [14*14, 768]
+    hw = np.sqrt(patch_embeddings.shape[0]).astype(int)
+    patch_embeddings = rearrange(patch_embeddings, "(h w) c -> c h w", h=hw)
+    patch_embeddings = F.interpolate(patch_embeddings.unsqueeze(0), size=(h, w), mode="nearest").squeeze(0)
+    patch_embeddings = rearrange(patch_embeddings, "c h w -> (h w) c")
+    embeddings = torch.cat([cls_embeddings.unsqueeze(0), patch_embeddings], dim=0)
+    return embeddings
+
 class CLIP(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -289,17 +392,12 @@ class CLIP(torch.nn.Module):
         from transformers import CLIPProcessor, CLIPModel
 
         model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
-        
-        # resample the patch embeddings to 64x64, take 1024x1024 input
+                
+        # resample the patch embeddings to 56x56, take 896x896 input
         embeddings = model.vision_model.embeddings.position_embedding.weight
-        cls_embeddings = embeddings[0]
-        patch_embeddings = embeddings[1:]  # [14*14, 768]
-        patch_embeddings = rearrange(patch_embeddings, "(h w) c -> c h w", h=14)
-        patch_embeddings = F.interpolate(patch_embeddings.unsqueeze(0), size=(64, 64), mode="bilinear", align_corners=False).squeeze(0)
-        patch_embeddings = rearrange(patch_embeddings, "c h w -> (h w) c")
-        embeddings = torch.cat([cls_embeddings.unsqueeze(0), patch_embeddings], dim=0)
+        embeddings = resample_position_embeddings(embeddings, 56, 56)
         model.vision_model.embeddings.position_embedding.weight = nn.Parameter(embeddings)
-        model.vision_model.embeddings.position_ids = torch.arange(0, 1+64*64)
+        model.vision_model.embeddings.position_ids = torch.arange(0, 1+56*56)
         
         # processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
         self.model = model.eval()
@@ -392,17 +490,12 @@ class MAE(timm.models.vision_transformer.VisionTransformer):
         msg = self.load_state_dict(checkpoint_model, strict=False)
         print(msg)
         
-        # resample the patch embeddings to 64x64, take 1024x1024 input
+        # resample the patch embeddings to 56x56, take 896x896 input
         pos_embed = self.pos_embed[0]
-        cls_embeddings = pos_embed[0]
-        patch_embeddings = pos_embed[1:]  # [14*14, 768]
-        patch_embeddings = rearrange(patch_embeddings, "(h w) c -> c h w", h=14)
-        patch_embeddings = F.interpolate(patch_embeddings.unsqueeze(0), size=(64, 64), mode="bilinear", align_corners=False).squeeze(0)
-        patch_embeddings = rearrange(patch_embeddings, "c h w -> (h w) c")
-        pos_embed = torch.cat([cls_embeddings.unsqueeze(0), patch_embeddings], dim=0)
+        pos_embed = resample_position_embeddings(pos_embed, 56, 56)
         self.pos_embed = nn.Parameter(pos_embed.unsqueeze(0))
-        self.img_size = (1024, 1024)
-        self.patch_embed.img_size = (1024, 1024)
+        self.img_size = (896, 896)
+        self.patch_embed.img_size = (896, 896)
 
         self.requires_grad_(False)
         self.eval()
@@ -441,6 +534,8 @@ def extract_features(images, model_name, node_type, layer):
     resolution = (1024, 1024)
     resolution_dict = {
         "DiNO(dinov2_vitb14_reg)": (896, 896),
+        'CLIP(openai/clip-vit-base-patch16)': (896, 896),
+        'MAE(vit_base)': (896, 896),
     }
     if model_name in resolution_dict:
         resolution = resolution_dict[model_name]

requirements.txt

@@ -6,5 +6,6 @@ decord
 transformers
 datasets
 segment-anything @ git+https://github.com/facebookresearch/segment-anything.git
+segment-anything-2 @ git+https://github.com/facebookresearch/segment-anything-2.git
 mobile-sam @ git+https://github.com/ChaoningZhang/MobileSAM.git
 timm