Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -2,9 +2,20 @@ import gradio as gr
|
|
2 |
from ultralytics import YOLO
|
3 |
import cv2
|
4 |
import numpy as np
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
5 |
|
6 |
# Load the document segmentation model
|
7 |
-
docseg_model = YOLO(
|
8 |
|
9 |
def process_image(image):
|
10 |
# Convert image to the format YOLO model expects
|
|
|
2 |
from ultralytics import YOLO
|
3 |
import cv2
|
4 |
import numpy as np
|
5 |
+
import os
|
6 |
+
|
7 |
+
# Ensure the model file is in the correct location
|
8 |
+
model_path = "yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
|
9 |
+
if not os.path.exists(model_path):
|
10 |
+
# Download the model file if it doesn't exist
|
11 |
+
import requests
|
12 |
+
model_url = "https://huggingface.co/DILHTWD/documentlayoutsegmentation_YOLOv8_ondoclaynet/resolve/main/yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
|
13 |
+
response = requests.get(model_url)
|
14 |
+
with open(model_path, "wb") as f:
|
15 |
+
f.write(response.content)
|
16 |
|
17 |
# Load the document segmentation model
|
18 |
+
docseg_model = YOLO(model_path)
|
19 |
|
20 |
def process_image(image):
|
21 |
# Convert image to the format YOLO model expects
|