Update app.py

app.py

@@ -1,28 +1,24 @@
 import gradio as gr
 import numpy as np
 from PIL import Image
-import tensorflow as tf
 from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
 
+# Segformer 모델 및 feature extractor 불러오기
 feature_extractor = SegformerFeatureExtractor.from_pretrained(
     "nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
-
 model = TFSegformerForSemanticSegmentation.from_pretrained(
     "nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
 
-
 # 모델 예측 함수를 정의합니다.
 def classify_image(img):
-    # 이미지를 모델 입력 크기에 맞게 조정합니다.
-    img = img.resize(input_size)
-    img_array = np.array(img) / 255.0  # 이미지를 0에서 1 사이로 정규화합니다.
-    img_array = np.expand_dims(img_array, axis=0)  # 배치 차원을 추가합니다.
-
+    # 이미지를 전처리합니다.
+    inputs = feature_extractor(images=img, return_tensors="tf")
+    
     # 모델로 예측을 수행합니다.
-    predictions = model.predict(img_array)
+    predictions = model(**inputs)
     
     # 예측 결과 중에서 가장 높은 확률을 가진 클래스를 선택합니다.
-    predicted_label = np.argmax(predictions)
+    predicted_label = tf.argmax(predictions.logits[0], axis=-1).numpy()
 
     # 라벨을 반환합니다.
     return predicted_label