Update app.py

app.py

@@ -1,43 +1,43 @@
 import gradio as gr
 import torch
 from PIL import Image
-
-# 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=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>",
+import torchvision.transforms as T
+
+# Load the trained model (YOLOv8n) with your weights
+model = torch.hub.load('ultralytics/yolov8', 'yolov8n')
+model.load_state_dict(torch.load("best_p6.pt"))
+model.eval()
+
+# Define the image transformation (if required, based on your dataset preprocessing)
+transform = T.Compose([T.ToTensor()])
+
+# Define the inference function
+def process_image(image):
+    # Convert the image to tensor and make inference
+    image_tensor = transform(image).unsqueeze(0)  # Add batch dimension
+    with torch.no_grad():
+        outputs = model(image_tensor)
+    
+    # Get the output image with bounding boxes (you can adjust this part based on your model's output)
+    result_image = outputs.render()[0]  # This will render bounding boxes on the image
+    
+    # Convert to PIL image for easy download
+    result_pil_image = Image.fromarray(result_image)
+    
+    # Save the output image for download
+    output_path = "/tmp/output_image.jpg"
+    result_pil_image.save(output_path)
+    
+    return output_path
+
+# Define Gradio interface
+iface = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="pil"),  # Image input from user
+    outputs=gr.File(label="Download Processed Image"),  # Provide the file output for download
+    title="Waste Detection",  # Interface title
+    description="Upload an image of floating waste, and the model will detect and label the objects in it."
 )
 
 # Launch the interface
-interface.launch(server_port=11111, server_name="localhost", enable_queue=True, allow_screenshot=False, allow_user_code=False)
+iface.launch()