Update app.py

app.py

@@ -1,62 +1,54 @@
 import gradio as gr
 from ultralytics import YOLO
-import cv2
-import numpy as np
 import os
-import requests
-import torch
-import huggingface_hub
-from accelerate import Accelerator
-from huggingface_hub import notebook_login  # Added this for HF login
-from huggingface_hub.utils import HfHubHTTPError  # Added this to catch HF login errors
-# Initialize Hugging Face Hub login
-notebook_login()
-# Initialize Accelerator
-accelerator = Accelerator()
-
-
-# Load the model file
-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"
-    try:
-        response = requests.get(model_url)
-        with open(model_path, "wb") as f:
-            f.write(response.content)
-    except HfHubHTTPError as e:
-        if e.response.status_code == 401:
-            print("Authentication error. Please login to Hugging Face Hub.")
-        else:
-            raise e
-# Load the document segmentation model
-docseg_model = YOLO(model_path) 
 
+# Load pre-trained YOLOv8 models
+docseg_model1 = YOLO("yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt")
+docseg_model2 = YOLO("path/to/your/second/model.pt")  # Replace with your second model's path
 
-docseg_model = accelerator.prepare(docseg_model)
+# Available models
+MODELS = {
+    "DocLayNet YOLOv8": docseg_model1,
+    # "Your Second Model": docseg_model2  # Uncomment and add more as needed
+}
 
-def process_image(image):
+def process_image(image, model_name):
     try:
-        # Convert image to the format YOLO model expects
-        image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-
-        # Move image to accelerator
-        image = torch.from_numpy(image).to(accelerator.device)
+        # Select the model
+        model = MODELS[model_name]
 
-        results = docseg_model.predict(image)
-        result = results[0]  # Get the first (and usually only) result
+        # Process the image
+        results = model(source=image, save=False, show_labels=True, show_conf=True, show_boxes=True)
+        result = results[0]
         
-        # Extract annotated image from results
-        annotated_img = result.plot() 
-        annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
-
-        # Prepare detected areas and labels as text output
+        # Extract the annotated image and the labels/confidence scores
+        annotated_image = result.plot()
         detected_areas_labels = "\n".join(
             [f"{box.label.upper()}: {box.conf:.2f}" for box in result.boxes]
         )
+
+        return annotated_image, detected_areas_labels
     except Exception as e:
-        return None, f"Error during processing: {e}"  # Error handling
+        return None, f"Error processing image: {e}"
+    
+
+
+# Create the Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Document Segmentation Demo")
+    
+    # Input Components
+    with gr.Row():
+        input_image = gr.Image(type="pil", label="Upload Image")
+        model_dropdown = gr.Dropdown(list(MODELS.keys()), label="Select Model", value=list(MODELS.keys())[0])
+
+    # Output Components
+    output_image = gr.Image(type="pil", label="Annotated Image")
+    output_text = gr.Textbox(label="Detected Areas and Labels")
 
-    return annotated_img, detected_areas_labels
+    # Button to trigger inference
+    btn = gr.Button("Run Document Segmentation")
+    btn.click(fn=process_image, inputs=[input_image, model_dropdown], outputs=[output_image, output_text])
 
-# The rest of the code remains the same (Gradio interface) 
+# Launch the demo
+demo.launch()