Update app.py

app.py

@@ -52,24 +52,25 @@ def imgnt_reg(img1,img2):#, model_selected):
                 'M_i' :M_i }
     #[moving_images/255, fixed_images/255]
     pred = IR_Model_tst(model_inputs)
-    img_out = wrap_imge_cropped(pred['Affine_mtrx'], fixed_images, dim1=224, dim2=128)
+    img_out = wrap_imge_cropped(pred['Affine_mtrx'].detach(), fixed_images, dim1=224, dim2=128)
     registered_img = torchvision.transforms.ToPILImage()(img_out[0])
     if with_points:
         x0_source, y0_source = generate_standard_elips(N_samples= 100)
         x_source = destandarize_point(x0_source, dim=dim, flip = False)
         y_source = destandarize_point(y0_source, dim=dim, flip = False)
-        source_im_w_points = wrap_points(fixed_images, x_source, y_source, l=1)
+        source_im_w_points = wrap_points(fixed_images.detach(), x_source, y_source, l=1)
+        
         M_Predicted = workaround_matrix(pred['Affine_mtrx'].detach(), acc = 0.5/crop_ratio)
         x0_transformed, y0_transformed = transform_standard_points(M_Predicted[0], x0_source, y0_source)
         x_transformed = destandarize_point(x0_transformed, dim=dim, flip = False)
         y_transformed = destandarize_point(y0_transformed, dim=dim, flip = False)
-        wrapped_img = wrap_points(img_out, x_transformed, y_transformed, l=1)
-        img_out2 = wrapped_img
+        wrapped_img = wrap_points(img_out.detach(), x_transformed, y_transformed, l=1)
+        img_out2 = wrapped_img.detach()
         marked_image = torchvision.transforms.ToPILImage()(img_out2[0])
     else:
-        marked_image = torchvision.transforms.ToPILImage()(torch.zeros(3,10,10))
-    return [registered_img, marked_image]
-
+        #source_im_w_points = torchvision.transforms.ToPILImage()(torch.zeros(3,128,128))
+        marked_image = torchvision.transforms.ToPILImage()(torch.zeros(3,128,128))
+    return [registered_img,source_im_w_points, marked_image]
 
 
 with gr.Blocks() as demo:
@@ -98,11 +99,14 @@ with gr.Blocks() as demo:
             #type = "filepath",
             elem_id = "image-out"
         )
-    out_image2 = gr.Image(label = "ٌMarked image",
-            elem_id = "image-out"
+    out_image2 = gr.Image(label = "ٌMarked source image",
+            elem_id = "image-out2"
+        )
+    out_image3 = gr.Image(label = "ٌMarked wrapped image",
+            elem_id = "image-out3"
         )
     inputs = [image_1, image_2]#, model_list]
-    out_image = [out_image1, out_image2]
+    out_image = [out_image1, out_image2, out_image3]
     iface = gr.Interface(fn=imgnt_reg, inputs=inputs,outputs=out_image,
     title="Imagenet registration V2",
     description="Upload 2 images to generate a registered one:",