app.py

"""import gradio as gr
import torch
from PIL import Image
from pathlib import Path
import numpy as np

# Load YOLO Model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='https://huggingface.co/ayoubsa/yolo_model/blob/main/best.pt')  # Replace 'model.pt' with your uploaded model's path

# Function to make predictions
def predict(image):
    # Convert input to PIL image if it's not
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    # Run inference
    results = model(image)
    
    # Extract predictions (bounding boxes, labels, confidence scores)
    predictions = results.pandas().xyxy[0]  # Pandas dataframe
    annotated_image = np.array(results.render()[0])  # Annotated image as NumPy array
    
    return annotated_image, predictions[['name', 'confidence']].to_dict(orient="records")

# Create Gradio Interface
image_input = gr.inputs.Image(type="numpy", label="Input Image")
output_image = gr.outputs.Image(type="numpy", label="Annotated Image")
output_text = gr.outputs.JSON(label="Predictions (Labels & Confidence)")

interface = gr.Interface(
    fn=predict,
    inputs=image_input,
    outputs=[output_image, output_text],
    title="YOLO Object Detection",
    description="Upload an image to detect objects using YOLO.",
    examples=["example1.jpg", "example2.jpg", "example3.jpg"]  # Provide paths to example images
)

interface.launch()"""


from datasets import load_dataset
import gradio as gr
import random
import numpy as np
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

# Load your dataset from Hugging Face
dataset = load_dataset("username/your_dataset_name")  # 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'
    random_examples = random.sample(images, num_examples)
    return random_examples

# 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 examples for Gradio app
example_images = get_random_examples(dataset, num_examples=3)

# Create the Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.inputs.Image(type="pil"),  # PIL Image for compatibility with YOLO
    outputs=gr.outputs.Image(type="pil"),
    examples=example_images  # Linking examples directly from Hugging Face dataset
)

# Launch the Gradio app
iface.launch()