Update app.py

app.py

@@ -12,6 +12,7 @@ from model import DocGeoNet
 from seg import U2NETP
 import glob
 
+
 warnings.filterwarnings('ignore')
 
 class Net(nn.Module):
@@ -52,18 +53,15 @@ def reload_rec_model(model, path=""):
         model.load_state_dict(model_dict)
         return model
 
-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)
-
-# Compile models (assuming PyTorch 2.0)
-net = torch.compile(net)
+def rec(input_image):
+    seg_model_path = './model_pretrained/preprocess.pth'
+    rec_model_path = './model_pretrained/DocGeoNet.pth'
 
-net.eval()
+    net = Net()
+    reload_rec_model(net.DocTr, rec_model_path)
+    reload_seg_model(net.msk, seg_model_path)
+    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))
@@ -80,7 +78,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)
@@ -92,3 +90,10 @@ 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()