Update app.py

app.py

@@ -1,53 +1,43 @@
 import gradio as gr
-import subprocess
-import os
+import torch
 from PIL import Image
 
-# YOLOv8 detection function
-def detect_objects(image):
-    # Save the uploaded image temporarily
-    input_image_path = "input.jpg"
-    image.save(input_image_path)
-
-    # Define the YOLOv8 command
-    output_dir = "./runs/detect"
-    os.makedirs(output_dir, exist_ok=True)
-
-    command = [
-        "yolo",  # YOLOv8 CLI command
-        "task=detect",  # Specify the task as detection
-        "mode=predict",  # Set mode to predict
-        f"model=best.pt",  # Path to your YOLOv8 model weights
-        f"source={input_image_path}",  # Input image path (uploaded by user)
-        f"project={output_dir}",  # Output directory
-        f"name=result",  # Subfolder name for results
-        "exist_ok=True"  # Allow overwriting existing results
-    ]
-
-    # Run YOLOv8 inference
-    try:
-        subprocess.run(command, check=True)
-    except Exception as e:
-        return f"Error during YOLO inference: {e}"
-
-    # Get the path to the output image
-    output_image_path = os.path.join(output_dir, "result", os.path.basename(input_image_path))
-
-    # Check if the image exists
-    if os.path.exists(output_image_path):
-        # Return the image directly to Gradio
-        return output_image_path  # Return file path for Gradio to display
-    else:
-        return "Error: Output image not found."
-
-# Create the Gradio interface
+# Define the function to load the YOLOv8 model and perform processing
+def process_image(image_path, model_path="waste-detection-yolov8/best_p6.pt"):
+  """
+  Processes an image using a YOLOv8 model and returns the processed image.
+
+  Args:
+      image_path (str): Path to the input image.
+      model_path (str, optional): Path to the YOLOv8 model weights file. Defaults to "waste-detection-yolov8/best_p6.pt".
+
+  Returns:
+      PIL.Image: The processed image.
+  """
+  # Load the YOLOv8 model from the specified path
+  model = torch.hub.load('ultralytics/yolov8n', 'custom', path=model_path)
+
+  # Read the input image
+  image = Image.open(image_path)
+
+  # Convert the image to a tensor
+  image = model(image) 
+
+  # Get the processed image from the results
+  processed_image = image.imgs[0] 
+
+  return processed_image
+
+# Define the Gradio interface
 interface = gr.Interface(
-    fn=detect_objects,  # The function to call when an image is uploaded
-    inputs=gr.Image(type="pil"),  # Accept images as input
-    outputs=gr.Image(type="filepath"),  # Return image file path for display
-    title="YOLOv8 Object Detection",
-    description="Upload an image of floating waste in water, and this app will detect it using YOLOv8."
+  fn=process_image,
+  inputs=gr.Image(label="Input Image", type="filepath"),
+  outputs="image",
+  title="Image Processing with YOLOv8n",
+  description="Upload an image to process it with the YOLOv8n model.",
+  thumbnail=None,
+  article="<p>This Gradio app allows you to upload an image and process it using a YOLOv8n model.</p>",
 )
 
-# Launch the Gradio app
-interface.launch(share=True)  # Use share=True to get a public URL
+# Launch the interface
+interface.launch(server_port=11111, server_name="localhost", enable_queue=True, allow_screenshot=False, allow_user_code=False)