import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import os
from io import BytesIO
from roboflow import Roboflow

roboflow_key = os.getenv("roboflow")
rf = Roboflow(api_key=roboflow_key)
project = rf.workspace('yudi-pratama-putra-rwuep').project("corn-pest-detection-2")
model = project.version(3).model


def predict_image(image, confidence, overlap):
    colors = [
        "Red",
        "Green",
        "Blue",
        "Yellow",
        "Cyan",
        "Magenta",
        "Orange",
        "Purple",
        "Brown",
        "Pink",
        "DarkRed",
        "Black",
        "White"
    ]

    pest_class = ['aphids', 'army worm', 'black cutworm', 'corn borer', 'grub', 'large cutworm', 'mole cricket', 'peach borer', 'potosiabre vitarsis', 'red spider', 'white margined moth', 'wireworm', 'yellow cutworm']

    prediction = model.predict(image, confidence=confidence, overlap=overlap).json()

    img = Image.open(image).convert("RGB")

    resize_img = 1
    
    img = img.resize((img.width * resize_img, img.height * resize_img))

    draw = ImageDraw.Draw(img)
    font = ImageFont.load_default(size=20)

    for result in prediction['predictions']:
        x0 = result['x'] - result['width'] / 2
        y0 = result['y'] - result['height'] / 2
        x1 = result['x'] + result['width'] / 2
        y1 = result['y'] + result['height'] / 2

        x0 *= resize_img
        y0 *= resize_img
        x1 *= resize_img
        y1 *= resize_img

        for i in range(len(pest_class)):
          if result['class'] == pest_class[i]:
              draw.rectangle([x0, y0, x1, y1], outline=colors[i], width=3)

              label = f"{result['class']} ({result['confidence']*100:.2f}%)"

              draw.text((x0, y0 - 25), label, fill=colors[i], font=font)

    if prediction['predictions'] == []:
        img = Image.open(image)

    return img

inputs_image = [
    gr.Image(type='filepath', label='input image'),
    gr.Slider(minimum=0, maximum=100, value=40, label='Confidence (%)'),
    gr.Slider(minimum=0, maximum=100, value=30, label='Overlap (%)')
]
outputs_image = [
    gr.Image(type='numpy', label='output image')
]

interface_image = gr.Interface(
    fn=predict_image,
    inputs=inputs_image,
    outputs=outputs_image,
    title="Corn Pest Detection",
    description=(
        "Upload an image and the model will detect pests.\n\n"
        "Model detected: aphids, army worm, black cutworm, corn borer, grub, large cutworm, mole cricket, peach borer, "
        "potosiabre vitarsis, red spider, white margined moth, wireworm, yellow cutworm.\n\n"
        "Confidence: The higher the confidence value, the more certain the model is about the detected object being correct. "
        "For example, a higher confidence threshold will filter out less certain predictions.\n\n"
        "Overlap: The higher the accepted overlap value, the more predictions are allowed, even if they overlap with each other. "
        "A higher overlap value can help detect multiple objects that are close together."
    )
)

interface_image.launch(share=True)