Update app.py

app.py

@@ -12,7 +12,7 @@ from huggingface_hub import snapshot_download
 from visualization import visualize_bbox
 
 # == download weights ==
-model_dir = snapshot_download('juliozhao/DocLayout-YOLO-DocStructBench', local_dir='./models/DocLayout-YOLO-DocStructBench')
+model_dir = snapshot_download('juliozhao/DocLayout-YOLO-DocStructBench-imgsz1280-2501', local_dir='./models/DocLayout-YOLO-DocStructBench-imgsz1280-2501')
 # == select device ==
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
@@ -33,7 +33,7 @@ id_to_names = {
 def recognize_image(input_img, conf_threshold, iou_threshold):
     det_res = model.predict(
         input_img,
-        imgsz=1024,
+        imgsz=1280,
         conf=conf_threshold,
         device=device,
     )[0]
@@ -60,7 +60,14 @@ if __name__ == "__main__":
     # == load model ==
     from doclayout_yolo import YOLOv10
     print(f"Using device: {device}")
-    model = YOLOv10(os.path.join(os.path.dirname(__file__), "models", "DocLayout-YOLO-DocStructBench", "doclayout_yolo_docstructbench_imgsz1024.pt"))  # load an official model
+    model = YOLOv10(
+        os.path.join(
+            os.path.dirname(__file__), 
+            "models", 
+            "DocLayout-YOLO-DocStructBench-imgsz1280-2501", 
+            "doclayout_yolo_docstructbench_imgsz1280_2501.pt"
+        )
+    )  # load an official model
     
     with open("header.html", "r") as file:
         header = file.read()