Update app.py

app.py

@@ -12,7 +12,6 @@ from model import DocGeoNet
 from seg import U2NETP
 import glob
 
-
 warnings.filterwarnings('ignore')
 
 class Net(nn.Module):
@@ -53,15 +52,18 @@ def reload_rec_model(model, path=""):
         model.load_state_dict(model_dict)
         return model
 
-def rec(input_image):
-    seg_model_path = './model_pretrained/preprocess.pth'
-    rec_model_path = './model_pretrained/DocGeoNet.pth'
+net = Net()
+seg_model_path = './model_pretrained/preprocess.pth'
+rec_model_path = './model_pretrained/DocGeoNet.pth'
+reload_rec_model(net.DocTr, rec_model_path)
+reload_seg_model(net.msk, seg_model_path)
 
-    net = Net()
-    reload_rec_model(net.DocTr, rec_model_path)
-    reload_seg_model(net.msk, seg_model_path)
-    net.eval()
+# Compile models (assuming PyTorch 2.0)
+net = torch.compile(net)
 
+net.eval()
+
+def rec(input_image):
     im_ori = np.array(input_image)[:, :, :3] / 255.  # read image 0-255 to 0-1
     h, w, _ = im_ori.shape
     im = cv2.resize(im_ori, (256, 256))
@@ -78,7 +80,7 @@ def rec(input_image):
         bm1 = cv2.blur(bm1, (3, 3))
         lbl = torch.from_numpy(np.stack([bm0, bm1], axis=2)).unsqueeze(0)  # h * w * 2
         out = F.grid_sample(torch.from_numpy(im_ori).permute(2, 0, 1).unsqueeze(0).float(), lbl, align_corners=True)
-        img_rec = ((out[0] * 255).permute(1, 2, 0).numpy())[:,:,::-1].astype(np.uint8)
+        img_rec = ((out[0] * 255).permute(1, 2, 0).numpy())[:, :, ::-1].astype(np.uint8)
 
         # Convert from BGR to RGB
         img_rec = cv2.cvtColor(img_rec, cv2.COLOR_BGR2RGB)
@@ -90,10 +92,3 @@ demo_img_files = glob.glob('./distorted/*.[jJ][pP][gG]') + glob.glob('./distorte
 # Gradio Interface
 input_image = gr.inputs.Image()
 output_image = gr.outputs.Image(type='pil')
-
-
-
-iface = gr.Interface(fn=rec, inputs=input_image, outputs=output_image, title="DocGeoNet",examples=demo_img_files)
-
-#iface.launch(server_port=8821, server_name="0.0.0.0")
-iface.launch()