Update app.py

app.py

@@ -24,7 +24,7 @@ from torchvision.models import resnet18
 
 
 #device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-
+with_points = True
 file_savingfolder = './modelbest/'
 ext = '_bestVal'
 core_model_tst = resnet18(pretrained=True)
@@ -37,6 +37,49 @@ IR_Model_tst.load_state_dict(torch.load(file_savingfolder+'IR_Model'+ext+'.pth',
 IR_Model_tst.eval()
 
 
+
+def standarize_point(d, dim=128, flip = False):
+    if flip:
+        d = -d
+    return d/dim - 0.5
+
+def destandarize_point(d, dim=128, flip = False):
+    if flip:
+        d = -d
+    return dim*(d + 0.5)
+
+def generate_standard_elips(N_samples = 100, a= 1,b = 1):
+    radius = 0.25
+    center = 0
+    N_samples1 = int(N_samples/2 - 1)
+    N_samples2 = N_samples - N_samples1
+    x1 = torch.distributions.uniform.Uniform(center-radius,center + radius).sample([N_samples1])
+    x1_ordered = torch.sort(x1).values
+    y1 = center + b*torch.sqrt(radius**2 - ((x1_ordered-center)/a)**2)
+    x2 = torch.distributions.uniform.Uniform(center-radius,center + radius).sample([N_samples2])
+    x2_ordered = torch.sort(x2, descending=True).values
+    y2 = center - b*torch.sqrt(radius**2 - ((x2_ordered-center)/a)**2)
+    x = torch.cat([x1_ordered, x2_ordered])
+    y = torch.cat([y1, y2])
+    return x, y
+
+def transform_standard_points(Affine_mat, x,y):
+    XY = torch.ones([3,x.shape[0]])
+    XY[0,:]= x
+    XY[1,:]= y
+    XYt = torch.matmul(Affine_mat.to('cpu').detach(),XY)
+    xt0 = XYt[0]
+    yt0 = XYt[1]
+    return xt0, yt0
+
+def wrap_points(img, x_source, y_source, l=1):
+ for i in range(len(y_source)):
+   x0 = x_source[i].int()
+   y0 = y_source[i].int()
+   img[:,:,x0-l:x0+l,y0-l:y0+l] = 0
+ return img
+
+
 def wrap_imge_cropped(Affine_mtrx, source_img, dim1=224, dim2=128):
   source_img224 = torch.nn.ZeroPad2d(int((dim1-dim2)/2))(source_img)
   grd = torch.nn.functional.affine_grid(Affine_mtrx, size=source_img224.shape,align_corners=False)
@@ -45,7 +88,7 @@ def wrap_imge_cropped(Affine_mtrx, source_img, dim1=224, dim2=128):
   wrapped_img = torchvision.transforms.CenterCrop((dim2, dim2))(wrapped_img)
   return wrapped_img
 
-def imgnt_reg(img1,img2, model_selected):
+def imgnt_reg(img1,img2):#, model_selected):
     #fixed_images = np.empty((1, 128, 128, 3))
     #moving_images = np.empty((1, 128, 128, 3))
     # prepare inputs:
@@ -59,7 +102,21 @@ def imgnt_reg(img1,img2, model_selected):
     pred = IR_Model_tst(model_inputs)
     img_out = wrap_imge_cropped(pred['Affine_mtrx'], fixed_images, dim1=224, dim2=128)
     registered_img = torchvision.transforms.ToPILImage()(img_out[0])
-    return registered_img
+    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)
+        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
+        marked_image = torchvision.transforms.ToPILImage()(img_out2[0])
+    else:
+        marked_image = torchvision.transforms.ToPILImage()(torch.zeros(3,10,10))
+    return [registered_img, marked_image]
 
 
 
@@ -79,21 +136,25 @@ with gr.Blocks() as demo:
             type = "filepath",
             elem_id = "image-in",
         )
-    
+    '''
     model_list = gr.Dropdown(
             ["Additive_Recurence", "Rawblock"], label="Model", info="select a model"
-        )
+        )'''
     
-    out_image = gr.Image(label = "ٌRegistered image",
+    out_image1 = gr.Image(label = "ٌRegistered image",
             #source = "upload", 
             #type = "filepath",
             elem_id = "image-out"
         )
-    inputs = [image_1, image_2, model_list]
+    out_image2 = gr.Image(label = "ٌMarked image",
+            elem_id = "image-out"
+        )
+    inputs = [image_1, image_2]#, model_list]
+    out_image = [out_image1, out_image2]
     iface = gr.Interface(fn=imgnt_reg, inputs=inputs,outputs=out_image,
     title="Imagenet registration V2",
     description="Upload 2 images to generate a registered one:",
     examples=[["./examples/ex1.jpg","./examples/ex2.jpg"]],
                          )
 
-demo.queue().launch()
+demo.queue().launch(share=True)