finish

app.py

@@ -5,16 +5,15 @@ import matplotlib.pyplot as plt
 import numpy as np
 from PIL import Image
 import tensorflow as tf
-from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation, AutoFeatureExtractor, AutoModelForImageCaptioning
+from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
 
 feature_extractor = SegformerFeatureExtractor.from_pretrained(
     "nvidia/segformer-b2-finetuned-cityscapes-1024-1024"
 )
-seg_model = TFSegformerForSemanticSegmentation.from_pretrained(
+model = TFSegformerForSemanticSegmentation.from_pretrained(
     "nvidia/segformer-b2-finetuned-cityscapes-1024-1024"
 )
 
-caption_model = AutoModelForImageCaptioning.from_pretrained("facebook/deit-base-cc-turbo")
 
 def ade_palette():
     """ADE20K palette that maps each class to RGB values."""
@@ -83,7 +82,7 @@ def sepia(input_img):
     input_img = Image.fromarray(input_img)
 
     inputs = feature_extractor(images=input_img, return_tensors="tf")
-    outputs = seg_model(**inputs)
+    outputs = model(**inputs)
     logits = outputs.logits
 
     logits = tf.transpose(logits, [0, 2, 3, 1])
@@ -106,12 +105,12 @@ def sepia(input_img):
     return fig
 
 
-def segment_and_caption(input_img):
+def segment_image(input_img):
     input_img = Image.fromarray(input_img)
 
     # 세그멘테이션 수행
     inputs = feature_extractor(images=input_img, return_tensors="tf")
-    outputs = seg_model(**inputs)
+    outputs = model(**inputs)
     logits = outputs.logits
 
     logits = tf.transpose(logits, [0, 2, 3, 1])
@@ -120,23 +119,12 @@ def segment_and_caption(input_img):
     )
     seg = tf.math.argmax(logits, axis=-1)[0]
 
-    # 세그멘테이션 결과를 텍스트로 변환
-    seg_text = ""
-    for label, label_name in enumerate(labels_list):
-        count = np.sum(seg.numpy() == label)
-        seg_text += f"{label_name}: {count} pixels\n"
+    return input_img, seg
 
-    # 이미지 캡션 생성
-    caption_input = caption_model.generate(input_img, max_length=20, num_return_sequences=1)
-    caption_text = caption_input[0]['text']
 
-    # 세그멘테이션 결과와 캡션을 반환
-    return input_img, seg_text, caption_text
-
-
-demo = gr.Interface(fn=segment_and_caption,
+demo = gr.Interface(fn=segment_image,
                     inputs=gr.Image(shape=(1024, 1024)),
-                    outputs=["image", "text", "text"],
+                    outputs=["image", "image"],
                     examples=["city-1.jpg", "city-2.jpg", "city-3.jpg", "city-4.jpg", "city-5.jpg"],
                     allow_flagging='never')