Update app.py

app.py

@@ -5,6 +5,14 @@ import numpy as np
 import os
 import requests
 import torch
+import huggingface_hub
+
+# Initialize ZeroGPU 
+zero_gpu_is_available = huggingface_hub.utils.is_google_colab() or huggingface_hub.utils.is_notebook()
+if zero_gpu_is_available:
+    from accelerate import Accelerator
+    accelerator = Accelerator()
+
 
 # Load the model file
 model_path = "yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
@@ -16,41 +24,35 @@ if not os.path.exists(model_path):
         f.write(response.content)
 
 # Load the document segmentation model
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-docseg_model = YOLO(model_path)  # Remove .to(device) to let ultralytics auto-detect
+docseg_model = YOLO(model_path) 
+
+if zero_gpu_is_available:
+    docseg_model.to(accelerator.device)  # Put the model on the accelerator's device.
+
 
 def process_image(image):
     try:
         # Convert image to the format YOLO model expects
         image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-        results = docseg_model.predict(image)  # Use predict for inference
+
+        # If Zero GPU, move image to accelerator
+        if zero_gpu_is_available:
+            image = torch.from_numpy(image).to(accelerator.device)
+
+        results = docseg_model.predict(image)
         result = results[0]  # Get the first (and usually only) result
         
         # Extract annotated image from results
-        annotated_img = result.plot()  # Simplified plotting
+        annotated_img = result.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.upper()}: {box.conf:.2f}" for box in result.boxes]  # Uppercase labels
+            [f"{box.label.upper()}: {box.conf:.2f}" for box in result.boxes]
         )
     except Exception as e:
         return None, f"Error during processing: {e}"  # Error handling
 
     return annotated_img, detected_areas_labels
 
-# 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]
-    )
-
-# Launch the interface (remove the conditional launch)
-interface.launch(share=True)  # Allow sharing for easier debugging
+# The rest of the code remains the same (Gradio interface)