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

from yolo import Yolo
model_link = "https://intranet-projects-files.s3.amazonaws.com/holbertonschool-ml/yolo.h5"

def download_model():
    if not os.path.exists('data'):
        os.makedirs('data') 
        
    if not os.path.exists('data/yolo.h5'):
        print("Downloading model...")
        r = requests.get(model_link)
        with open('data/yolo.h5', 'wb') as f:
            f.write(r.content)
        print("Model downnloaded")
    else:
        print("Model already exists locally.")
def run(img):
    np.random.seed(0)
    
    anchors = np.array([[[116, 90], [156, 198], [373, 326]],
                        [[30, 61], [62, 45], [59, 119]],
                        [[10, 13], [16, 30], [33, 23]]])
    yolo = Yolo('data/yolo.h5', 'coco_classes.txt', 0.6, 0.5, anchors)
    
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    boxes, box_classes, box_scores = yolo.predict_frame(img)
    
    for idx, box in enumerate(boxes):
        top_left_x = int(box[0])
        top_left_y = int(box[1])
        bottom_right_x = int(box[2])
        bottom_right_y = int(box[3])
        class_name = yolo.class_names[box_classes[idx]]
        score = box_scores[idx]
        color = (255, 0, 0)
        cv2.rectangle(img, (top_left_x, top_left_y),
                      (bottom_right_x, bottom_right_y),
                      color, 2)
        text = f"{class_name} {score:.2f}"
        cv2.putText(img, text, (top_left_x, top_left_y - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1,
                    cv2.LINE_AA)
    
    # Convert image back from RGB to BGR for displaying with OpenCV
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    
    return img

demo = gr.Interface(run, "image", "image")

if __name__ == "__main__":
    download_model()
    demo.launch()