Update app.py

app.py

@@ -36,7 +36,7 @@ interface = gr.Interface(
     examples=["example1.jpg", "example2.jpg", "example3.jpg"]  # Provide paths to example images
 )
 
-interface.launch()"""
+interface.launch()
 
 
 from datasets import load_dataset
@@ -47,14 +47,14 @@ from PIL import Image
 import torch
 
 # Load the YOLO model
-model = torch.hub.load('ultralytics/yolov5', 'custom', path='model.pt')  # Replace with your uploaded model's path
+model = torch.hub.load('ultralytics/yolov5', 'custom', path='https://huggingface.co/ayoubsa/yolo_model/blob/main/best.pt')  # Replace with your uploaded model's path
 
 # Load your dataset from Hugging Face
-dataset = load_dataset("username/your_dataset_name")  # Replace with your dataset's repository name on Hugging Face
+dataset = load_dataset("https://huggingface.co/datasets/ayoubsa/Sign_Road_Detection_Dataset/tree/main")  # Replace with your dataset's repository name on Hugging Face
 
 # Function to get random examples from the dataset
 def get_random_examples(dataset, num_examples=3):
-    images = dataset['train'][:]['image']  # Assuming the images are in the 'train' split and column 'image'
+    images = dataset['test'][:]['image']  # Assuming the images are in the 'train' split and column 'image'
     random_examples = random.sample(images, num_examples)
     return random_examples
 
@@ -77,5 +77,55 @@ iface = gr.Interface(
 )
 
 # Launch the Gradio app
-iface.launch()
+iface.launch()"""
+
+import gradio as gr
+import torch
+from PIL import Image
+import zipfile
+import os
+import random
+
+# Define the paths for the model and dataset
+MODEL_PATH = 'https://huggingface.co/ayoubsa/yolo_model/resolve/main/best.pt'  # YOLO model file
+DATASET_PATH = 'test.zip'  # The name of the uploaded test dataset zip file
 
+# Load the YOLO model
+model = torch.hub.load('ultralytics/yolov5', 'custom', path=MODEL_PATH)
+
+# Unzip the dataset
+if not os.path.exists("unzipped_test"):
+    with zipfile.ZipFile(DATASET_PATH, 'r') as zip_ref:
+        zip_ref.extractall("unzipped_test")  # Extract images to this folder
+
+# Get all image paths from the unzipped folder
+image_folder = "unzipped_test"
+all_images = [os.path.join(image_folder, f) for f in os.listdir(image_folder) if f.endswith(('.jpg', '.png', '.jpeg'))]
+
+# Function to get random examples
+def get_random_examples(num_examples=3):
+    if len(all_images) >= num_examples:
+        return random.sample(all_images, num_examples)
+    else:
+        return all_images  # Return whatever is available if less than required
+
+# Define the prediction function
+def predict(image):
+    results = model(image)  # Perform object detection using YOLO
+    results.render()  # Render bounding boxes on the image
+    output_image = Image.fromarray(results.imgs[0])  # Convert to PIL image
+    return output_image
+
+# Get example images
+example_images = get_random_examples(num_examples=3)
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.inputs.Image(type="pil", label="Upload an Image"),
+    outputs=gr.outputs.Image(type="pil", label="Predicted Image with Bounding Boxes"),
+    examples=example_images  # Link the example images
+)
+
+# Launch the Gradio app
+iface.launch()