temp

app.py

@@ -37,13 +37,21 @@ transform = transforms.Compose([
         ])
 
 
-# which sf
-sf_idx_ = [55, 14, 5, 4, 52, 57, 40, 9]
+# sf_idx_ = [55, 14, 5, 4, 52, 57, 40, 9]
 
 col = plt.get_cmap('tab10')
 
-def generate_matching_superfeatures(im1, im2, scale_id=6, threshold=50):
-
+def generate_matching_superfeatures(im1, im2, scale_id=6, threshold=50, sf_ids=''):
+    print('im1:', im1.size)
+    print('im2:', im2.size)
+    # which sf 
+    sf_idx_ = [55, 14, 5, 4, 52, 57, 40, 9]
+    if sf_ids.lower().startswith('r'):
+        n_sf_ids = int(sf_ids[1:])
+        sf_idx_ = np.random.randint(256, size=n_sf_ids)
+    elif sf_ids != '':
+        sf_idx_ = map(int, sf_ids.strip().split(','))
+    
     im1_tensor = transform(im1).unsqueeze(0)
     im2_tensor = transform(im2).unsqueeze(0)
 
@@ -74,7 +82,7 @@ def generate_matching_superfeatures(im1, im2, scale_id=6, threshold=50):
         att_heat = np.array(attns1[0,i,:,:].numpy(), dtype=np.float32)
         att_heat = np.uint8(att_heat / np.max(att_heat[:]) * 255.0)
         att_heat_bin  = np.where(att_heat>threshold, 255, 0)
-        print(att_heat_bin)
+        # print(att_heat_bin)
         all_att_bin1.append(att_heat_bin)
 
         att_heat = np.array(attns2[0,i,:,:].numpy(), dtype=np.float32)
@@ -86,7 +94,7 @@ def generate_matching_superfeatures(im1, im2, scale_id=6, threshold=50):
     fin_img = []
     img1rsz = np.copy(im1_cv)
     print('im1:', im1.size)
-    print(img1rsz.shape)
+    print('img1rsz:', img1rsz.shape)
     for j, att in enumerate(all_att_bin1):
         att = cv2.resize(att, im1.size, interpolation=cv2.INTER_NEAREST)
         # att = cv2.resize(att, imgz[i].shape[:2][::-1], interpolation=cv2.INTER_CUBIC)
@@ -102,6 +110,8 @@ def generate_matching_superfeatures(im1, im2, scale_id=6, threshold=50):
     fin_img.append(img1rsz)
             
     img2rsz = np.copy(im2_cv)
+    print('im2:', im2.size)
+    print('img2rsz:', img2rsz.shape)
     for j, att in enumerate(all_att_bin2):
         att = cv2.resize(att, im2.size, interpolation=cv2.INTER_NEAREST)
         # att = cv2.resize(att, imgz[i].shape[:2][::-1], interpolation=cv2.INTER_CUBIC)
@@ -116,19 +126,21 @@ def generate_matching_superfeatures(im1, im2, scale_id=6, threshold=50):
             img2rsz[m,n, :] = col_[::-1]   
     fin_img.append(img2rsz)
 
-    fig1 = plt.figure()
-    fig1.show(cv2.cvtColor(img1rsz, cv2.COLOR_BGR2RGB))
+    fig1 = plt.figure(1)
+    plt.imshow(cv2.cvtColor(img1rsz, cv2.COLOR_BGR2RGB))
     ax1 = plt.gca()
-    ax1.axis('scaled')
+    # ax1.axis('scaled')
     ax1.axis('off')
-
     plt.tight_layout()    
+    # fig1.canvas.draw()
     
-    fig2 = plt.figure()
-    fig2.show(cv2.cvtColor(img2rsz, cv2.COLOR_BGR2RGB))
+    fig2 = plt.figure(2)
+    plt.imshow(cv2.cvtColor(img2rsz, cv2.COLOR_BGR2RGB))
     ax2 = plt.gca()
-    ax2.axis('scaled')
+    # ax2.axis('scaled')
     ax2.axis('off')
+    plt.tight_layout()    
+    # fig2.canvas.draw()
        
     # fig = plt.figure()
     # grid = ImageGrid(fig, 111, nrows_ncols=(2, 1),  axes_pad=0.1)
@@ -143,7 +155,7 @@ def generate_matching_superfeatures(im1, im2, scale_id=6, threshold=50):
     # # 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
+    return fig1, fig2, ','.join(map(str, sf_idx_))
 
 
 # GRADIO APP
@@ -155,25 +167,34 @@ article = "<p style='text-align: center'><a href='https://github.com/naver/fire'
 # css = ".output-image, .input-image {height: 40rem !important; width: 100% !important;}"
 # css = "@media screen and (max-width: 600px) { .output_image, .input_image {height:20rem !important; width: 100% !important;} }"
 # css = ".output_image, .input_image {hieght: 1000px !important}"
-css = ".input_image {height: 600px !important; width: 600px !important;} .output_image {height: 1200px !important; width: 600px !important;}"
+css = ".input_image, .input_image {height: 600px !important; width: 600px !important;} "
 # css = ".output-image, .input-image {height: 40rem !important; width: 100% !important;}"
 
 
 iface = gr.Interface(
     fn=generate_matching_superfeatures,
     inputs=[
-        gr.inputs.Image(shape=(1024, 1024), type="pil"),
-        gr.inputs.Image(shape=(1024, 1024), type="pil"),
-        gr.inputs.Slider(minimum=1, maximum=7, step=1, default=3, label="Scale"),
-        gr.inputs.Slider(minimum=1, maximum=255, step=25, default=100, label="Binarization Threshold")],
-    outputs=["plot", "plot"],
+#        gr.inputs.Image(shape=(1024, 1024), type="pil", label="First Image"),
+#        gr.inputs.Image(shape=(1024, 1024), type="pil", label="Second Image"),
+        gr.inputs.Image(type="pil", label="First Image"),
+        gr.inputs.Image(type="pil", label="Second Image"),
+        gr.inputs.Slider(minimum=0, maximum=6, step=1, default=2, label="Scale"),
+        gr.inputs.Slider(minimum=1, maximum=255, step=25, default=100, label="Binarization Threshold"),
+        gr.inputs.Textbox(lines=1, default="", label="SF IDs to show (comma separated numbers from 0-255; typing 'rX' will return X random SFs", optional=True)],
+    outputs=[
+        "plot", 
+        "plot", 
+        gr.outputs.Textbox(label="SFs")],
     # outputs=gr.outputs.Image(shape=(1024,2048), type="plot"),
     title=title,
-    theme='dark-peach',
+    theme='peach',
     layout="horizontal",
     description=description,
     article=article,
     css=css,
-    examples=[["chateau_1.png", "chateau_2.png", 3, 100]],
+    examples=[
+        ["chateau_1.png", "chateau_2.png", 2, 100, '55,14,5,4,52,57,40,9'],
+        ["anafi1.jpeg", "anafi2.jpeg", 4, 50, '99,100,142,213,236']        
+    ],
 )
 iface.launch(enable_queue=True)