from typing import Optional, Tuple from einops import rearrange import torch import torch.nn.functional as F from PIL import Image from torch import nn import numpy as np import os import time import gradio as gr import spaces USE_CUDA = torch.cuda.is_available() print("CUDA is available:", USE_CUDA) def transform_images(images, resolution=(1024, 1024)): images = [image.convert("RGB").resize(resolution) for image in images] # Convert to torch tensor images = [torch.tensor(np.array(image).transpose(2, 0, 1)).float() / 255 for image in images] # Normalize images = [(image - 0.5) / 0.5 for image in images] images = torch.stack(images) return images 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) 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.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 mobilesam = MobileSAM() def image_mobilesam_feature( images, node_type="block", layer=-1, ): print("Running MobileSAM") global USE_CUDA if USE_CUDA: images = images.cuda() global mobilesam feat_extractor = mobilesam if USE_CUDA: feat_extractor = feat_extractor.cuda() print("images shape:", images.shape) # attn_outputs, mlp_outputs, block_outputs = [], [], [] outputs = [] for i in range(images.shape[0]): attn_output, mlp_output, block_output = feat_extractor( images[i].unsqueeze(0) ) out_dict = { "attn": attn_output, "mlp": mlp_output, "block": block_output, } out = out_dict[node_type] out = out[layer] outputs.append(out) outputs = torch.cat(outputs, dim=0) return outputs class SAM(torch.nn.Module): def __init__(self, **kwargs): super().__init__(**kwargs) from segment_anything import sam_model_registry, SamPredictor from segment_anything.modeling.sam import Sam checkpoint = "sam_vit_b_01ec64.pth" if not os.path.exists(checkpoint): checkpoint_url = 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth' import requests r = requests.get(checkpoint_url) with open(checkpoint, 'wb') as f: f.write(r.content) sam: Sam = sam_model_registry["vit_b"](checkpoint=checkpoint) from segment_anything.modeling.image_encoder import ( window_partition, window_unpartition, ) def new_block_forward(self, x: torch.Tensor) -> torch.Tensor: shortcut = x x = self.norm1(x) # Window partition if self.window_size > 0: H, W = x.shape[1], x.shape[2] x, pad_hw = window_partition(x, self.window_size) x = self.attn(x) # Reverse window partition if self.window_size > 0: x = window_unpartition(x, self.window_size, pad_hw, (H, W)) self.attn_output = x.clone() x = shortcut + x mlp_outout = self.mlp(self.norm2(x)) self.mlp_output = mlp_outout.clone() x = x + mlp_outout self.block_output = x.clone() return x setattr(sam.image_encoder.blocks[0].__class__, "forward", new_block_forward) self.image_encoder = sam.image_encoder self.image_encoder.eval() @torch.no_grad() def forward(self, x: torch.Tensor) -> torch.Tensor: 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, blk in enumerate(self.image_encoder.blocks): attn_outputs.append(blk.attn_output) mlp_outputs.append(blk.mlp_output) block_outputs.append(blk.block_output) attn_outputs = torch.stack(attn_outputs) mlp_outputs = torch.stack(mlp_outputs) block_outputs = torch.stack(block_outputs) return attn_outputs, mlp_outputs, block_outputs sam = SAM() def image_sam_feature( images, node_type="block", layer=-1, ): global USE_CUDA if USE_CUDA: images = images.cuda() global sam feat_extractor = sam if USE_CUDA: feat_extractor = feat_extractor.cuda() # attn_outputs, mlp_outputs, block_outputs = [], [], [] outputs = [] for i in range(images.shape[0]): attn_output, mlp_output, block_output = feat_extractor( images[i].unsqueeze(0) ) out_dict = { "attn": attn_output, "mlp": mlp_output, "block": block_output, } out = out_dict[node_type] out = out[layer] outputs.append(out) outputs = torch.cat(outputs, dim=0) return outputs class DiNOv2(torch.nn.Module): def __init__(self, ver="dinov2_vitb14_reg"): super().__init__() self.dinov2 = torch.hub.load("facebookresearch/dinov2", ver) self.dinov2.requires_grad_(False) self.dinov2.eval() def new_block_forward(self, x: torch.Tensor) -> torch.Tensor: def attn_residual_func(x): return self.ls1(self.attn(self.norm1(x))) def ffn_residual_func(x): return self.ls2(self.mlp(self.norm2(x))) attn_output = attn_residual_func(x) self.attn_output = attn_output.clone() x = x + attn_output mlp_output = ffn_residual_func(x) self.mlp_output = mlp_output.clone() x = x + mlp_output block_output = x self.block_output = block_output.clone() return x setattr(self.dinov2.blocks[0].__class__, "forward", new_block_forward) @torch.no_grad() def forward(self, x): out = self.dinov2(x) attn_outputs, mlp_outputs, block_outputs = [], [], [] for i, blk in enumerate(self.dinov2.blocks): attn_outputs.append(blk.attn_output) mlp_outputs.append(blk.mlp_output) block_outputs.append(blk.block_output) attn_outputs = torch.stack(attn_outputs) mlp_outputs = torch.stack(mlp_outputs) block_outputs = torch.stack(block_outputs) return attn_outputs, mlp_outputs, block_outputs dinov2 = DiNOv2() def image_dino_feature(images, node_type="block", layer=-1): global USE_CUDA if USE_CUDA: images = images.cuda() global dinov2 feat_extractor = dinov2 if USE_CUDA: feat_extractor = feat_extractor.cuda() # attn_outputs, mlp_outputs, block_outputs = [], [], [] outputs = [] for i in range(images.shape[0]): attn_output, mlp_output, block_output = feat_extractor( images[i].unsqueeze(0) ) out_dict = { "attn": attn_output, "mlp": mlp_output, "block": block_output, } out = out_dict[node_type] out = out[layer] outputs.append(out) outputs = torch.cat(outputs, dim=0) outputs = rearrange(outputs[:, 5:, :], "b (h w) c -> b h w c", h=32, w=32) return outputs class CLIP(torch.nn.Module): def __init__(self): super().__init__() from transformers import CLIPProcessor, CLIPModel model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16") # processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16") self.model = model.eval() def new_forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, causal_attention_mask: torch.Tensor, output_attentions: Optional[bool] = False, ) -> Tuple[torch.FloatTensor]: residual = hidden_states hidden_states = self.layer_norm1(hidden_states) hidden_states, attn_weights = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, ) hw = np.sqrt(hidden_states.shape[1]-1).astype(int) self.attn_output = rearrange(hidden_states.clone()[:, 1:], "b (h w) c -> b h w c", h=hw) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.layer_norm2(hidden_states) hidden_states = self.mlp(hidden_states) self.mlp_output = rearrange(hidden_states.clone()[:, 1:], "b (h w) c -> b h w c", h=hw) hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (attn_weights,) self.block_output = rearrange(hidden_states.clone()[:, 1:], "b (h w) c -> b h w c", h=hw) return outputs setattr(self.model.vision_model.encoder.layers[0].__class__, "forward", new_forward) @torch.no_grad() def forward(self, x): out = self.model.vision_model(x) attn_outputs, mlp_outputs, block_outputs = [], [], [] for i, blk in enumerate(self.model.vision_model.encoder.layers): attn_outputs.append(blk.attn_output) mlp_outputs.append(blk.mlp_output) block_outputs.append(blk.block_output) attn_outputs = torch.stack(attn_outputs) mlp_outputs = torch.stack(mlp_outputs) block_outputs = torch.stack(block_outputs) return attn_outputs, mlp_outputs, block_outputs clip = CLIP() def image_clip_feature( images, node_type="block", layer=-1 ): global USE_CUDA if USE_CUDA: images = images.cuda() global clip feat_extractor = clip if USE_CUDA: feat_extractor = feat_extractor.cuda() # attn_outputs, mlp_outputs, block_outputs = [], [], [] outputs = [] for i in range(images.shape[0]): attn_output, mlp_output, block_output = feat_extractor( images[i].unsqueeze(0) ) out_dict = { "attn": attn_output, "mlp": mlp_output, "block": block_output, } out = out_dict[node_type] out = out[layer] outputs.append(out) outputs = torch.cat(outputs, dim=0) return outputs import hashlib import pickle import sys from collections import OrderedDict # Cache dictionary with limited size class LimitedSizeCache(OrderedDict): def __init__(self, max_size_bytes): self.max_size_bytes = max_size_bytes self.current_size_bytes = 0 super().__init__() def __setitem__(self, key, value): item_size = self.get_item_size(value) # Evict items until there is enough space while self.current_size_bytes + item_size > self.max_size_bytes: self.popitem(last=False) super().__setitem__(key, value) self.current_size_bytes += item_size def __delitem__(self, key): value = self[key] super().__delitem__(key) self.current_size_bytes -= self.get_item_size(value) def get_item_size(self, value): """Estimate the size of the value in bytes.""" return sys.getsizeof(value) # Initialize the cache with a 4GB limit cache = LimitedSizeCache(max_size_bytes=4 * 1024 * 1024 * 1024) # 4GB def compute_hash(*args, **kwargs): """Compute a unique hash based on the function arguments.""" hasher = hashlib.sha256() pickled_args = pickle.dumps((args, kwargs)) hasher.update(pickled_args) return hasher.hexdigest() def run_model_on_image(images, model_name="sam", node_type="block", layer=-1): global USE_CUDA USE_CUDA = True if model_name == "SAM(sam_vit_b)": if not USE_CUDA: gr.warning("GPU not detected. Running SAM on CPU, ~30s/image.") result = image_sam_feature(images, node_type=node_type, layer=layer) elif model_name == 'MobileSAM': result = image_mobilesam_feature(images, node_type=node_type, layer=layer) elif model_name == "DiNO(dinov2_vitb14_reg)": result = image_dino_feature(images, node_type=node_type, layer=layer) elif model_name == "CLIP(openai/clip-vit-base-patch16)": result = image_clip_feature(images, node_type=node_type, layer=layer) else: raise ValueError(f"Model {model_name} not supported.") return result def extract_features(images, model_name="MobileSAM", node_type="block", layer=-1): resolution_dict = { "MobileSAM": (1024, 1024), "SAM(sam_vit_b)": (1024, 1024), "DiNO(dinov2_vitb14_reg)": (448, 448), "CLIP(openai/clip-vit-base-patch16)": (224, 224), } images = transform_images(images, resolution=resolution_dict[model_name]) # Compute the cache key cache_key = compute_hash(images, model_name, node_type, layer) # Check if the result is already in the cache if cache_key in cache: print("Cache hit!") return cache[cache_key] result = run_model_on_image(images, model_name=model_name, node_type=node_type, layer=layer) # Store the result in the cache cache[cache_key] = result return result def compute_ncut( features, num_eig=100, num_sample_ncut=10000, affinity_focal_gamma=0.3, knn_ncut=10, knn_tsne=10, embedding_method="UMAP", num_sample_tsne=1000, perplexity=500, n_neighbors=500, min_dist=0.1, ): from ncut_pytorch import NCUT, rgb_from_tsne_3d, rgb_from_umap_3d start = time.time() eigvecs, eigvals = NCUT( num_eig=num_eig, num_sample=num_sample_ncut, device="cuda" if USE_CUDA else "cpu", affinity_focal_gamma=affinity_focal_gamma, knn=knn_ncut, ).fit_transform(features.reshape(-1, features.shape[-1])) print(f"NCUT time: {time.time() - start:.2f}s") start = time.time() if embedding_method == "UMAP": X_3d, rgb = rgb_from_umap_3d( eigvecs, n_neighbors=n_neighbors, min_dist=min_dist, ) print(f"UMAP time: {time.time() - start:.2f}s") elif embedding_method == "t-SNE": X_3d, rgb = rgb_from_tsne_3d( eigvecs, num_sample=num_sample_tsne, perplexity=perplexity, knn=knn_tsne, ) print(f"t-SNE time: {time.time() - start:.2f}s") else: raise ValueError(f"Embedding method {embedding_method} not supported.") rgb = rgb.reshape(features.shape[:3] + (3,)) return rgb def dont_use_too_much_green(image_rgb): # make sure the foval 40% of the image is red leading x1, x2 = int(image_rgb.shape[1] * 0.3), int(image_rgb.shape[1] * 0.7) y1, y2 = int(image_rgb.shape[2] * 0.3), int(image_rgb.shape[2] * 0.7) sum_values = image_rgb[:, x1:x2, y1:y2].mean((0, 1, 2)) sorted_indices = sum_values.argsort(descending=True) image_rgb = image_rgb[:, :, :, sorted_indices] return image_rgb def to_pil_images(images): return [ Image.fromarray((image * 255).cpu().numpy().astype(np.uint8)).resize((256, 256), Image.NEAREST) for image in images ] @spaces.GPU(duration=30) def main_fn( images, model_name="SAM(sam_vit_b)", layer=-1, num_eig=100, node_type="block", affinity_focal_gamma=0.3, num_sample_ncut=10000, knn_ncut=10, embedding_method="UMAP", num_sample_tsne=1000, knn_tsne=10, perplexity=500, n_neighbors=500, min_dist=0.1, ): if perplexity >= num_sample_tsne: # raise gr.Error("Perplexity must be less than the number of samples for t-SNE.") gr.Warning("Perplexity must be less than the number of samples for t-SNE.\n" f"Setting perplexity to {num_sample_tsne-1}.") perplexity = num_sample_tsne - 1 images = [image[0] for image in images] start = time.time() features = extract_features( images, model_name=model_name, node_type=node_type, layer=layer ) print(f"Feature extraction time (gpu): {time.time() - start:.2f}s") rgb = compute_ncut( features, num_eig=num_eig, num_sample_ncut=num_sample_ncut, affinity_focal_gamma=affinity_focal_gamma, knn_ncut=knn_ncut, knn_tsne=knn_tsne, num_sample_tsne=num_sample_tsne, embedding_method=embedding_method, perplexity=perplexity, n_neighbors=n_neighbors, min_dist=min_dist, ) rgb = dont_use_too_much_green(rgb) return to_pil_images(rgb) default_images = ['./images/image_0.jpg', './images/image_1.jpg', './images/image_2.jpg', './images/image_3.jpg', './images/image_5.jpg'] default_outputs = ['./images/ncut_0.jpg', './images/ncut_1.jpg', './images/ncut_2.jpg', './images/ncut_3.jpg', './images/ncut_5.jpg'] 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(["MobileSAM", "SAM(sam_vit_b)", "DiNO(dinov2_vitb14_reg)", "CLIP(openai/clip-vit-base-patch16)"], label="Model", value="MobileSAM", elem_id="model_name"), 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'), ], gr.Gallery(value=default_outputs, label="NCUT Embedding", show_label=False, elem_id="ncut", columns=[3], rows=[1], object_fit="contain", height="auto"), additional_inputs=[ 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.01, 1, step=0.01, label="Affinity focal gamma", value=0.3, elem_id="affinity_focal_gamma", info="decrease for more aggressive cleaning on the affinity matrix"), gr.Slider(100, 50000, step=100, label="num_sample (NCUT)", value=25000, elem_id="num_sample_ncut", info="for Nyström approximation"), gr.Slider(1, 100, step=1, label="KNN (NCUT)", value=10, elem_id="knn_ncut", info="for Nyström approximation"), gr.Dropdown(["t-SNE", "UMAP"], label="Embedding method", value="t-SNE", elem_id="embedding_method"), gr.Slider(100, 1000, step=100, label="num_sample (t-SNE/UMAP)", value=300, elem_id="num_sample_tsne", info="for Nyström approximation. Adding will slow down quite a lot"), gr.Slider(1, 100, step=1, label="KNN (t-SNE/UMAP)", value=10, elem_id="knn_tsne", info="for Nyström approximation"), gr.Slider(10, 500, step=10, label="Perplexity (t-SNE)", value=150, elem_id="perplexity", info="for t-SNE"), gr.Slider(10, 500, step=10, label="n_neighbors (UMAP)", value=150, elem_id="n_neighbors", info="for UMAP"), gr.Slider(0.1, 1, step=0.1, label="min_dist (UMAP)", value=0.1, elem_id="min_dist", info="for UMAP"), ] ) demo.launch()