Update app.py

app.py

@@ -4,6 +4,7 @@
 import numpy as np
 import streamlit as st
 from PIL import Image
+import cv2
 
 from rapid_layout import RapidLayout, VisLayout
 
@@ -79,8 +80,11 @@ select_model = st.selectbox("选择版面分析模型：", model_types.keys())
 st.write("支持检测类型：")
 st.code(model_types[select_model], language="python")
 
+conf_threshold = st.slider("Confidence Threshold", 0, 1, 0.5)
+iou_threshold = st.slider("IoU Threshold", 0, 1, 0.5)
+
 with st.spinner(f"Downloading {select_model} model..."):
-    layout_engine = RapidLayout(model_type=select_model)
+    layout_engine = RapidLayout(model_type=select_model, conf_thres=conf_threshold, iou_thres=iou_threshold)
 
 st.write("请上传图像：")
 img_suffix = ["png", "jpg", "jpeg"]
@@ -92,7 +96,7 @@ if img_file_buffer:
     image = Image.open(img_file_buffer)
     img = np.array(image)
     with st.spinner("推理中...."):
-        boxes, scores, class_names, *elapse = layout_engine(img)
+        boxes, scores, class_names, *elapse = layout_engine(img, )
         ploted_img = VisLayout.draw_detections(img, boxes, scores, class_names)
 
     st.image(ploted_img, use_column_width=True)