Update app

app.py

@@ -1,91 +1,199 @@
-import os, os.path
-from os.path import splitext
-import numpy as np
+import os
+import os.path
 import sys
+from os.path import splitext
+
+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
-import wget 
+import torchvision.transforms.functional as TF
+from gradio.inputs import Image as GradioInputImage
+from gradio.outputs import Image as GradioOutputImage
+from PIL import Image
+from scipy.sparse.linalg import eigsh
+from skimage.color import label2rgb
+from torch.utils.hooks import RemovableHandle
+from torchvision import transforms
+from torchvision.utils import make_grid
 
 
-destination_folder = "output"
-destination_for_weights = "weights"
+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
 
-if os.path.exists(destination_for_weights):
-    print("The weights are at", destination_for_weights)
-else:
-    print("Creating folder at ", destination_for_weights, " to store weights")
-    os.mkdir(destination_for_weights)
-    
-segmentationWeightsURL = 'https://github.com/douyang/EchoNetDynamic/releases/download/v1.0.0/deeplabv3_resnet50_random.pt'
 
-if not os.path.exists(os.path.join(destination_for_weights, os.path.basename(segmentationWeightsURL))):
-    print("Downloading Segmentation Weights, ", segmentationWeightsURL," to ",os.path.join(destination_for_weights, os.path.basename(segmentationWeightsURL)))
-    filename = wget.download(segmentationWeightsURL, out = destination_for_weights)
-else:
-    print("Segmentation Weights already present")
+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.ToTensor(), normalize])
+    else:
+        raise NotImplementedError()
+    return transform
 
-torch.cuda.empty_cache()
 
-def collate_fn(x):
-    x, f = zip(*x)
-    i = list(map(lambda t: t.shape[1], x))
-    x = torch.as_tensor(np.swapaxes(np.concatenate(x, 1), 0, 1))
-    return x, f, i
+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
 
-model = torchvision.models.segmentation.deeplabv3_resnet50(pretrained=False, aux_loss=False)
-model.classifier[-1] = torch.nn.Conv2d(model.classifier[-1].in_channels, 1, kernel_size=model.classifier[-1].kernel_size)
 
-print("loading weights from ", os.path.join(destination_for_weights, "deeplabv3_resnet50_random"))
+# Parameters
+model_name = 'dino_vitb16'  # TODOL Figure out how to make this user-editable
+K = 5
 
+
+# Load model
+model, val_transform, patch_size, num_heads = get_model(model_name)
+
+
+# GPU
 if torch.cuda.is_available():
-    print("cuda is available, original weights")
+    print("CUDA is available, using GPU.")
     device = torch.device("cuda")
-    model = torch.nn.DataParallel(model)
     model.to(device)
-    checkpoint = torch.load(os.path.join(destination_for_weights, os.path.basename(segmentationWeightsURL)))
-    model.load_state_dict(checkpoint['state_dict'])
 else:
-    print("cuda is not available, cpu weights")
+    print("CUDA is not available, using CPU.")
     device = torch.device("cpu")
-    checkpoint = torch.load(os.path.join(destination_for_weights, os.path.basename(segmentationWeightsURL)), map_location = "cpu")
-    state_dict_cpu = {k[7:]: v for (k, v) in checkpoint['state_dict'].items()}
-    model.load_state_dict(state_dict_cpu)
 
-model.eval()
 
-def segment(inp):
-    x = inp.transpose([2, 0, 1])  #  channels-first
-    x = np.expand_dims(x, axis=0)  # adding a batch dimension    
-    
-    mean = x.mean(axis=(0, 2, 3))
-    std = x.std(axis=(0, 2, 3))
-    x = x - mean.reshape(1, 3, 1, 1)
-    x = x / std.reshape(1, 3, 1, 1)
-    
-    with torch.no_grad():
-        x = torch.from_numpy(x).type('torch.FloatTensor').to(device)
-        output = model(x)    
-    
-    y = output['out'].numpy()
-    y = y.squeeze()
-    
-    out = y>0    
-    
-    mask = inp.copy()
-    mask[out] = np.array([0, 0, 255])
+@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)
+
+    # 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)
+
+    # 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)
+
+    # Remove hook from the model
+    handle.remove()
+
+    # Normalize features
+    normalize = True
+    if normalize:
+        feats = F.normalize(feats, p=2, dim=-1)
+
+    # Compute affinity matrix
+    W_feat = (feats @ feats.T)
     
-    return mask
+    # 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()
 
-import gradio as gr
+    # # 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]
 
-i = gr.inputs.Image(shape=(112, 112))
-o = gr.outputs.Image()
+    # Arrange eigenvectors into grid
+    output_image_grid = []
+    for i in range(1, K):
+        eigenvector = eigenvectors[i].reshape(1, 1, H_pad, W_pad)
+        eigenvector = F.interpolate(eigenvector, size=(H, W), mode='nearest')  # slightly off, but for visualizations this is okay
+        # plt.imsave('./tmp.png', eigenvector.squeeze().numpy())  # save to a temporary location
+        # eigenvector = Image.open('./tmp.png').convert('RGB') # load back from our temporary location
+        output_image_grid.append(eigenvector)
+    img_tensor_grid = make_grid(output_image_grid, nrow=8, pad_value=1)
 
-examples = [["img1.jpg"], ["img2.jpg"]]
-title = None #"Left Ventricle Segmentation"
-description = "This semantic segmentation model identifies the left ventricle in echocardiogram images."
-# videos. Accurate evaluation of the motion and size of the left ventricle is crucial for the assessment of cardiac function and ejection fraction. In this interface, the user inputs apical-4-chamber images from echocardiography videos and the model will output a prediction of the localization of the left ventricle in blue. This model was trained on the publicly released EchoNet-Dynamic dataset of 10k echocardiogram videos with 20k expert annotations of the left ventricle and published as part of ‘Video-based AI for beat-to-beat assessment of cardiac function’ by Ouyang et al. in Nature, 2020."
+    # Postprocess for Gradio
+    img_tensor_grid.numpy().squeeze()
+    
+    return img_tensor_grid
+
+# Placeholders
+input_placeholders = GradioInputImage(shape=(256, 256), source="upload", tool="editor", type="pil")
+output_placeholders = GradioOutputImage(type="numpy", label=f"Eigenvectors")
+# alternatively: [GradioOutputImage(type="numpy", label=f"Eigenvector {i}") for i in range(K)]
+
+# Metadata
+examples = [["images/img1.jpg"], ["images/img2.jpg"]]
+title = "Deep Spectral Segmentation"
+description = "Deep spectral segmentation..."
 thumbnail = "https://raw.githubusercontent.com/gradio-app/hub-echonet/master/thumbnail.png"
-gr.Interface(segment, i, o, examples=examples, allow_flagging=False, analytics_enabled=False,
-	title=title, description=description, thumbnail=thumbnail).launch()
+
+# Gradio
+gr.Interface(
+    segment, 
+    input_placeholders,
+    output_placeholders,
+    examples=examples,
+    allow_flagging=False,
+    analytics_enabled=False,
+	title=title,
+    description=description,
+    thumbnail=thumbnail
+).launch()