import gradio as gr
import requests
import cv2
import numpy as np
import os

API_URL = "https://detect.roboflow.com"
API_KEY = os.getenv("API_KEY")
MODEL_ID = "biofarma-x-mit-hacking-medicine-hackathon/1"

def annotate_image(image):
    # Convert the input image to a format suitable for OpenCV
    image = np.array(image)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Save the input image
    cv2.imwrite("input_image.jpg", image)
    
    # Prepare the request
    url = f"{API_URL}/{MODEL_ID}?api_key={API_KEY}"
    with open("input_image.jpg", "rb") as file:
        response = requests.post(url, files={"file": file})
    
    result = response.json()
    
    # Annotate the image
    for prediction in result['predictions']:
        x, y, width, height = prediction['x'], prediction['y'], prediction['width'], prediction['height']
        left = int(x - width / 2)
        top = int(y - height / 2)
        right = int(x + width / 2)
        bottom = int(y + height / 2)
        cv2.rectangle(image, (left, top), (right, bottom), (0, 0, 255), 2)
        cv2.putText(image, f"{prediction['confidence']:.2f}", (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    # Convert the image back to RGB for display in Gradio
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image

# Create the Gradio interface
iface = gr.Interface(
    fn=annotate_image,
    inputs=gr.Image(type="pil"),
    outputs="image",
    title="TB-Bacillus Detection",
    description="Upload an image to get annotated results from the model."
)

# Launch the app
iface.launch()