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

# .h5 파일을 이용한 모델 로드
model = tf.keras.models.load_model('keras_model.h5')

# 클래스 레이블
class_labels = ['disposablecup', 'envmark', 'label', 'mug', 'nonlabel', 'reusablecup']

# 이미지 전처리 함수
def preprocess_image(image):
    image = Image.fromarray(image)
    image = image.resize((224, 224))  # 모델 입력 크기에 맞게 조정
    image = np.array(image) / 255.0   # 정규화
    image = np.expand_dims(image, axis=0)
    return image

# 예측 함수
def predict(image):
    processed_image = preprocess_image(image)
    predictions = model.predict(processed_image)
    predicted_class_index = np.argmax(predictions, axis=-1)[0]
    predicted_class_label = class_labels[predicted_class_index]
    return predicted_class_label

# Gradio 인터페이스 설정
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Textbox(label="Predicted Class"),
    live=True
)

# Gradio 웹 인터페이스 실행
if __name__ == "__main__":
    interface.launch()