Create app.py

app.py

@@ -0,0 +1,52 @@
+import gradio as gr
+import cv2
+from model_fasterrcnn import LDRNet_fasterrcnn
+import torch
+import numpy as np
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = LDRNet_fasterrcnn()
+model.load_state_dict(torch.load("last_0603_92.pth", map_location=device))
+model.eval()
+inputs = [
+    gr.Image(type="filepath", label="Input Image"),]
+outputs = gr.Image(type="filepath", label="Output Image")
+def PerspectiveTransform(img,coord):
+    # img = cv2.resize(img, (224, 224))
+    height,width,_ = img.shape
+    src_points = np.float32([[coord[0], coord[1]],[coord[2], coord[3]],[coord[4], coord[5]],[coord[6], coord[7]]])
+    dst_points = np.float32([[0, 0],[width, 0],[width, height],[0, height]]) 
+    transformation_matrix = cv2.getPerspectiveTransform(src_points, dst_points)
+    transformed_img = cv2.warpPerspective(img, transformation_matrix, (width, height))
+    return transformed_img
+def detect(image):
+    image = cv2.imread(image)
+    original_size = image.shape[:2]
+    img = cv2.resize(image, (224, 224))
+    img_show = np.copy(image)
+    img = torch.tensor(img.astype(np.float32)).permute(2, 0, 1).unsqueeze(0)
+    img = (img - 0.5 * 255) / (0.5 * 255)
+    result = model(img)
+    results = result[0].detach().numpy()[0]
+    coord = results[0:8]
+    lines = result[1].detach().numpy()[0]
+       
+    try:
+        coord = [int(x * 224) for x in coord]
+        lines = [int(x * 224) for x in lines]
+        img_show = cv2.cvtColor(img_show, cv2.COLOR_BGR2RGB)
+        coord = np.array(coord,dtype=np.int32)
+        coord = coord.reshape(4,2)
+        coord = np.floor(coord / np.array([224/original_size[1], 224/original_size[0]])).astype(np.int32)
+        coord = coord.reshape(1,8)[0]
+        origin = PerspectiveTransform(img_show,coord)
+    except Exception as e: print(e)
+        
+    return origin
+title = "Scan Documment 📄"
+demo_app = gr.Interface(
+    fn=detect,
+    inputs=inputs,
+    outputs=outputs,
+)
+demo_app.launch(debug=True)