feat: add models

(cherry picked from commit 539fb282fd3e9e4030fc4bd3e867ba730702deec)

app.py

@@ -1,11 +1,12 @@
 import gradio as gr
+from rapid_undistorted.inference import InferenceEngine
 from rapid_unwrap.inference import DocUnwrapper
-
+import numpy as np
 # 初始化模型
-model_path = "models/uvdoc.onnx"
-doc_unwrapper = DocUnwrapper(model_path=model_path)
+engine = InferenceEngine()
 # 添加示例
 example_images = [
+    "images/demo.jpg",
     "images/demo1.jpg",
     "images/demo1.png",
     "images/demo2.png",
@@ -14,10 +15,25 @@ example_images = [
 ]
 
 
-def process_image(img_path):
-    unwrapped_img, elapse = doc_unwrapper(img_path)
+# 定义任务和模型选项
+tasks = {
+    "unwrap": ["UVDoc", None],
+    "unshadow": ["GCDnet",  None],
+    "unblur": ["OpenCvBilateral", "NAFDPM", None]
+}
+
+def process_image(img_path, unwrap_model, unshadow_model, unblur_model):
+    task_list = []
+    if unwrap_model:
+        task_list.append(("unwrap", unwrap_model))
+    if unshadow_model:
+        task_list.append(("unshadow", unshadow_model))
+    if unblur_model:
+        task_list.append(("unblur", unblur_model))
+
+    unwrapped_img, elapse = engine(img_path, task_list)
     print(f"doc unwrap elapse: {elapse}")
-    return unwrapped_img
+    return unwrapped_img.astype(np.uint8),elapse
 
 
 def main():
@@ -37,7 +53,7 @@ def main():
         }
     """) as demo:
         gr.HTML(
-            "<h1 style='text-align: center;'><a href='https://github.com/Joker1212/RapidUnWrap'>RapidUnWrap</a></h1>"
+            "<h1 style='text-align: center;'><a href='https://github.com/Joker1212/RapidUnWrap'>RapidUnDistort</a></h1>"
         )
         gr.HTML('''
                                             <div class="header-links">
@@ -60,12 +76,18 @@ def main():
                     outputs=img_input,
                     cache_examples=False
                 )
+
+                unwrap_dropdown = gr.Dropdown(label="Select Unwrap Model", choices=tasks["unwrap"], value="UVDoc")
+                unshadow_dropdown = gr.Dropdown(label="Select Unshadow Model", choices=tasks["unshadow"], value="GCDnet")
+                unblur_dropdown = gr.Dropdown(label="Select Unblur Model", choices=tasks["unblur"], value="OpenCvBilateral")
+
                 run_button = gr.Button("summit")
+
             with gr.Column(scale=2):  # 右边占2/3
                 output_image = gr.Image(label="output")
-
+                elapse_textbox = gr.Textbox(label="Elapsed Time", interactive=False)
         # 绑定按钮点击事件
-        run_button.click(fn=process_image, inputs=img_input, outputs=output_image)
+        run_button.click(fn=process_image, inputs=[img_input, unwrap_dropdown, unshadow_dropdown, unblur_dropdown],outputs=[output_image, elapse_textbox])
 
     # 启动应用
     demo.launch()

config.yaml

@@ -0,0 +1,38 @@
+tasks:
+  unwrap:
+    models:
+      - type: "UVDoc"
+        path: "models/uvdoc.onnx"
+        use_cuda: false
+        use_dml: false
+
+  unshadow:
+    models:
+      - type: "GCDnet"
+        sub_models:
+          - type: "GCDnet"
+            path: "models/gcnet.onnx"
+            use_cuda: false
+            use_dml: false
+          - type: "DRnet"
+            path: "models/drnet.onnx"
+            use_cuda: false
+            use_dml: false
+
+  binarize:
+    models:
+      - type: "UnetCnn"
+        path: "models/unetcnn.onnx"
+        use_cuda: false
+        use_dml: false
+
+  unblur:
+    models:
+      - type: "OpenCvBilateral"
+        path:
+      - type: "NAFDPM"
+        path: "models/nafdpm.onnx"
+        use_cuda: false
+        use_dml: false
+
+

requirements.txt

@@ -1,2 +1,2 @@
-rapid_unwrap
+rapid-undistorted
 gradio