app.py

# Show VLAD clustering for set of example images or a user image
"""
    User input:
    - Domain: Indoor, Aerial, or Urban
    - Image: Image to be clustered
    - Cluster numbers (to visualize)
    - Pixel coordinates (to pick further clusters)
    - A unique cache ID (to store the DINO forward passes)
    
    There are example images for each domain.
    
    Output:
    - All images with cluster assignments
    
    Some Gradio links:
    - Controlling layout
        - https://www.gradio.app/guides/quickstart#blocks-more-flexibility-and-control
    - Data state (persistence)
        - https://www.gradio.app/guides/interface-state
        - https://www.gradio.app/docs/state
    - Layout control
        - https://www.gradio.app/guides/controlling-layout
        - https://www.gradio.app/guides/blocks-and-event-listeners
"""

# %%
import os
import gradio as gr
import numpy as np
import cv2 as cv
import torch
from torch import nn
from torch.nn import functional as F
from torchvision import transforms as tvf
from torchvision.transforms import functional as T
from PIL import Image
import matplotlib.pyplot as plt
import distinctipy as dipy
from typing import Literal, List
import gradio as gr
import time
import glob
import shutil
from copy import deepcopy
# DINOv2 imports
from utilities import DinoV2ExtractFeatures
from utilities import VLAD

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# %%
# Configurations
T1 = Literal["query", "key", "value", "token"]
T2 = Literal["aerial", "indoor", "urban"]
DOMAINS = ["aerial", "indoor", "urban"]
T3 = Literal["dinov2_vits14", "dinov2_vitb14", "dinov2_vitl14", 
                "dinov2_vitg14"]
_ex = lambda x: os.path.realpath(os.path.expanduser(x))
dino_model: T3 = "dinov2_vitg14"
desc_layer: int = 31
desc_facet: T1 = "value"
num_c: int = 8
cache_dir: str = _ex("./cache") # Directory containing program cache
max_img_size: int = 1024    # Image resolution (max dim/size)
max_num_imgs: int = 10      # Max number of images to upload
share: bool = False          # Share application using .gradio link

# Verify inputs
assert os.path.isdir(cache_dir), "Cache directory not found"

# %%
# Model and transforms
print("Loading DINO model")
# extractor = DinoV2ExtractFeatures(dino_model, desc_layer, desc_facet, 
#                                     device=device)
extractor = None
print("DINO model loaded")
# VLAD path (directory)
ext_s = f"{dino_model}/l{desc_layer}_{desc_facet}_c{num_c}"
vc_dir = os.path.join(cache_dir, "vocabulary", ext_s)
# Base image transformations
base_tf = tvf.Compose([
    tvf.ToTensor(),
    tvf.Normalize(mean=[0.485, 0.456, 0.406], 
                    std=[0.229, 0.224, 0.225])
])


# %%
# Get VLAD object
def get_vlad_clusters(domain, pr = gr.Progress()):
    dm: T2 = str(domain).lower()
    assert dm in DOMAINS, "Invalid domain"
    # Load VLAD cluster centers
    pr(0, desc="Loading VLAD clusters")
    c_centers_file = os.path.join(vc_dir, dm, "c_centers.pt")
    if not os.path.isfile(c_centers_file):
        return f"Cluster centers not found for: {domain}", None
    c_centers = torch.load(c_centers_file)
    pr(0.5)
    num_c = c_centers.shape[0]
    desc_dim = c_centers.shape[1]
    vlad = VLAD(num_c, desc_dim, 
            cache_dir=os.path.dirname(c_centers_file))
    vlad.fit(None)  # Restore the cache
    pr(1)
    return f"VLAD clusters loaded for: {domain}", vlad


# %%
# Get VLAD descriptors
@torch.no_grad()
def get_descs(imgs_batch, pr = gr.Progress()):
    imgs_batch: List[np.ndarray] = imgs_batch
    pr(0, desc="Extracting descriptors")
    patch_descs = []
    for i, img in enumerate(imgs_batch):
        # Convert to PIL image
        pil_img = Image.fromarray(img)
        img_pt = base_tf(pil_img).to(device)
        if max(img_pt.shape[-2:]) > max_img_size:
            print(f"Image {i+1}: {img_pt.shape[-2:]}, outside")
            c, h, w = img_pt.shape
            # Maintain aspect ratio
            if h == max(img_pt.shape[-2:]):
                w = int(w * max_img_size / h)
                h = max_img_size
            else:
                h = int(h * max_img_size / w)
                w = max_img_size
            img_pt = T.resize(img_pt, (h, w), 
                interpolation=T.InterpolationMode.BICUBIC)
            pil_img = pil_img.resize((w, h))    # Backup
        # Make image patchable
        c, h, w = img_pt.shape
        h_new, w_new = (h // 14) * 14, (w // 14) * 14
        img_pt = tvf.CenterCrop((h_new, w_new))(img_pt)[None, ...]
        # Extract descriptors
        ret = extractor(img_pt).cpu()  # [1, n_p, d]
        patch_descs.append({"img": pil_img, "descs": ret})
        pr((i+1) / len(imgs_batch))
    return patch_descs, \
            f"Descriptors extracted for {len(imgs_batch)} images"


# %%
# Assign VLAD clusters (descriptor assignment)
def assign_vlad(patch_descs, vlad, pr = gr.Progress()):
    vlad: VLAD = vlad
    img_patch_descs = [pd["descs"] for pd in patch_descs]
    pr(0, desc="Assigning VLAD clusters")
    desc_assignments = []   # List[Tensor;shape=('h', 'w');int]
    for i, qu_desc in enumerate(img_patch_descs):
        # Residual vectors; 'n' could differ (based on img sizes)
        res = vlad.generate_res_vec(qu_desc[0]) # ['n', n_c, d]
        img = patch_descs[i]["img"]
        h, w, c = np.array(img).shape
        h_p, w_p = h // 14, w // 14
        h_new, w_new = h_p * 14, w_p * 14
        assert h_p * w_p == res.shape[0], "Residual incorrect!"
        # Descriptor assignments
        da = res.abs().sum(dim=2).argmin(dim=1).reshape(h_p, w_p)
        da = F.interpolate(da[None, None, ...].to(float),
                (h_new, w_new), mode="nearest")[0, 0].to(da.dtype)
        desc_assignments.append(da)
        pr((i+1) / len(img_patch_descs))
    pr(1.0)
    return desc_assignments, "VLAD clusters assigned"


# %%
# Cluster assignments to images
def get_ca_images(desc_assignments, patch_descs, alpha,
            pr = gr.Progress()):
    if desc_assignments is None or len(desc_assignments) == 0:
        return None, "First load images"
    c_colors = dipy.get_colors(num_c, rng=928, 
            colorblind_type="Deuteranomaly")
    np_colors = (np.array(c_colors) * 255).astype(np.uint8)
    # Get images with clusters
    pil_imgs = [pd["img"] for pd in patch_descs]
    res_imgs = []   # List[PIL.Image]
    pr(0, desc="Generating cluster assignment images")
    for i, pil_img in enumerate(pil_imgs):
        # Descriptor assignment image: [h, w, 3]
        da: torch.Tensor = desc_assignments[i]    # ['h', 'w']
        da_img = np.zeros((*da.shape, 3), dtype=np.uint8)
        for c in range(num_c):
            da_img[da == c] = np_colors[c]
        # Background image: [h, w, 3]
        img_np = np.array(pil_img, dtype=np.uint8)
        h, w, c = np.array(img_np).shape
        h_p, w_p = (h // 14), (w // 14)
        h_new, w_new = h_p * 14, w_p * 14
        img_np = F.interpolate(torch.tensor(img_np)\
                .permute(2, 0, 1)[None, ...], (h_new, w_new),
                mode='nearest')[0].permute(1, 2, 0).numpy()
        res_img = cv.addWeighted(img_np, 1 - alpha, da_img, alpha, 0.)
        res_imgs.append(Image.fromarray(res_img))
        pr((i+1) / len(pil_imgs))
    pr(1.0)
    return res_imgs, "Cluster assignment images generated"


# %%
print("Interface build started")
# Build the interface
with gr.Blocks() as demo:
    # ---- Helper functions ----
    # Variable number of input images
    def var_num_img(s):
        n = int(s)  # Slider value as int
        return [gr.Image.update(label=f"Image {i+1}", visible=True) \
                for i in range(n)] + [gr.Image.update(visible=False) \
                        for _ in range(max_num_imgs - n)]
    
    # ---- State declarations ----
    vlad = gr.State()   # VLAD object
    desc_assignments = gr.State()   # Cluster assignments
    imgs_batch = gr.State() # Images as batch
    patch_descs = gr.State()    # Patch descriptors
    
    # ---- All UI elements ----
    d_vals = [k.title() for k in DOMAINS]
    domain = gr.Radio(d_vals, value=d_vals[0])
    nimg_s = gr.Slider(1, max_num_imgs, value=1, step=1, 
            label="How many images?")   # How many images?
    with gr.Row():  # Dynamic row (images in columns)
        imgs = [gr.Image(label=f"Image {i+1}", visible=True) \
                for i in range(nimg_s.value)] + \
                [gr.Image(visible=False) \
                for _ in range(max_num_imgs - nimg_s.value)]
        for i, img in enumerate(imgs):  # Set image as "input"
            img.change(lambda _: None, img)
    with gr.Row():  # Dynamic row of output (cluster) images
        imgs2 = [gr.Image(label=f"VLAD Clusters {i+1}", 
                visible=False) for i in range(max_num_imgs)]
    nimg_s.change(var_num_img, nimg_s, imgs)
    blend_alpha = gr.Slider(0, 1, 0.4, step=0.01, # Cluster centers
        label="Blend alpha (weight for cluster centers)")
    bttn1 = gr.Button("Click Me!")  # Cluster assignment
    out_msg1 = gr.Markdown("Select domain and upload images")
    out_msg2 = gr.Markdown("For descriptor extraction")
    out_msg3 = gr.Markdown("Followed by VLAD assignment")
    out_msg4 = gr.Markdown("Followed by cluster images")
    
    # ---- Utility functions ----
    # A wrapper to batch the images
    def batch_images(data):
        sv = data[nimg_s]
        images: List[np.ndarray] = [data[imgs[k]] \
                for k in range(sv)]
        return images
    # A wrapper to unbatch images (and pad to max)
    def unbatch_images(imgs_batch):
        ret = [gr.Image.update(visible=False) \
                for _ in range(max_num_imgs)]
        if imgs_batch is None or len(imgs_batch) == 0:
            return ret
        for i, img_pil in enumerate(imgs_batch):
            img_np = np.array(img_pil)
            ret[i] = gr.Image.update(img_np, visible=True)
        return ret
    
    # ---- Main pipeline ----
    # Get the VLAD cluster assignment images on click
    bttn1.click(get_vlad_clusters, domain, [out_msg1, vlad])\
        .then(batch_images, {nimg_s, *imgs, imgs_batch}, imgs_batch)\
        .then(get_descs, imgs_batch, [patch_descs, out_msg2])\
        .then(assign_vlad, [patch_descs, vlad], 
                [desc_assignments, out_msg3])\
        .then(get_ca_images, 
                [desc_assignments, patch_descs, blend_alpha],
                [imgs_batch, out_msg4])\
        .then(unbatch_images, imgs_batch, imgs2)
    # If the blending changes now, update the cluster images
    blend_alpha.change(get_ca_images, 
            [desc_assignments, patch_descs, blend_alpha],
            [imgs_batch, out_msg4])\
        .then(unbatch_images, imgs_batch, imgs2)

print("Interface build completed")


# %%
# Deploy application
demo.queue().launch(share=share)
print("Application deployment ended, exiting...")