Update app.py

app.py

@@ -1,9 +1,12 @@
-import spaces
-import datetime
+import gradio as gr
+from ultralytics import YOLO
+import cv2
+import numpy as np
 import os
-import subprocess
+import requests
 import torch
-import gradio as gr
+import datetime
+import subprocess
 
 CUSTOM_CSS = """
 #output_box textarea {
@@ -11,6 +14,34 @@ CUSTOM_CSS = """
 }
 """
 
+# 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
+    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(model_path)
+
+def process_image(image):
+    # Convert image to the format YOLO model expects
+    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    results = docseg_model(image)
+
+    # Extract annotated image from results
+    annotated_img = results[0].plot()
+    annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
+
+    # Prepare detected areas and labels as text output
+    detected_areas_labels = "\n".join(
+        [f"{box.label}: {box.conf:.2f}" for box in results[0].boxes]
+    )
+
+    return annotated_img, detected_areas_labels
+
 zero = torch.Tensor([0]).cuda()
 print(zero.device) # <-- 'cpu' 🤔
 
@@ -49,4 +80,22 @@ with gr.Blocks(css=CUSTOM_CSS) as demo:
 
     check.change(run, inputs=[check], outputs=output, every=1)
 
+# Define the Gradio interface
+with gr.Blocks() as interface:
+    gr.Markdown("### Document Segmentation using YOLOv8")
+    input_image = gr.Image(type="pil", label="Input Image")
+    output_image = gr.Image(type="pil", label="Annotated Image")
+    output_text = gr.Textbox(label="Detected Areas and Labels")
+
+    gr.Button("Run").click(
+        fn=process_image,
+        inputs=input_image,
+        outputs=[output_image, output_text]
+    )
+
 demo.queue().launch(show_api=False)
+interface.launch()
+
+if __name__ == "__main__":
+    demo.launch()
+    interface.launch()