Update app.py

app.py

@@ -66,8 +66,6 @@ class GeoTr_Seg(nn.Module):
         return bm
 
 
-        
-
 # Initialize models
 GeoTr_Seg_model = GeoTr_Seg()
 #IllTr_model = IllTr()
@@ -83,27 +81,27 @@ GeoTr_Seg_model = torch.compile(GeoTr_Seg_model)
 
 def process_image(input_image):
     GeoTr_Seg_model.eval()
-    #IllTr_model.eval()
 
     im_ori = np.array(input_image)[:, :, :3] / 255.
     h, w, _ = im_ori.shape
-    im = cv2.resize(im_ori, (288, 288))
+    new_height = int(h * (288 / w))
+    im = cv2.resize(im_ori, (288, new_height))
     im = im.transpose(2, 0, 1)
     im = torch.from_numpy(im).float().unsqueeze(0)
 
     with torch.no_grad():
         bm = GeoTr_Seg_model(im)
         bm = bm.cpu()
-        bm0 = cv2.resize(bm[0, 0].numpy(), (w, h))
-        bm1 = cv2.resize(bm[0, 1].numpy(), (w, h))
+        bm0 = cv2.resize(bm[0, 0].numpy(), (288, new_height))
+        bm1 = cv2.resize(bm[0, 1].numpy(), (288, new_height))
         bm0 = cv2.blur(bm0, (3, 3))
         bm1 = cv2.blur(bm1, (3, 3))
         lbl = torch.from_numpy(np.stack([bm0, bm1], axis=2)).unsqueeze(0)
 
         out = F.grid_sample(torch.from_numpy(im_ori).permute(2, 0, 1).unsqueeze(0).float(), lbl, align_corners=True)
         img_geo = ((out[0] * 255).permute(1, 2, 0).numpy()).astype(np.uint8)
-        
-        ill_rec=False
+
+        ill_rec = False
 
         if ill_rec:
             img_ill = rec_ill(IllTr_model, img_geo)
@@ -111,6 +109,7 @@ def process_image(input_image):
         else:
             return Image.fromarray(img_geo)
 
+
 # Define Gradio interface
 input_image = gr.inputs.Image()
 output_image = gr.outputs.Image(type='pil')