import gradio as gr
from PIL import Image
from ultralytics import YOLO
import asyncio
import aiohttp
import json
import logging
import cv2
from io import BytesIO
import numpy as np

logging.basicConfig(level=logging.INFO)

model_detection = YOLO('./detection_best.pt')
model_classification = YOLO('./classification_best.pt')

async def fetch_image(url):
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            if response.status == 200:
                image_data = await response.read()
                print("1")
                return Image.open(BytesIO(image_data))
            else:
                logging.error(f"Failed to load image from {url}")
                return None

async def detect_objects(images):
    classes = {2: "Positive", 1: "Negative"}
    results = []

    processed_images = [cv2.resize(np.array(image), (640, 640)) for image in images]

    results_detection = model_detection(processed_images)
    print(results_detection)
    for image, detection in zip(processed_images, results_detection):
        names = []
        if detection:
            i = 0
            for box in detection.boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                cropped_img = image[y1:y2, x1:x2]
                resized_img = cv2.resize(cropped_img, (640, 640))
                resized_img = cv2.cvtColor(resized_img, cv2.COLOR_BGR2RGB)
                cv2.imwrite(f'resized_{i}.png',resized_img)
                results_classification = model_classification.predict(resized_img)
                i+=1
                if results_classification:
                    top1_class = results_classification[0].probs.top1
                    names.append(classes[top1_class])
        if not names:
            names.append("None")
        results.append(names)
    return results

def create_solutions(image_urls, names, file_ids):
    return [
        {"image": url, "answer": name, "qcUserId": None, "normalfileID": file_id}
        for url, name, file_id in zip(image_urls, names, file_ids)
    ]

async def process_images_async(params):
    try:
        params = json.loads(params)
    except json.JSONDecodeError as e:
        logging.error(f"Invalid JSON input: {e}")
        return {"error": f"Invalid JSON input: {e}"}

    image_urls = params.get("urls", [])
    file_ids = params.get("normalfileID", [None] * len(image_urls))

    if not image_urls:
        logging.error("Missing required parameters: 'urls'")
        return {"error": "Missing required parameters: 'urls'"}

    images = await asyncio.gather(*[fetch_image(url) for url in image_urls])

    if not any(images):
        logging.error("No valid images were loaded.")
        return {"error": "No valid images were loaded."}

    names = await detect_objects(images)
    solutions = create_solutions(image_urls, names, file_ids)

    return json.dumps({"solutions": solutions})

def process_images(params):
    return asyncio.run(process_images_async(params))

inputt = gr.Textbox(label="Parameters (JSON format) Eg. img_url:['','']")
outputs = gr.JSON()

application = gr.Interface(fn=process_images, inputs=inputt, outputs=outputs, title="ART +ve -ve Detection with API Integration")
application.launch()