edits

app.py

@@ -147,6 +147,7 @@ def generate_matching_superfeatures(im1, im2, Imagenet_model=False, scale_id=6,
         sf_idx_ = np.random.randint(256, size=n_sf_ids)
     else:
         sf_idx_ = map(int, sf_ids.strip().split(','))
+        n_sf_ids = len(sf_idx_)
         
     if only_matching:
         if random_mode:
@@ -238,7 +239,23 @@ def generate_matching_superfeatures(im1, im2, Imagenet_model=False, scale_id=6,
     # # Now we can save it to a numpy array.
     # data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
     # data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
-    return fig1, fig2, ','.join(map(str, sf_idx_))
+    
+
+    
+    f = lambda m,c: plt.plot([],[],marker=m, color=c, ls="none")[0]
+    handles = [f("s", col[i]) for i in range(n_sf_ids)]
+    labels = sf_idx_
+    legend = plt.legend(handles, sf_idx_, loc=3, framealpha=1, frameon=False)
+
+    def export_legend(legend, filename="legend.png"):
+        fig  = legend.figure
+        fig.canvas.draw()
+        # bbox  = legend.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
+        # fig.savefig(filename, dpi="figure", bbox_inches=bbox)
+        return fig
+
+    return fig1, fig2, export_legend(legend)
+    # ','.join(map(str, sf_idx_))
 
 
 # GRADIO APP
@@ -259,11 +276,13 @@ iface = gr.Interface(
         gr.inputs.Checkbox(default=False, label="Show random (matching) SFs"),
         gr.inputs.Textbox(lines=1, default="", label="...or show specific SF IDs:", optional=True),
         gr.inputs.Checkbox(default=True, label="Show only matching SFs"),
+
         ],
     outputs=[
         gr.outputs.Image(type="plot", label="First Image SFs"),
         gr.outputs.Image(type="plot", label="Second Image SFs"),
-        gr.outputs.Textbox(label="SFs")],
+        gr.outputs.Image(type="plot", label="SF legend")],
+        # gr.outputs.Textbox(label="SFs")],
     # outputs=gr.outputs.Image(shape=(1024,2048), type="plot"),
     title=title,
     theme='peach',