import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
import os

# Завантаження моделі
model = load_model('model.keras')

# Отримання класів з файлу (якщо вони були збережені)
try:
    with open('class_names.txt', 'r') as f:
        class_names = f.read().splitlines()
except:
    # За замовчуванням використовуємо числові імена класів
    num_classes = model.output_shape[-1]
    class_names = [f"Class {i}" for i in range(num_classes)]

# Визначення функції передбачення
def predict(image):
    """
    Виконує передбачення класу за допомогою навченої моделі
    
    Args:
        image: Вхідне зображення
    
    Returns:
        dict: Словник з ймовірностями для кожного класу
    """
    # Отримання форми вхідних даних з моделі
    input_shape = model.input_shape[1:]
    
    # Перевірка кількості каналів
    if len(input_shape) == 3:  # Для зображень (CNN)
        # Масштабуємо зображення до потрібного розміру
        resized_image = tf.image.resize(image, [input_shape[0], input_shape[1]])
        
        # Перетворення на одноканальне зображення, якщо потрібно
        if input_shape[2] == 1 and resized_image.shape[-1] == 3:
            resized_image = tf.image.rgb_to_grayscale(resized_image)
        
        # Нормалізація та додавання розмірності пакету
        img_array = tf.expand_dims(resized_image, 0) / 255.0
    else:  # Для плоских даних
        # Розгортаємо зображення в одновимірний вектор
        img_array = tf.reshape(image, (1, -1)) / 255.0
    
    # Передбачення
    predictions = model.predict(img_array)
    
    # Формування результату
    result = {class_names[i]: float(predictions[0][i]) for i in range(len(class_names))}
    return result

# Створення інтерфейсу Gradio
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=len(class_names)),
    title="Модель багатокласової класифікації",
    description="Завантажте зображення для класифікації."
)

# Запуск інтерфейсу
interface.launch()