import os
import gradio as gr
import ultralytics
from IPython import display
from ultralytics import YOLO
import cv2
import numpy as np
from sklearn.cluster import KMeans
import webcolors

CURRENT_DIR = os.getcwd()

ultralytics.checks()
display.clear_output()

model = YOLO('yolov8s.pt')

def load_model(model_path):
    model = YOLO(model_path)
    return model

def detect_cats(model, image):
    results = model(image)
    return results

def crop_regions(image, detections):
    cropped_regions = []
    for det in detections:
        x1, y1, x2, y2 = map(int, det[:4])
        cropped_region = image[y1:y2, x1:x2]
        cropped_regions.append(cropped_region)
    return cropped_regions

def get_dominant_color(image, k=1):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = image.reshape((-1, 3))

    kmeans = KMeans(n_clusters=k,n_init='auto')
    kmeans.fit(image)

    dominant_color = kmeans.cluster_centers_[0]
    return dominant_color

def get_color_name(rgb_color):
    try:
        closest_name = webcolors.rgb_to_name(rgb_color)
    except ValueError:
        min_colors = {}
        for key, name in webcolors.CSS3_HEX_TO_NAMES.items():
            r_c, g_c, b_c = webcolors.hex_to_rgb(key)
            rd = (r_c - rgb_color[0]) ** 2
            gd = (g_c - rgb_color[1]) ** 2
            bd = (b_c - rgb_color[2]) ** 2
            min_colors[(rd + gd + bd)] = name
        closest_name = min_colors[min(min_colors.keys())]
    return closest_name

def draw_results(image, detections, colors_list):
    for det, color in zip(detections, colors_list):
        x1, y1, x2, y2 = map(int, det[:4])
        rgb_color = (int(color[0]), int(color[1]), int(color[2]))
        color_name = get_color_name(rgb_color)
        hex_color = "#{:02x}{:02x}{:02x}".format(rgb_color[0], rgb_color[1], rgb_color[2])
        color_info = f"{color_name}"

        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(image, color_info, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    return image, color_info

def main(image_path):
    #model = load_model(model_path)
    image = cv2.imread(image_path)

    results = detect_cats(model, image)
    detections = results[0].boxes.data.cpu().numpy()  # Get bounding boxes

    cropped_regions = crop_regions(image, detections)
    colors_list = [get_dominant_color(region) for region in cropped_regions]

    output_image, color = draw_results(image, detections, colors_list)

    return output_image, color

def result(input_image):
    input_image_path = f"{CURRENT_DIR}/input_image.jpg"
    cv2.imwrite(input_image_path, input_image)
    img = input_image_path
    processed_img, cat_color = main(img)

    return processed_img, cat_color

demo = gr.Interface(
    fn=result,
    inputs=gr.Image(),
    outputs=[gr.Image(),gr.Text()]
)

demo.launch()