update demo

CHMCorr.py

@@ -494,6 +494,7 @@ def export_visualizations_results(
         "chm-prediction": pfn,
         "chm-prediction-confidence": pr,
         "chm-nearest-neighbors": rfiles,
+        "chm-nearest-neighbors-all": reranked_nns,
         "correspondance_map": cmaps,
         "masked_cos_values": MASKED_COSINE_VALUES,
         "src-keypoints": list_of_source_points,

app.py

@@ -13,7 +13,7 @@ from PIL import Image
 from SimSearch import FaissCosineNeighbors, SearchableTrainingSet
 from ExtractEmbedding import QueryToEmbedding
 from CHMCorr import chm_classify_and_visualize
-from visualization import plot_from_reranker_output
+from visualization import plot_from_reranker_corrmap
 
 csv.field_size_limit(sys.maxsize)
 
@@ -74,7 +74,7 @@ id_to_bird_name = {
 }
 
 
-def search(query_image, draw_arcs, searcher=searcher):
+def search(query_image, searcher=searcher):
     query_embedding = QueryToEmbedding(query_image)
     scores, indices, labels = searcher.search(query_embedding, k=50)
 
@@ -101,7 +101,7 @@ def search(query_image, draw_arcs, searcher=searcher):
         query_image, kNN_results, support, training_folder
     )
 
-    fig = plot_from_reranker_output(chm_output, draw_arcs=draw_arcs)
+    fig, chm_output_label = plot_from_reranker_corrmap(chm_output)
 
     # Resize the output
 
@@ -117,35 +117,58 @@ def search(query_image, draw_arcs, searcher=searcher):
     right = (width + new_width) / 2
     bottom = (height + new_height) / 2
 
-    viz_image = image.crop((left + 540, top + 40, right - 492, bottom - 100))
+    viz_image = image.crop((left + 310, top + 60, right - 248, bottom - 80))
 
-    return viz_image, predicted_labels
+    chm_output_labels = Counter(
+        [
+            x.split("/")[-2].replace(".", " ").replace("_", " ")
+            for x in chm_output["chm-nearest-neighbors-all"][:20]
+        ]
+    )
+
+    return viz_image, {l: s / 20.0 for l, s in chm_output_labels.items()}
 
 
 blocks = gr.Blocks()
 
 with blocks:
     gr.Markdown(""" # CHM-Corr DEMO""")
-    gr.Markdown(""" ### Parameters: N=50, k=20 - Using ResNet50 features""")
+    gr.Markdown(
+        """ ### Parameters: N=50, k=20 - Using ``ImageNet Pretrained 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")
-
+    gr.Markdown(""" ### CHM-Corr Output Visualization """)
+    viz_plot = gr.Image(type="pil", label="Visualization")
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown(""" ### CHM-Corr Prediction """)
+            labels = gr.Label(label="Prediction")
+        with gr.Column():
+            gr.Markdown(""" ### Examples """)
+            examples = gr.Examples(
+                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"],
+                ],
+                inputs=[input_image],
+                outputs=[viz_plot, labels],
+                fn=search,
+                cache_examples=False,
+            )
     run_btn.click(
         search,
-        inputs=[input_image, arcs_checkbox],
-        outputs=[viz_plot, predicted_labels],
+        inputs=[input_image],
+        outputs=[viz_plot, labels],
     )
 
+
 if __name__ == "__main__":
     blocks.launch(
         debug=True,

visualization.py

@@ -38,7 +38,6 @@ def arg_topK(inputarray, topK=5):
     return np.argsort(inputarray.T.reshape(-1))[::-1][:topK]
 
 
-# FOR MULTI
 def plot_from_reranker_output(reranker_output, draw_box=True, draw_arcs=True):
     """
     visualize chm results from a reranker output dict
@@ -261,14 +260,122 @@ 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,
-    # )
 
     return fig
+
+
+def plot_from_reranker_corrmap(reranker_output, draw_box=True):
+    """
+    visualize chm results from a reranker output dict
+    """
+
+    ### SET COLORS
+    cmap = matplotlib.cm.get_cmap("gist_rainbow")
+    rgba = cmap(0.5)
+    colors = []
+    for k in range(5):
+        colors.append(cmap(k / 5.0))
+
+    ### SET POINTS
+    A = np.linspace(1 + 17, 240 - 17 - 1, 7)
+    point_list = list(product(A, A))
+
+    fig, axes = plt.subplots(
+        2,
+        7,
+        figsize=(25, 8),
+        gridspec_kw={
+            "wspace": 0,
+            "hspace": 0,
+            "width_ratios": [1, 0.28, 1, 1, 1, 1, 1],
+        },
+        facecolor=(1, 1, 1),
+    )
+
+    for i in range(2):
+        for j in range(7):
+            axes[i][j].axis("off")
+
+    axes[0][0].imshow(
+        display_transform(Image.open(reranker_output["q"]).convert("RGB"))
+    )
+
+    for i in range(min(5, reranker_output["chm-prediction-confidence"])):
+        axes[0][2 + i].imshow(
+            display_transform(Image.open(reranker_output["q"]).convert("RGB"))
+        )
+
+    # Lower ROWs CHM Top5
+    for i in range(min(5, reranker_output["chm-prediction-confidence"])):
+        axes[1][2 + i].imshow(
+            display_transform(
+                Image.open(reranker_output["chm-nearest-neighbors"][i]).convert("RGB")
+            )
+        )
+
+    if reranker_output["chm-prediction-confidence"] < 5:
+        for i in range(reranker_output["chm-prediction-confidence"], 5):
+            axes[0][2 + i].imshow(Image.new(mode="RGB", size=(240, 240), color="white"))
+            axes[1][2 + i].imshow(Image.new(mode="RGB", size=(240, 240), color="white"))
+
+    nzm = reranker_output["non_zero_mask"]
+    # Go throught top 5 nearest images
+
+    # #################################################################################
+    if draw_box:
+        # SQUARAES
+        for NC in range(min(5, reranker_output["chm-prediction-confidence"])):
+            # ON SOURCE
+            valid_patches_source = arg_topK(
+                reranker_output["masked_cos_values"][NC], topK=nzm
+            )
+
+            # ON QUERY
+            target_masked_patches = arg_topK(
+                reranker_output["masked_cos_values"][NC], topK=nzm
+            )
+            valid_patches_target = [
+                reranker_output["correspondance_map"][NC][x]
+                for x in target_masked_patches
+            ]
+            valid_patches_target = [(x[0] * 7) + x[1] for x in valid_patches_target]
+
+            patch_colors = [c for c in colors]
+            overlaps = [
+                item
+                for item, count in Counter(valid_patches_target).items()
+                if count > 1
+            ]
+
+            for O in overlaps:
+                indices = [i for i, val in enumerate(valid_patches_target) if val == O]
+                for ii in indices[1:]:
+                    patch_colors[ii] = patch_colors[indices[0]]
+
+            for i in valid_patches_source:
+                Psource = point_list[i]
+                rect = patches.Rectangle(
+                    (Psource[0] - 16, Psource[1] - 16),
+                    32,
+                    32,
+                    linewidth=2,
+                    edgecolor=patch_colors[valid_patches_source.tolist().index(i)],
+                    facecolor="none",
+                    alpha=1,
+                )
+                axes[0][2 + NC].add_patch(rect)
+
+            for i in valid_patches_target:
+                Psource = point_list[i]
+                rect = patches.Rectangle(
+                    (Psource[0] - 16, Psource[1] - 16),
+                    32,
+                    32,
+                    linewidth=2,
+                    edgecolor=patch_colors[valid_patches_target.index(i)],
+                    facecolor="none",
+                    alpha=1,
+                )
+                axes[1][2 + NC].add_patch(rect)
+
+    return fig, reranker_output["chm-prediction"]