Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
+import tensorflow as tf
 
 # Segformer 모델 및 feature extractor 불러오기
 feature_extractor = SegformerFeatureExtractor.from_pretrained(
@@ -11,17 +12,21 @@ model = TFSegformerForSemanticSegmentation.from_pretrained(
 
 # 모델 예측 함수를 정의합니다.
 def classify_image(img):
-    # 이미지를 전처리합니다.
-    inputs = feature_extractor(images=img, return_tensors="tf")
-    
-    # 모델로 예측을 수행합니다.
-    predictions = model(**inputs)
-    
-    # 예측 결과 중에서 가장 높은 확률을 가진 클래스를 선택합니다.
-    predicted_label = tf.argmax(predictions.logits[0], axis=-1).numpy()
+    try:
+        # 이미지를 전처리합니다.
+        inputs = feature_extractor(images=img, return_tensors="tf")
 
-    # 라벨을 반환합니다.
-    return predicted_label
+        # 모델로 예측을 수행합니다.
+        predictions = model(**inputs)
+
+        # 예측 결과 중에서 가장 높은 확률을 가진 클래스를 선택합니다.
+        predicted_label = tf.argmax(predictions.logits[0], axis=-1).numpy()
+
+        # 라벨을 반환합니다.
+        return predicted_label
+    except Exception as e:
+        # 예외가 발생하면 예외 메시지를 반환합니다.
+        return str(e)
 
 # Gradio UI를 생성합니다.
 iface = gr.Interface(fn=classify_image,