Pushed the second version of HF App

app.py

@@ -23,6 +23,34 @@
         - https://www.gradio.app/guides/blocks-and-event-listeners
 """
 
+# A markdown string shown at the top of the app
+header_markdown = """
+# AnyLoc Demo
+
+\|  [Website](https://anyloc.github.io/) \| \
+    [GitHub](https://github.com/AnyLoc/AnyLoc) \| \
+    [YouTube](https://youtu.be/ITo8rMInatk) \|
+
+
+This space contains a collection of demos for AnyLoc. Each demo is a \
+self-contained application in the tabs below. The following \
+applications are included
+
+1. **GeM t-SNE Projection**: Upload a set of images and see where \
+    they land on a t-SNE projection of GeM descriptors from many \
+    domains. This can be used to guide domain selection (from a few \
+    representative images).
+2. **Cluster Visualization**: This visualizes the VLAD cluster \
+    assignments for the patch descriptors. You need to select the \
+    domain for loading VLAD cluster centers (vocabulary).
+
+We do **not** save any images uploaded to the demo. Some errors may \
+leave a log. We do not collect any information about the user.
+
+🥳 Thanks to HuggingFace for providing a free GPU for this demo.
+
+"""
+
 # %%
 import os
 import gradio as gr
@@ -35,12 +63,15 @@ from torchvision import transforms as tvf
 from torchvision.transforms import functional as T
 from PIL import Image
 import matplotlib.pyplot as plt
+from sklearn.manifold import TSNE
 import distinctipy as dipy
+import joblib
 from typing import Literal, List
 import gradio as gr
 import time
 import glob
 import shutil
+import matplotlib.pyplot as plt
 from copy import deepcopy
 # DINOv2 imports
 from utilities import DinoV2ExtractFeatures
@@ -62,22 +93,33 @@ 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
+max_num_imgs: int = 16      # 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 = None
+# extractor = None  # FIXME: For quick testing only
 extractor = DinoV2ExtractFeatures(dino_model, desc_layer, desc_facet, 
                                     device=device)
 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)
+assert os.path.isdir(vc_dir), f"VLAD directory: {vc_dir} not found"
+# GeM path (cache)
+gem_cf = os.path.join(cache_dir, "gem_cache", "result_dino_v2.gz")
+assert os.path.isfile(gem_cf), f"GeM cache: {gem_cf} not found"
+gem_cache = joblib.load(gem_cf)
+assert gem_cache["model"]["type"] == dino_model
+assert gem_cache["model"]["layer"] == desc_layer
+assert gem_cache["model"]["facet"] == desc_facet
+fig = plt.figure()  # Main figure
+fig.clear()
 # Base image transformations
 base_tf = tvf.Compose([
     tvf.ToTensor(),
@@ -115,6 +157,9 @@ def get_descs(imgs_batch, pr = gr.Progress()):
     pr(0, desc="Extracting descriptors")
     patch_descs = []
     for i, img in enumerate(imgs_batch):
+        if img is None:
+            print(f"Image {i+1} is None")
+            continue
         # Convert to PIL image
         pil_img = Image.fromarray(img)
         img_pt = base_tf(pil_img).to(device)
@@ -139,6 +184,7 @@ def get_descs(imgs_batch, pr = gr.Progress()):
         ret = extractor(img_pt).cpu()  # [1, n_p, d]
         patch_descs.append({"img": pil_img, "descs": ret})
         pr((i+1) / len(imgs_batch))
+    pr(1.0)
     return patch_descs, \
             f"Descriptors extracted for {len(imgs_batch)} images"
 
@@ -173,7 +219,10 @@ def assign_vlad(patch_descs, vlad, pr = gr.Progress()):
 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"
+        if not 0 <= alpha <= 1:
+            return None, f"Invalid alpha value: {alpha} (should be "\
+                    "between 0 and 1)"
+        return None, "First load the images"
     c_colors = dipy.get_colors(num_c, rng=928, 
             colorblind_type="Deuteranomaly")
     np_colors = (np.array(c_colors) * 255).astype(np.uint8)
@@ -202,43 +251,177 @@ def get_ca_images(desc_assignments, patch_descs, alpha,
     return res_imgs, "Cluster assignment images generated"
 
 
+# %%
+# Get GeM descriptors from cache
+def get_gem_descs_cache(use_d, pr = gr.Progress()):
+    use_d: List[str] = use_d
+    if len(use_d) == 0:
+        return "Select at least one domain", None
+    else:
+        use_d = [d.lower() for d in use_d]
+    indoor_datasets = ["baidu_datasets", "gardens", "17places"]
+    urban_datasets = ["pitts30k", "st_lucia", "Oxford"]
+    aerial_datasets = ["Tartan_GNSS_test_rotated", 
+            "Tartan_GNSS_test_notrotated", "VPAir"]
+    pr(0, desc="Loading GeM descriptors from cache")
+    gem_descs = {
+        "labels": [],
+        "descs": [],
+    }
+    for i, ds in enumerate(gem_cache["data"]):
+        # GeM descriptors from data: n_desc, desc_dim
+        d: np.ndarray = gem_cache["data"][ds]["descriptors"]
+        if ds in indoor_datasets and "indoor" in use_d:
+            gem_descs["labels"].extend(["indoor"] * d.shape[0])
+        elif ds in urban_datasets and "urban" in use_d:
+            gem_descs["labels"].extend(["urban"] * d.shape[0])
+        elif ds in aerial_datasets and "aerial" in use_d:
+            gem_descs["labels"].extend(["aerial"] * d.shape[0])
+        else:
+            continue
+        gem_descs["descs"].append(d)
+        pr((i+1) / len(gem_cache["data"]))
+    gem_descs["descs"] = np.concatenate(gem_descs["descs"], axis=0)
+    pr(1.0)
+    return "GeM descriptors loaded from cache", gem_descs
+
+
+# %%
+# Get GeM pooled features of the uploaded images
+def get_add_gem_descs(imgs_batch, gem_descs, pr = gr.Progress()):
+    imgs_batch: List[np.ndarray] = imgs_batch
+    gem_descs: dict = gem_descs
+    pr(0, desc="Extracting GeM descriptors")
+    num_imgs_extracted = 0
+    for i, img in enumerate(imgs_batch):
+        if img is None:
+            print(f"Image {i+1} is None")
+            continue
+        # 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]
+        # Get the GeM pooled descriptor
+        x = torch.mean(ret**3, dim=-2)
+        g_res = x.to(torch.complex64) ** (1/3)
+        g_res = torch.abs(g_res) * torch.sign(x)    # [1, d]
+        g_res = g_res.numpy()
+        # Add to state
+        gem_descs["labels"].append(f"Image{i+1}")
+        gem_descs["descs"] = np.concatenate([gem_descs["descs"], 
+                                            g_res])
+        num_imgs_extracted += 1
+        pr((i+1) / len(imgs_batch))
+    pr(1.0)
+    gem_descs["num_uimgs"] = num_imgs_extracted
+    return gem_descs, "GeM descriptors extracted"
+
+
+# %%
+# Apply tSNE to the GeM descriptors
+def get_tsne_fm_gem(gem_descs, pr = gr.Progress()):
+    pr(0, desc="Applying tSNE to GeM descriptors")
+    desc_all: np.ndarray = gem_descs["descs"]   # [n, d_dim]
+    labels_all: List[str] = gem_descs["labels"] # [n]
+    # tSNE projection
+    tsne = TSNE(n_components=2, random_state=30, perplexity=50, 
+            learning_rate=200, init='random')
+    desc_2d = tsne.fit_transform(desc_all)
+    # Result
+    tsne_pts = {
+        "labels": labels_all,
+        "pts": desc_2d,
+        "num_uimgs": gem_descs["num_uimgs"],    # Number of user imgs
+    }
+    pr(1.0)
+    return tsne_pts, "tSNE projection done"
+
+
+# %%
+# Plot tSNE to matplotlib figure
+def plot_tsne(tsne_pts):
+    colors = {
+        "aerial": (80/255,  0/255,  80/255),
+        "indoor": ( 0/255, 76/255, 204/255),
+        "urban":  ( 0/255, 204/255,  0/255),
+    }
+    ni = int(tsne_pts["num_uimgs"])
+    # Custom colors for user images
+    ucs = dipy.get_colors(ni, exclude_colors=list(colors.values())\
+            .extend([(0, 0, 0), (1, 1, 1)]), 
+            colorblind_type="Deuteranomaly")
+    for i in range(ni):
+        colors[f"Image{i+1}"] = ucs[i]
+    fig.clear()
+    gs = fig.add_gridspec(1, 1)
+    ax = fig.add_subplot(gs[0, 0])
+    ax.set_title("tSNE Projection")
+    for i, domain in enumerate(list(colors.keys())):
+        pts = tsne_pts["pts"][np.array(tsne_pts["labels"]) == domain]
+        if domain.startswith("Image"):
+            m = "x"
+        else:
+            m = "o"
+        ax.scatter(pts[:, 0], pts[:, 1], label=domain, marker=m,
+                color=colors[domain])
+    # Put legend at the bottom of axis
+    ax.legend()
+    ax.set_xticks([])
+    ax.set_yticks([])
+    fig.set_tight_layout(True)
+    # fig.set_tight_layout(True)
+    return fig, "tSNE plot created"
+
+
 # %%
 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 ----
+
+
+# Tab for VLAD cluster assignment visualization
+def tab_cluster_viz():
     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?
+    domain = gr.Radio(d_vals, value=d_vals[0], label="Domain",
+            info="The domain of images (for loading VLAD vocabulary)")
+    nimg_s = gr.Number(2, label="How many images?", precision=0,
+            info=f"Between '1' and '{max_num_imgs}' images. Press "\
+                    "enter/return to register")
     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)] + \
+                for i in range(int(nimg_s.value))] + \
                 [gr.Image(visible=False) \
-                for _ in range(max_num_imgs - nimg_s.value)]
+                for _ in range(max_num_imgs - int(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)")
+    nimg_s.submit(var_num_img, nimg_s, imgs)
+    blend_alpha = gr.Number(0.4, label="Blending alpha",
+        info="Weight for cluster centers (between 0 and 1). "\
+            "Higher (close to 1) means greater emphasis on cluster "\
+                "visibility. Lower (closer to 0) will show the "\
+                "underlying image more. "\
+            "Press enter/return to register")
     bttn1 = gr.Button("Click Me!")  # Cluster assignment
+    gr.Markdown("### Status strings")
     out_msg1 = gr.Markdown("Select domain and upload images")
     out_msg2 = gr.Markdown("For descriptor extraction")
     out_msg3 = gr.Markdown("Followed by VLAD assignment")
@@ -247,21 +430,40 @@ with gr.Blocks() as demo:
     # ---- Utility functions ----
     # A wrapper to batch the images
     def batch_images(data):
-        sv = data[nimg_s]
+        sv = int(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):
+    def unbatch_images(imgs_batch, nimg):
         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)
+        for i in range(nimg):   # nimg only to match input layout
+            if i < len(imgs_batch):
+                img_np = np.array(imgs_batch[i])
+            else:
+                img_np = None
             ret[i] = gr.Image.update(img_np, visible=True)
         return ret
     
+    # ---- Examples ----
+    # Two images from each domain
+    gr.Examples(
+        [
+        ["Aerial", 2, 
+            "ex_aerial_nardo-air_db-42.png",
+            "ex_aerial_nardo-air_qu-42.png",],
+        ["Indoor", 2,
+            "ex_indoor_17places_db-75.jpg",
+            "ex_indoor_17places_qu-75.jpg"],
+        ["Urban", 2,
+            "ex_urban_oxford_db-75.png",
+            "ex_urban_oxford_qu-75.png"],],
+        [domain, nimg_s, *imgs],
+    )
+    
     # ---- Main pipeline ----
     # Get the VLAD cluster assignment images on click
     bttn1.click(get_vlad_clusters, domain, [out_msg1, vlad])\
@@ -272,12 +474,113 @@ with gr.Blocks() as demo:
         .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, 
+        .then(unbatch_images, [imgs_batch, nimg_s], imgs2)
+    # If the blending changes now, update the cluster images only
+    blend_alpha.submit(get_ca_images, 
             [desc_assignments, patch_descs, blend_alpha],
             [imgs_batch, out_msg4])\
-        .then(unbatch_images, imgs_batch, imgs2)
+        .then(unbatch_images, [imgs_batch, nimg_s], imgs2)
+
+
+# Tab for GeM t-SNE projection plot
+def tab_gem_tsne():
+    d_vals = [k.title() for k in DOMAINS]
+    dms = gr.CheckboxGroup(d_vals, value=d_vals, label="Domains",
+            info="The domains to use for the t-SNE projection")
+    nimg_s = gr.Number(2, label="How many images?", precision=0,
+            info=f"Between '1' and '{max_num_imgs}' images. Press "\
+                    "enter/return to register")
+    with gr.Row():  # Dynamic row (images in columns)
+        imgs = [gr.Image(label=f"Image {i+1}", visible=True) \
+                for i in range(int(nimg_s.value))] + \
+                [gr.Image(visible=False) \
+                for _ in range(max_num_imgs - int(nimg_s.value))]
+        for i, img in enumerate(imgs):  # Set image as "input"
+            img.change(lambda _: None, img)
+    nimg_s.submit(var_num_img, nimg_s, imgs)
+    tsne_plot = gr.Plot(None, label="tSNE Plot")
+    out_msg1 = gr.Markdown("Select domains")
+    out_msg2 = gr.Markdown("Upload images")
+    out_msg3 = gr.Markdown("Wait for tSNE plots")
+    
+    # A wrapper to batch the images
+    def batch_images(data):
+        sv = int(data[nimg_s])
+        # images: List[np.ndarray] = [data[imgs[k]] \
+        #         for k in range(sv)]
+        images: List[np.ndarray] = []
+        for k in range(sv):
+            img = data[imgs[k]]
+            if img is None:
+                return None, f"Image {k+1} is None!"
+            images.append(img)
+        return images, "Images batched"
+    
+    bttn1 = gr.Button("Click Me!")
+    
+    # ---- Main pipeline ----
+    # Get the tSNE plot
+    bttn1.click(get_gem_descs_cache, dms, [out_msg1, gem_descs])\
+        .then(batch_images, {nimg_s, *imgs, imgs_batch}, 
+                [imgs_batch, out_msg2])\
+        .then(get_add_gem_descs, [imgs_batch, gem_descs],
+                [gem_descs, out_msg2])\
+        .then(get_tsne_fm_gem, gem_descs, [tsne_pts, out_msg3])\
+        .then(plot_tsne, tsne_pts, [tsne_plot, out_msg3])
+
+
+# Build the interface
+with gr.Blocks() as demo:
+    # Main header
+    gr.Markdown(header_markdown)
+    
+    # ---- Helper functions ----
+    # Variable number of input images (show/hide UI image array)
+    def var_num_img(s):
+        n = int(s)  # Slider (string) value as int
+        assert 1 <= n <= max_num_imgs, f"Invalid num of images: {n}!"
+        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
+    gem_descs = gr.State()  # GeM descriptors (of each state)
+    tsne_pts = gr.State()   # tSNE points
+    
+    # ---- All UI elements ----
+    with gr.Tab("GeM t-SNE Projection"):
+        gr.Markdown(
+            """
+            ## GeM t-SNE Projection
+            
+            Select the domains (toggle visibility) for t-SNE plot. \
+            Enter the number of images to upload and upload images. \
+            Then click the button to get the t-SNE plot.
+            
+            """)
+        tab_gem_tsne()
+    
+    with gr.Tab("Cluster Visualization"):
+        gr.Markdown(
+            """
+            ## Cluster Visualizations
+            
+            Select the domain for the images (all should be from the \
+            same domain). Enter the number of images to upload. \
+            Upload the images. Then click the button to get the \
+            cluster assignment images.
+            
+            You can also directly click on one of the examples (at \
+            the bottom) to load the data and then click the button \
+            to get the cluster assignment images.
+            
+            """)
+        tab_cluster_viz()
 
 print("Interface build completed")
 

cache/gem_cache/result_dino_v2.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d27de9e8552eeeca6fd3f99fca135429169adcf926388d65eb44bb1ba9391f5
+size 1990740

cache/gem_cache/result_dino_v2_tsne.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:77d978ec8fcb5f10075db999d037925846f2cbad522e5d0266ddcbfb50cd99de
+size 3192

requirements.txt

@@ -8,4 +8,5 @@ matplotlib
 distinctipy
 einops
 fast_pytorch_kmeans
-
+joblib==1.2.0
+sklearn==1.0.2