import os
import gradio as gr
from CoinCounter.model import predict
import matplotlib.pyplot as plt
import numpy as np

# Получаем все пути к изображениям в папке 'data'
image_paths = [os.path.join('data', filename) for filename in os.listdir('data') if filename.endswith(('.png', '.jpg', '.jpeg'))]

def process_image(image, conf, iou):
    # Сохраняем изображение во временный буфер для обработки
    image_path = "temp_image.jpg"
    image.save(image_path)
    
    # Выполняем предсказание с моделью с заданными параметрами
    results = predict(path=image_path, conf=conf, iou=iou)
    
    # Получаем изображение с размеченными объектами
    annotated_image = results.image(0)
    
    # Конвертируем результат в формат, подходящий для вывода
    annotated_image = np.array(annotated_image)
    
    # Получаем сумму денег
    total_sum = results.total(0)
    
    # Извлекаем количество монет из DataFrame
    df = results.df
    coin_counts = df.iloc[0, :-1]  # исключаем последнюю колонку 'total'
    
    # Формируем строку с количеством монет
    coin_info = "\n".join([f"{coin}: {int(count)}" for coin, count in zip(coin_counts.index, coin_counts) if count > 0])
    
    # Формируем итоговый текст
    result_text = f"Общая сумма на изображении: {total_sum:.2f} $\n\nКоличество объектов:\n{coin_info}"
    
    return annotated_image, result_text

# Определяем интерфейс Gradio
with gr.Blocks() as demo:
    gr.Markdown("## Coin Counter")
    
    with gr.Row():
        # Левый блок (ввод)
        with gr.Column():
            input_image = gr.Image(label="Загрузите изображение", type="pil")
            conf_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.3, label="Порог уверенности (conf)")
            iou_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.45, label="Порог перекрытия (IoU)")
        
        # Правый блок (вывод)
        with gr.Column():
            output_image = gr.Image(label="Размеченное изображение")
            output_text = gr.Textbox(label="Результат", lines=5)
    
    # Обработка изображения автоматически при изменении изображения или тумблеров
    input_image.change(process_image, inputs=[input_image, conf_slider, iou_slider], outputs=[output_image, output_text])
    conf_slider.change(process_image, inputs=[input_image, conf_slider, iou_slider], outputs=[output_image, output_text])
    iou_slider.change(process_image, inputs=[input_image, conf_slider, iou_slider], outputs=[output_image, output_text])
    
    # Добавление примеров из папки 'data'
    gr.Examples(examples=image_paths, inputs=input_image)

# Запускаем интерфейс
demo.launch()