import gradio as gr
from ultralytics import YOLO
from PIL import Image
import requests
import torch

# Periksa apakah GPU tersedia dan pilih perangkat GPU jika ada
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# Step 1: Download the YOLO model from Hugging Face
url = 'https://huggingface.co/Yudsky/pest-detection-yolo11/resolve/main/best.pt'
model_path = 'best.pt'

# Check if the model file already exists, if not, download it
if not os.path.exists(model_path):
    print("Downloading the model...")
    response = requests.get(url)
    with open(model_path, 'wb') as f:
        f.write(response.content)
    print("Download completed.")

# Step 2: Load the model using YOLO from ultralytics
# print("Loading the model...")
model = YOLO(model_path)
# print("Model loaded successfully.")

# Step 3: Define the prediction functions for images and videos
def predict_image(image):
    # Run inference
    results = model(image, device=device)
    # Plot results on the image
    annotated_image = results[0].plot()  # Get the annotated image with bounding boxes
    return Image.fromarray(annotated_image)

inputs_image = [
    gr.Image(type='filepath', label='input image')
]
outputs_image = [
    gr.Image(type='numpy', label='output image')
]

# Step 4: Define the Gradio Interface
interface_image = gr.Interface(
    fn=predict_image,
    inputs=inputs_image,
    outputs=outputs_image,
    title="Pest Detection",
    description="Upload an image and the model will detect pests."
)

# Step 5: Launch the interface
interface_image.launch()