update the UI

app.py

@@ -1,3 +1,4 @@
+import io
 import csv
 import sys
 import pickle
@@ -7,6 +8,7 @@ import gradio as gr
 import gdown
 import torchvision
 from torchvision.datasets import ImageFolder
+from PIL import Image
 
 from SimSearch import FaissCosineNeighbors, SearchableTrainingSet
 from ExtractEmbedding import QueryToEmbedding
@@ -17,42 +19,42 @@ csv.field_size_limit(sys.maxsize)
 
 concat = lambda x: np.concatenate(x, axis=0)
 
-# Embeddings
-gdown.cached_download(
-    url="https://static.taesiri.com/chm-corr/embeddings.pickle",
-    path="./embeddings.pickle",
-    quiet=False,
-    md5="002b2a7f5c80d910b9cc740c2265f058",
-)
-
-# embeddings
-# gdown.download(id="116CiA_cXciGSl72tbAUDoN-f1B9Frp89")
-
-# labels
-gdown.download(id="1SDtq6ap7LPPpYfLbAxaMGGmj0EAV_m_e")
-
-# CUB training set
-gdown.cached_download(
-    url="https://static.taesiri.com/chm-corr/CUB_train.zip",
-    path="./CUB_train.zip",
-    quiet=False,
-    md5="1bd99e73b2fea8e4c2ebcb0e7722f1b1",
-)
-
-# EXTRACT training set
-torchvision.datasets.utils.extract_archive(
-    from_path="CUB_train.zip",
-    to_path="data/",
-    remove_finished=False,
-)
-
-# CHM Weights
-gdown.cached_download(
-    url="https://static.taesiri.com/chm-corr/pas_psi.pt",
-    path="pas_psi.pt",
-    quiet=False,
-    md5="6b7b4d7bad7f89600fac340d6aa7708b",
-)
+# # Embeddings
+# gdown.cached_download(
+#     url="https://static.taesiri.com/chm-corr/embeddings.pickle",
+#     path="./embeddings.pickle",
+#     quiet=False,
+#     md5="002b2a7f5c80d910b9cc740c2265f058",
+# )
+
+# # embeddings
+# # gdown.download(id="116CiA_cXciGSl72tbAUDoN-f1B9Frp89")
+
+# # labels
+# gdown.download(id="1SDtq6ap7LPPpYfLbAxaMGGmj0EAV_m_e")
+
+# # CUB training set
+# gdown.cached_download(
+#     url="https://static.taesiri.com/chm-corr/CUB_train.zip",
+#     path="./CUB_train.zip",
+#     quiet=False,
+#     md5="1bd99e73b2fea8e4c2ebcb0e7722f1b1",
+# )
+
+# # EXTRACT training set
+# torchvision.datasets.utils.extract_archive(
+#     from_path="CUB_train.zip",
+#     to_path="data/",
+#     remove_finished=False,
+# )
+
+# # CHM Weights
+# gdown.cached_download(
+#     url="https://static.taesiri.com/chm-corr/pas_psi.pt",
+#     path="pas_psi.pt",
+#     quiet=False,
+#     md5="6b7b4d7bad7f89600fac340d6aa7708b",
+# )
 
 
 # Caluclate Accuracy
@@ -72,7 +74,7 @@ id_to_bird_name = {
 }
 
 
-def search(query_image, searcher=searcher):
+def search(query_image, draw_arcs, searcher=searcher):
     query_embedding = QueryToEmbedding(query_image)
     scores, indices, labels = searcher.search(query_embedding, k=50)
 
@@ -99,27 +101,53 @@ def search(query_image, searcher=searcher):
         query_image, kNN_results, support, training_folder
     )
 
-    viz_plot = plot_from_reranker_output(chm_output, draw_arcs=False)
+    fig = plot_from_reranker_output(chm_output, draw_arcs=draw_arcs)
 
-    return viz_plot, predicted_labels, gallery_images
+    # Resize the output
 
+    img_buf = io.BytesIO()
+    fig.savefig(img_buf, format="jpg")
+    image = Image.open(img_buf)
+    width, height = image.size
+    new_width = width
+    new_height = height
 
-demo = gr.Interface(
-    search,
-    gr.Image(type="filepath"),
-    ["plot", "label", "gallery"],
-    examples=[
-        ["./examples/bird.jpg"],
-        ["./examples/Red_Winged_Blackbird_0012_6015.jpg"],
-        ["./examples/Red_Winged_Blackbird_0025_5342.jpg"],
-        ["./examples/sample1.jpeg"],
-        ["./examples/sample2.jpeg"],
-        ["./examples/Yellow_Headed_Blackbird_0020_8549.jpg"],
-        ["./examples/Yellow_Headed_Blackbird_0026_8545.jpg"],
-    ],
-    description="WIP - kNN on CUB dataset",
-    title="Work in Progress - CHM-Corr",
-)
+    left = (width - new_width) / 2
+    top = (height - new_height) / 2
+    right = (width + new_width) / 2
+    bottom = (height + new_height) / 2
+
+    viz_image = image.crop((left + 540, top + 40, right - 492, bottom - 100))
+
+    return viz_image, predicted_labels
+
+
+blocks = gr.Blocks()
+
+with blocks:
+    gr.Markdown(""" # CHM-Corr DEMO""")
+    gr.Markdown(""" ### Parameters: N=50, k=20 - Using ResNet50 features""")
+
+    # with gr.Row():
+    input_image = gr.Image(type="filepath")
+    with gr.Column():
+        arcs_checkbox = gr.Checkbox(label="Draw Arcs")
+    run_btn = gr.Button("Classify")
+
+    # with gr.Column():
+    gr.Markdown(""" ### CHM-Corr Output """)
+    viz_plot = gr.Image(type="pil")
+    gr.Markdown(""" ### kNN Predicted Labels """)
+    predicted_labels = gr.Label(label="kNN Prediction")
+
+    run_btn.click(
+        search,
+        inputs=[input_image, arcs_checkbox],
+        outputs=[viz_plot, predicted_labels],
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    blocks.launch(
+        debug=True,
+        enable_queue=True,
+    )

requirements.txt

@@ -1,4 +1,3 @@
-faiss-cpu==1.7.2
 gdown
 gradio
 numpy
@@ -8,3 +7,4 @@ torchvision
 tqdm
 tensorboardX==2.5
 matplotlib
+faiss-cpu==1.7.2

visualization.py

@@ -261,14 +261,14 @@ def plot_from_reranker_output(reranker_output, draw_box=True, draw_arcs=True):
         color="black",
         fontsize=22,
     )
-    fig.text(
-        0.8,
-        0.95,
-        f"KNN: {reranker_output['knn-prediction']}",
-        ha="right",
-        va="bottom",
-        color="black",
-        fontsize=22,
-    )
+    # fig.text(
+    #     0.8,
+    #     0.95,
+    #     f"KNN: {reranker_output['knn-prediction']}",
+    #     ha="right",
+    #     va="bottom",
+    #     color="black",
+    #     fontsize=22,
+    # )
 
     return fig