Spaces:
Runtime error
Runtime error
File size: 9,726 Bytes
cd30264 e11cde7 cd30264 94e6f4c cfe4337 95dc30b cfe4337 95dc30b cfe4337 94e6f4c cfe4337 95dc30b cfe4337 95dc30b cfe4337 cd30264 c7054d6 cd30264 cfe4337 95dc30b cfe4337 95dc30b cd30264 95dc30b cfe4337 94e6f4c cfe4337 95dc30b cfe4337 cd30264 cfe4337 95dc30b cfe4337 95dc30b cfe4337 95dc30b cfe4337 e11cde7 cfe4337 95dc30b cfe4337 95dc30b cfe4337 95dc30b cfe4337 cd30264 e11cde7 59b4598 e11cde7 59b4598 e11cde7 94e6f4c e11cde7 cd30264 e11cde7 cfe4337 e11cde7 94e6f4c cfe4337 e11cde7 94e6f4c 95dc30b cfe4337 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 |
import gc
import io
from collections import namedtuple
from typing import Tuple
import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import scipy.sparse
import torch
import torch.nn.functional as F
import torchvision.transforms.functional as TF
from gradio.inputs import Image as GradioInputImage
from gradio.outputs import Image as GradioOutputImage
from matplotlib.pyplot import get_cmap
from PIL import Image
from scipy.sparse.linalg import eigsh
from torch.utils.hooks import RemovableHandle
from torchvision import transforms
from torchvision.utils import make_grid
def get_model(name: str):
if 'dino' in name:
model = torch.hub.load('facebookresearch/dino:main', name)
model.fc = torch.nn.Identity()
val_transform = get_transform(name)
patch_size = model.patch_embed.patch_size
num_heads = model.blocks[0].attn.num_heads
elif name in ['mocov3_vits16', 'mocov3_vitb16']:
model = torch.hub.load('facebookresearch/dino:main', name.replace('mocov3', 'dino'))
checkpoint_file, size_char = {
'mocov3_vits16': ('vit-s-300ep-timm-format.pth', 's'),
'mocov3_vitb16': ('vit-b-300ep-timm-format.pth', 'b'),
}[name]
url = f'https://dl.fbaipublicfiles.com/moco-v3/vit-{size_char}-300ep/vit-{size_char}-300ep.pth.tar'
checkpoint = torch.hub.load_state_dict_from_url(url)
model.load_state_dict(checkpoint['model'])
model.fc = torch.nn.Identity()
val_transform = get_transform(name)
patch_size = model.patch_embed.patch_size
num_heads = model.blocks[0].attn.num_heads
else:
raise ValueError(f'Unsupported model: {name}')
model = model.eval()
return model, val_transform, patch_size, num_heads
def get_transform(name: str):
if any(x in name for x in ('dino', 'mocov3', 'convnext', )):
normalize = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
transform = transforms.Compose([
transforms.Resize(size=512, interpolation=TF.InterpolationMode.BICUBIC, max_size=1024),
transforms.ToTensor(),
normalize
])
else:
raise NotImplementedError()
return transform
def get_diagonal(W: scipy.sparse.csr_matrix, threshold: float = 1e-12):
D = W.dot(np.ones(W.shape[1], W.dtype))
D[D < threshold] = 1.0 # Prevent division by zero.
D = scipy.sparse.diags(D)
return D
# Cache
torch.cuda.empty_cache()
# Parameters
model_name = 'dino_vitb16' # TODO: Figure out how to make this user-editable
K = 5
# Load model
model, val_transform, patch_size, num_heads = get_model(model_name)
# Add hook
which_block = -1
if 'dino' in model_name or 'mocov3' in model_name:
feat_out = {}
def hook_fn_forward_qkv(module, input, output):
feat_out["qkv"] = output
handle: RemovableHandle = model._modules["blocks"][which_block]._modules["attn"]._modules["qkv"].register_forward_hook(
hook_fn_forward_qkv
)
else:
raise ValueError(model_name)
# GPU
if torch.cuda.is_available():
print("CUDA is available, using GPU.")
device = torch.device("cuda")
model.to(device)
else:
print("CUDA is not available, using CPU.")
device = torch.device("cpu")
@torch.no_grad()
def segment(inp: Image):
# NOTE: The image is already resized to the desired size.
# Preprocess image
images: torch.Tensor = val_transform(inp)
images = images.unsqueeze(0).to(device)
# Reshape image
P = patch_size
B, C, H, W = images.shape
H_patch, W_patch = H // P, W // P
H_pad, W_pad = H_patch * P, W_patch * P
T = H_patch * W_patch + 1 # number of tokens, add 1 for [CLS]
# Crop image to be a multiple of the patch size
images = images[:, :, :H_pad, :W_pad]
# Extract features
if 'dino' in model_name or 'mocov3' in model_name:
model.get_intermediate_layers(images)[0].squeeze(0)
output_qkv = feat_out["qkv"].reshape(B, T, 3, num_heads, -1 // num_heads).permute(2, 0, 3, 1, 4)
feats = output_qkv[1].transpose(1, 2).reshape(B, T, -1)[:, 1:, :].squeeze(0)
else:
raise ValueError(model_name)
# Normalize features
normalize = True
if normalize:
feats = F.normalize(feats, p=2, dim=-1)
# Compute affinity matrix
W_feat = (feats @ feats.T)
# Feature affinities
threshold_at_zero = True
if threshold_at_zero:
W_feat = (W_feat * (W_feat > 0))
W_feat = W_feat / W_feat.max() # NOTE: If features are normalized, this naturally does nothing
W_feat = W_feat.cpu().numpy()
# # NOTE: Here is where we would add the color information. For simplicity, we will not add it here.
# W_comb = W_feat + W_color * image_color_lambda # combination
# D_comb = np.array(get_diagonal(W_comb).todense()) # is dense or sparse faster? not sure, should check
# Diagonal
W_comb = W_feat
D_comb = np.array(get_diagonal(W_comb).todense()) # is dense or sparse faster? not sure, should check
# Compute eigenvectors
try:
eigenvalues, eigenvectors = eigsh(D_comb - W_comb, k=(K + 1), sigma=0, which='LM', M=D_comb)
except:
eigenvalues, eigenvectors = eigsh(D_comb - W_comb, k=(K + 1), which='SM', M=D_comb)
eigenvalues = torch.from_numpy(eigenvalues)
eigenvectors = torch.from_numpy(eigenvectors.T).float()
# Resolve sign ambiguity
for k in range(eigenvectors.shape[0]):
if 0.5 < torch.mean((eigenvectors[k] > 0).float()).item() < 1.0: # reverse segment
eigenvectors[k] = 0 - eigenvectors[k]
# Arrange eigenvectors into grid
# cmap = get_cmap('viridis')
output_images = []
# eigenvectors_upscaled = []
for i in range(1, K + 1):
eigenvector = eigenvectors[i].reshape(1, 1, H_patch, W_patch) # .reshape(1, 1, H_pad, W_pad)
eigenvector: torch.Tensor = F.interpolate(eigenvector, size=(H_pad, W_pad), mode='bilinear', align_corners=False) # slightly off, but for visualizations this is okay
buffer = io.BytesIO()
plt.imsave(buffer, eigenvector.squeeze().numpy(), format='png') # save to a temporary location
buffer.seek(0)
eigenvector_vis = Image.open(buffer).convert('RGB')
# eigenvector_vis = TF.to_tensor(eigenvector_vis).unsqueeze(0)
eigenvector_vis = np.array(eigenvector_vis)
# eigenvectors_upscaled.append(eigenvector)
output_images.append(eigenvector_vis)
# output_images = torch.cat(output_images, dim=0)
# output_images = make_grid(output_images, nrow=8, pad_value=1)
# Also add CRF
if False:
# Imports
import denseCRF
# Parameters
ParamsCRF = namedtuple('ParamsCRF', 'w1 alpha beta w2 gamma it')
DEFAULT_CRF_PARAMS = ParamsCRF(
w1 = 6, # weight of bilateral term # 10.0,
alpha = 40, # spatial std # 80,
beta = 13, # rgb std # 13,
w2 = 3, # weight of spatial term # 3.0,
gamma = 3, # spatial std # 3,
it = 5.0, # iteration # 5.0,
)
# Get unary potentials
unary_potentials = eigenvectors_upscaled[0].squeeze(1).squeeze(0)
unary_potentials = (unary_potentials - unary_potentials.min()) / (unary_potentials.max() - unary_potentials.min())
unary_potentials_np = torch.stack((1 - unary_potentials, unary_potentials), dim=-1).cpu().numpy()
img_np = images.cpu().numpy().transpose(0, 2, 3, 1)
img_np = (img_np * 255).astype(np.uint8)[0]
# Return result of CRF
out = denseCRF.densecrf(img_np, unary_potentials_np, DEFAULT_CRF_PARAMS)
out = out * 255
output_images.append(out)
# # Postprocess for Gradio
# output_images = np.array(TF.to_pil_image(output_images))
print(f'{len(output_images)=}')
# Garbage collection and other memory-related things
gc.collect()
del eigenvector, eigenvector_vis, eigenvectors, W_comb, D_comb
return output_images
# Placeholders
input_placeholders = GradioInputImage(source="upload", tool="editor", type="pil")
# output_placeholders = GradioOutputImage(type="numpy", label=f"Eigenvectors")
output_placeholders = [GradioOutputImage(type="numpy", label=(f"Eigenvector {i}")) for i in range(K)]
# Metadata
examples = [f"examples/{stem}.jpg" for stem in [
'2008_000099', '2008_000499', '2007_009446', '2007_001586', '2010_001256', '2008_000764', '2008_000705', # '2007_000039'
]]
title = "Demo: Deep Spectral Methods for Unsupervised Localization and Segmentation"
description = """
This is a demo of <a href="https://lukemelas.github.io/deep-spectral-segmentation/">Deep Spectral Methods: A Surprisingly Strong Baseline for Unsupervised Semantic Segmentation and Localization
</a> (CVPR 2022 Oral).
Our method decomposes an image into a set of soft segments in a <em>completely unsupervised</em> manner. Specifically, we extract the Laplacian eigenvectors of a feature affinity matrix from a large self-supervised network, and we find that we find that these eigenvectors can be readily used to localize and segment objects.
Below, you can upload an image (or select one of the examples) and see how our method decomposes it into soft segments. Hopefully it will localize some of the objects or semantic segments in your image!
"""
thumbnail = "https://raw.githubusercontent.com/gradio-app/hub-echonet/master/thumbnail.png"
# Gradio
gr.Interface(
segment,
input_placeholders,
output_placeholders,
examples=examples,
allow_flagging=False,
analytics_enabled=False,
title=title,
description=description,
thumbnail=thumbnail
).launch()
|