import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image

# =================================================================
# Модельді жүктеу үшін жолды көрсету
# СЕНІҢ МОДЕЛЬ АТАУЫҢ 'model.h5' БОЛУЫ КЕРЕК
# =================================================================
MODEL_PATH = 'model.h5'
try:
model = tf.keras.models.load_model(MODEL_PATH)
MODEL_LOADED = True
except Exception as e:
print(f"Модельді жүктеу қатесі: {e}")
# Егер модель жүктелмесе, жалған (dummy) функцияны қолданамыз
MODEL_LOADED = False

# Класс атаулары (Colab-тағыдай)
CLASS_NAMES = ['Empty', 'Full']
IMG_SIZE = (224, 224)

def classify_trash_bin(image):
# Егер модель жүктелмесе, қате туралы хабар береміз
if not MODEL_LOADED:
return {"Модель жүктелмеді": 1.0}

# Суретті өңдеу
if image is None:
return {CLASS_NAMES[0]: 0.5, CLASS_NAMES[1]: 0.5} # Бос нәтиже

img = Image.fromarray(image).convert('RGB')
img = img.resize(IMG_SIZE)
img_array = tf.keras.preprocessing.image.img_to_array(img)
img_array = np.expand_dims(img_array, axis=0)
img_array /= 255.0

# Болжам жасау
predictions = model.predict(img_array)[0]

# Нәтижені Gradio-ға арналған сөздікке айналдыру
results = {CLASS_NAMES[i]: float(predictions[i]) for i in range(len(CLASS_NAMES))}

return results

# Gradio интерфейсін құру
iface = gr.Interface(
fn=classify_trash_bin,
inputs=gr.Image(type="numpy", label="Қоқыс жәшігінің суреті"),
outputs=gr.Label(num_top_classes=2, label="Болжам"),
title="Қоқыс жәшігінің толуын анықтау (Full/Empty)",
description="Камера арқылы алынған суретті жүктеңіз немесе түсіріңіз. Hugging Face Spaces арқылы жұмыс істейді."
)

iface.launch()