Update app.py

app.py

@@ -1,17 +1,19 @@
 import tensorflow as tf
 import gradio as gr
-from transformers import ViTFeatureExtractor
 import numpy as np
 from PIL import Image
 
 # .h5 파일을 이용한 모델 로드
-model = tf.keras.models.load_model('keras_model.h5')
+model = tf.keras.models.load_model('path_to_your_model.h5')
+
+# 클래스 레이블 (예: 0은 'cat', 1은 'dog' 등)
+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 = image.resize((224, 224))  # 모델 입력 크기에 맞게 조정
+    image = np.array(image) / 255.0   # 정규화
     image = np.expand_dims(image, axis=0)
     return image
 
@@ -19,14 +21,15 @@ def preprocess_image(image):
 def predict(image):
     processed_image = preprocess_image(image)
     predictions = model.predict(processed_image)
-    predicted_class = np.argmax(predictions, axis=-1)[0]
-    return predicted_class
+    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.inputs.Image(source="webcam", tool="editor"),
-    outputs="label",
+    inputs=gr.Image(source="webcam", tool="editor"),
+    outputs=gr.Textbox(label="Predicted Class"),
     live=True
 )