import os
import sys
import tensorflow as tf
import numpy as np
from tensorflow.keras import layers, models
import requests
from PIL import Image
from io import BytesIO
import gradio as gr

#Класс для нейросети
class FoodQualityAI:
    def __init__(self):
        self.model = None
        self.image_size = (256, 256)
        self.train_dir = "./classes"

    def test(self, image_url):
        if self.model is None:
            raise Exception("Модель не загрузилась.")

        try:
            response = requests.get(image_url)
            response.raise_for_status()
        except Exception as ex:
            raise Exception("Ошибка при скачивании изображения: " + str(ex))
            
        try:
            img = Image.open(BytesIO(response.content)).convert("RGB")
        except Exception as e:
            raise Exception("Ошибка при обработке изображения: " + str(e))

        img = img.resize(self.image_size)
        img_array = np.array(img) / 255.0
        img_array = np.expand_dims(img_array, axis=0)

        prediction = self.model.predict(img_array)
        
        file_names = [f for f in os.listdir(self.train_dir) 
              if os.path.isfile(os.path.join(self.train_dir, f))]
        
        class_names = list(set([os.path.splitext(f)[0].split('_')[0] for f in file_names]))
            
        predicted_class = class_names[np.argmax(prediction)]
        confidence = float(np.max(prediction))

        return predicted_class, confidence
    #загрузка модели
    def load(self, filename='food_quality_model.h5'):
        try:
            self.model = models.load_model(filename)
        except Exception as e:
            raise Exception(f"Не удалсь загрузить модель ИИ {filename}: {str(e)}")


food_quality_ai = FoodQualityAI()

try:
    food_quality_ai.load('food_quality_model.h5')
    print("Model loaded successfully.")
except Exception as e:
    print("Error loading model:", e)
    sys.exit(1)
    
# Основная функция
def classify_image_url(image_url):
    try:
        predicted_class, confidence = food_quality_ai.test(image_url)
        return f"Модель думает что это: {predicted_class}\nС вероятностю в: {confidence:.2f}"
    except Exception as err:
        return f"Ошибка: {err}"
        
#Запуск интерфейса gradio
interface = gr.Interface(
    fn=classify_image_url,
    inputs=gr.Textbox(label="Ссылка на тестовую картинку", placeholder="Введите ссылку..."),
    outputs=gr.Textbox(label="Мнение модели"),
    title="ИИ для определения состояния еды.",
    description="Вставьте ссылку чтобы проверить как работает модель."
)

if __name__ == "__main__":
    interface.launch()