import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img, img_to_array
import autokeras as ak
import pickle
import gradio as gr
import cv2
import io
import zipfile
import base64
#Загружаем модель
loaded_model = load_model('best_model', custom_objects=ak.CUSTOM_OBJECTS)
# Загружаем label_encoder
with open('label_encoder.pkl', 'rb') as le_file:
    label_encoder = pickle.load(le_file)
def translator(Defect):
    translation_dict = {
        "Crazing": "Трещины",
        "Inclusion": "Вкрапления",
        "Patches": "Пятна",
        "Pitted": "Рябь",
        "Rolled": "Замятие",
        "Scratches": "Царапины"
    }
    return translation_dict.get(Defect)
def predict(zip_bytes):
    predictions_html = "<table><tr><th>Файл</th><th>Класс</th><th>Точность предсказания</th></tr>"
    with zipfile.ZipFile(zip_bytes, 'r') as zip_ref:
        for filename in zip_ref.namelist():
            if filename.lower().endswith(('.png', '.jpg', '.jpeg','.bmp')):
                with zip_ref.open(filename) as file:
                    img_bytes = file.read()
                    img_np = np.frombuffer(img_bytes, np.uint8)
                    img = cv2.imdecode(img_np, cv2.IMREAD_COLOR)
                    img_resized = cv2.resize(img, (200, 200))
                    img_array = np.expand_dims(img_resized, axis=0)
                    prediction = loaded_model.predict(img_array)
                    predicted_class = np.argmax(prediction)
                    decoded_class = label_encoder.inverse_transform([predicted_class])[0]
                    predicted_class_translation = translator(decoded_class)
                    accuracy = round(((np.max(prediction).item())*100),2)
                    predictions_html += f"<tr><th>{filename}</th><th>{predicted_class_translation}</th><th>{accuracy}%</th></tr>"
    predictions_html += "</table>"
    return predictions_html
interface = gr.Interface(
  fn=predict,
  inputs=gr.File(label="Выберите zip-файл (только с изображениями)"),
  outputs=gr.HTML(label="Предсказанные классы в таблице"),
  title="Нейронная сеть для классификации брака на металлопрокате"
) 

interface.launch(share=True)