Update app.py

app.py

@@ -1,26 +1,80 @@
-from transformers import AutoModelForSemanticSegmentation, AutoTokenizer
+import gradio as gr
 from PIL import Image
-import requests
-from io import BytesIO
+import numpy as np
+import tensorflow as tf
+from transformers import AutoFeatureExtractor, TFAutoModelForSemanticSegmentation
 
-# Hugging Face 모델 이름
+# Hugging Face 모델 및 토크나이저
 model_name = "nvidia/segformer-b0-finetuned-cityscapes-1024-1024"
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+model = TFAutoModelForSemanticSegmentation.from_pretrained(model_name)
 
-# 모델 및 토크나이저 로드
-model = AutoModelForSemanticSegmentation.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+def label_to_color_image(label, colormap):
+    color_seg = np.zeros(
+        (label.shape[0], label.shape[1], 3), dtype=np.uint8
+    )  # height, width, 3
+    for i in range(len(colormap)):
+        color_seg[label.numpy() == i, :] = colormap[i]
+    return color_seg
 
-# 이미지 다운로드 및 전처리
-image_url = "https://example.com/your_image.jpg"  # 이미지 URL을 실제 이미지 URL로 바꿔주세요.
-image = Image.open(BytesIO(requests.get(image_url).content))
+def draw_plot(pred_img, seg, colormap, labels_list):
+    # your existing draw_plot function, unchanged
 
-# 이미지를 모델 입력 형식으로 변환
-inputs = tokenizer(image, return_tensors="pt")
+def huggingface_model(input_img):
+    input_img = Image.fromarray(input_img)
 
-# 모델 실행
-outputs = model(**inputs)
-predictions = outputs.logits
+    inputs = feature_extractor(images=input_img, return_tensors="tf")
+    outputs = model(**inputs)
+    logits = outputs.logits
 
-# 여기서 predictions는 segmentation 맵입니다.
-# 이를 원하는 형식으로 처리하고 시각화할 수 있습니다.
-# 예를 들어, PIL Image로 변환하여 저장하거나 출력할 수 있습니다.
+    logits = tf.transpose(logits, [0, 2, 3, 1])
+    logits = tf.image.resize(
+        logits, input_img.size[::-1]
+    )  # We reverse the shape of `image` because `image.size` returns width and height.
+    seg = tf.math.argmax(logits, axis=-1)[0]
+
+    # Define the colormap for the cityscapes dataset
+    colormap = [
+        [128, 64, 128],
+        [244, 35, 232],
+        [70, 70, 70],
+        [102, 102, 156],
+        [190, 153, 153],
+        [153, 153, 153],
+        [250, 170, 30],
+        [220, 220, 0],
+        [107, 142, 35],
+        [152, 251, 152],
+        [0, 130, 180],
+        [220, 20, 60],
+        [255, 0, 0],
+        [0, 0, 142],
+        [0, 0, 70],
+        [0, 60, 100],
+        [0, 80, 100],
+        [0, 0, 230],
+        [119, 11, 32],
+    ]
+
+    color_seg = label_to_color_image(seg, colormap)
+
+    # Show image + mask
+    pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
+    pred_img = pred_img.astype(np.uint8)
+
+    # Draw plot
+    fig = draw_plot(pred_img, seg, colormap, labels_list)
+    return fig
+
+# 여러분이 가진 labels.txt 파일의 내용을 labels_list에 할당하세요.
+labels_list = ["label1", "label2", ...]
+
+demo = gr.Interface(
+    fn=huggingface_model,
+    inputs=gr.Image(shape=(1024, 1024)),  # 입력 이미지 크기는 모델의 입력 크기에 맞게 조절해야 합니다.
+    outputs=["plot"],
+    examples=["person-1.jpg", "person-2.jpg", "person-3.jpg", "person-4.jpg", "person-5.jpg"],
+    allow_flagging='never'
+)
+
+demo.launch()