update MobileSAM

app.py

@@ -1,6 +1,7 @@
 from typing import Optional, Tuple
 from einops import rearrange
 import torch
+import torch.nn.functional as F
 from PIL import Image
 import torchvision.transforms as transforms
 from torch import nn
@@ -12,8 +13,168 @@ import gradio as gr
 
 use_cuda = torch.cuda.is_available()
 
+# use_cuda = False
+
 print("CUDA is available:", use_cuda)
 
+class MobileSAM(nn.Module):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        from mobile_sam import sam_model_registry
+
+        url = 'https://raw.githubusercontent.com/ChaoningZhang/MobileSAM/master/weights/mobile_sam.pt'
+        model_type = "vit_t"
+        sam_checkpoint = "mobile_sam.pt"
+        if not os.path.exists(sam_checkpoint):
+            import requests
+            r = requests.get(url)
+            with open(sam_checkpoint, 'wb') as f:
+                f.write(r.content)
+
+        device = 'cuda' if use_cuda else 'cpu'
+
+        mobile_sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+
+        def new_forward_fn(self, x):
+            shortcut = x
+
+            x = self.conv1(x)
+            x = self.act1(x)
+
+            x = self.conv2(x)
+            x = self.act2(x)
+            
+            self.attn_output = rearrange(x.clone(), "b c h w -> b h w c")
+
+            x = self.conv3(x)
+            
+            self.mlp_output = rearrange(x.clone(), "b c h w -> b h w c")
+
+            x = self.drop_path(x)
+
+            x += shortcut
+            x = self.act3(x)
+            
+            self.block_output = rearrange(x.clone(), "b c h w -> b h w c")
+
+            return x
+
+        setattr(mobile_sam.image_encoder.layers[0].blocks[0].__class__, "forward", new_forward_fn)
+
+        def new_forward_fn2(self, x):
+            H, W = self.input_resolution
+            B, L, C = x.shape
+            assert L == H * W, "input feature has wrong size"
+            res_x = x
+            if H == self.window_size and W == self.window_size:
+                x = self.attn(x)
+            else:
+                x = x.view(B, H, W, C)
+                pad_b = (self.window_size - H %
+                            self.window_size) % self.window_size
+                pad_r = (self.window_size - W %
+                            self.window_size) % self.window_size
+                padding = pad_b > 0 or pad_r > 0
+
+                if padding:
+                    x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
+
+                pH, pW = H + pad_b, W + pad_r
+                nH = pH // self.window_size
+                nW = pW // self.window_size
+                # window partition
+                x = x.view(B, nH, self.window_size, nW, self.window_size, C).transpose(2, 3).reshape(
+                    B * nH * nW, self.window_size * self.window_size, C)
+                x = self.attn(x)
+                # window reverse
+                x = x.view(B, nH, nW, self.window_size, self.window_size,
+                            C).transpose(2, 3).reshape(B, pH, pW, C)
+
+                if padding:
+                    x = x[:, :H, :W].contiguous()
+
+                x = x.view(B, L, C)
+
+            hw = np.sqrt(x.shape[1]).astype(int)
+            self.attn_output = rearrange(x.clone(), "b (h w) c -> b h w c", h=hw)
+            
+            x = res_x + self.drop_path(x)
+
+            x = x.transpose(1, 2).reshape(B, C, H, W)
+            x = self.local_conv(x)
+            x = x.view(B, C, L).transpose(1, 2)
+
+            mlp_output = self.mlp(x)
+            self.mlp_output = rearrange(mlp_output.clone(), "b (h w) c -> b h w c", h=hw)
+            
+            x = x + self.drop_path(mlp_output)
+            self.block_output = rearrange(x.clone(), "b (h w) c -> b h w c", h=hw)
+            return x
+
+        setattr(mobile_sam.image_encoder.layers[1].blocks[0].__class__, "forward", new_forward_fn2)
+
+
+        mobile_sam.to(device=device)
+        mobile_sam.eval()
+        self.image_encoder = mobile_sam.image_encoder
+        
+    
+    @torch.no_grad()
+    def forward(self, x):
+        with torch.no_grad():
+            x = torch.nn.functional.interpolate(x, size=(1024, 1024), mode="bilinear")
+        out = self.image_encoder(x)
+
+        attn_outputs, mlp_outputs, block_outputs = [], [], []
+        for i_layer in range(len(self.image_encoder.layers)):
+            for i_block in range(len(self.image_encoder.layers[i_layer].blocks)):
+                blk = self.image_encoder.layers[i_layer].blocks[i_block]
+                attn_outputs.append(blk.attn_output)
+                mlp_outputs.append(blk.mlp_output)
+                block_outputs.append(blk.block_output)
+        return attn_outputs, mlp_outputs, block_outputs
+
+
+def image_mobilesam_feature(
+    images,
+    resolution=(1024, 1024),
+    node_type="block",
+    layer=-1,
+):
+
+    transform = transforms.Compose(
+        [
+            transforms.Resize(resolution),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+    
+
+    feat_extractor = MobileSAM()
+
+    # attn_outputs, mlp_outputs, block_outputs = [], [], []
+    outputs = []
+    for i, image in enumerate(images):
+        torch_image = transform(image)
+        if use_cuda:
+            torch_image = torch_image.cuda()
+        attn_output, mlp_output, block_output = feat_extractor(
+            torch_image.unsqueeze(0)
+        )
+        out_dict = {
+            "attn": attn_output,
+            "mlp": mlp_output,
+            "block": block_output,
+        }
+        out = out_dict[node_type]
+        out = out[layer]
+        outputs.append(out.cpu())
+    outputs = torch.cat(outputs, dim=0)
+    return outputs
+
+
 
 class SAM(torch.nn.Module):
     def __init__(self, checkpoint="/data/sam_model/sam_vit_b_01ec64.pth", **kwargs):
@@ -307,6 +468,8 @@ def image_clip_feature(
 def extract_features(images, model_name="sam", node_type="block", layer=-1):
     if model_name == "SAM(sam_vit_b)":
         return image_sam_feature(images, node_type=node_type, layer=layer)
+    elif model_name == 'MobileSAM':
+        return image_mobilesam_feature(images, node_type=node_type, layer=layer)
     elif model_name == "DiNO(dinov2_vitb14_reg)":
         return image_dino_feature(images, node_type=node_type, layer=layer)
     elif model_name == "CLIP(openai/clip-vit-base-patch16)":
@@ -346,6 +509,9 @@ def compute_ncut(
     )
     print(f"t-SNE time: {time.time() - start:.2f}s")
     
+    # print("input shape:", features.shape)
+    # print("output shape:", rgb.shape)
+    
     rgb = rgb.reshape(features.shape[:3] + (3,))
     return rgb
 
@@ -413,7 +579,7 @@ demo = gr.Interface(
     main_fn,
     [
         gr.Gallery(value=default_images, label="Select images", show_label=False, elem_id="images", columns=[3], rows=[1], object_fit="contain", height="auto", type="pil"),
-        gr.Dropdown(["SAM(sam_vit_b)", "DiNO(dinov2_vitb14_reg)", "CLIP(openai/clip-vit-base-patch16)"], label="Model", value="SAM(sam_vit_b)", elem_id="model_name"),
+        gr.Dropdown(["MobileSAM", "SAM(sam_vit_b)", "DiNO(dinov2_vitb14_reg)", "CLIP(openai/clip-vit-base-patch16)"], label="Model", value="MobileSAM", elem_id="model_name"),
         gr.Dropdown(["attn", "mlp", "block"], label="Node type", value="block", elem_id="node_type", info="attn: attention output, mlp: mlp output, block: sum of residual stream"),
         gr.Slider(0, 11, step=1, label="Layer", value=11, elem_id="layer", info="which layer of the image backbone features"),
         gr.Slider(1, 1000, step=1, label="Number of eigenvectors", value=100, elem_id="num_eig", info='increase for more object parts, decrease for whole object'),

requirements.txt

@@ -1,6 +1,7 @@
 ncut-pytorch
 transformers
 segment-anything @ git+https://github.com/facebookresearch/segment-anything.git
+mobile-sam @ git+https://github.com/ChaoningZhang/MobileSAM.git
 --extra-index-url https://download.pytorch.org/whl/cpu
 torch 
 torchvision