Update app.py

app.py

@@ -1,39 +1,62 @@
 import gradio as gr
-
 import os
-
 #os.system('pip install torch  -q')
 
 import sys
 import numpy as np
-import torch
 from PIL import Image
-
+import itertools
+import glob
+import random
+import torch
+import torchvision
+import torchvision.transforms as transforms
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
 from codes import *
 
 ## Print samples
-vxm_model_loaded = create_model(dim = 128)
-vxm_model_loaded.load_weights('./modelbest/')
 
-vxm_model_loaded_affine = create_model(dim = 128)
-vxm_model_loaded_affine.load_weights('./modelbest/')
+from torchvision.models import resnet18
 
 
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+file_savingfolder = './modelbest/'
+ext = '_bestVal'
+core_model_tst = resnet18(pretrained=True)
+core_model_tst.fc = Identity()
+core_model_tst.load_state_dict(torch.load(file_savingfolder+'core_model'+ext+'.pth'))
+core_model_tst.to(device)
+IR_Model_tst = Build_IRmodel_Resnet(core_model_tst, registration_method)
+IR_Model_tst.load_state_dict(torch.load(file_savingfolder+'IR_Model'+ext+'.pth'))
+IR_Model_tst.to(device)
+IR_Model_tst.eval()
+
+
+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)
+  wrapped_img = torch.nn.functional.grid_sample(source_img224, grid=grd,
+                                                mode='bilinear', padding_mode='zeros', align_corners=False)
+  wrapped_img = torchvision.transforms.CenterCrop((dim2, dim2))(wrapped_img)
+  return wrapped_img
 
 def imgnt_reg(img1,img2, model_selected):
-    fixed_images = np.empty((1, 128, 128, 3))
-    moving_images = np.empty((1, 128, 128, 3))
+    #fixed_images = np.empty((1, 128, 128, 3))
+    #moving_images = np.empty((1, 128, 128, 3))
     # prepare inputs:
-    fixed_images[0] = preprocess_image(img1, dim = 128)
-    moving_images[0] = preprocess_image(img2, dim = 128)
-    
-    model_inputs = [moving_images/255, fixed_images/255]
-    if model_selected == "Imagenet-wild":
-        Y_predict = vxm_model_loaded(model_inputs)
-    else:
-        Y_predict = vxm_model_loaded_affine(model_inputs)
-    img_out = Y_predict[0][0].numpy()*255
-    registered_img = Image.fromarray(img_out.astype(np.uint8))
+    fixed_images = preprocess_image(img1, dim = 128)
+    moving_images = preprocess_image(img2, dim = 128)
+    M_i = torch.normal(torch.zeros([2,3]), torch.ones([2,3]))
+    model_inputs = {'source':fixed_images,
+                'target':moving_images,
+                'M_i' :M_i.usqueeze(0) }
+    #[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)
+    registered_img = torchvision.transforms.ToPILImage()(img_out[0])
     return registered_img
 
 
@@ -66,9 +89,9 @@ with gr.Blocks() as demo:
         )
     inputs = [image_1, image_2, model_list]
     iface = gr.Interface(fn=imgnt_reg, inputs=inputs,outputs=out_image,
-    title="Imagenet registration V1",
+    title="Imagenet registration V2",
     description="Upload 2 images to generate a registered one:",
     examples=[["./examples/ex1.jpg","./examples/ex2.jpg"]],
                          )
 
-demo.queue(default_enabled = True).launch(debug = True)
+demo.queue(default_enabled = True).launch(debug = True)