Update app.py

app.py

@@ -2,9 +2,20 @@ import gradio as gr
 from ultralytics import YOLO
 import cv2
 import numpy as np
+import os
+
+# Ensure the model file is in the correct location
+model_path = "yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
+if not os.path.exists(model_path):
+    # Download the model file if it doesn't exist
+    import requests
+    model_url = "https://huggingface.co/DILHTWD/documentlayoutsegmentation_YOLOv8_ondoclaynet/resolve/main/yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
+    response = requests.get(model_url)
+    with open(model_path, "wb") as f:
+        f.write(response.content)
 
 # Load the document segmentation model
-docseg_model = YOLO("https://huggingface.co/DILHTWD/documentlayoutsegmentation_YOLOv8_ondoclaynet/blob/main/yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt")
+docseg_model = YOLO(model_path)
 
 def process_image(image):
     # Convert image to the format YOLO model expects