finish

app.py

@@ -6,14 +6,16 @@ import numpy as np
 from PIL import Image
 import tensorflow as tf
 from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
+from transformers import BigBirdForImageCaptioning
 
 feature_extractor = SegformerFeatureExtractor.from_pretrained(
     "nvidia/segformer-b2-finetuned-cityscapes-1024-1024"
 )
-model = TFSegformerForSemanticSegmentation.from_pretrained(
+seg_model = TFSegformerForSemanticSegmentation.from_pretrained(
     "nvidia/segformer-b2-finetuned-cityscapes-1024-1024"
 )
 
+caption_model = BigBirdForImageCaptioning.from_pretrained("bigbird/image-captioning-base")
 
 def ade_palette():
     """ADE20K palette that maps each class to RGB values."""
@@ -82,7 +84,7 @@ def sepia(input_img):
     input_img = Image.fromarray(input_img)
 
     inputs = feature_extractor(images=input_img, return_tensors="tf")
-    outputs = model(**inputs)
+    outputs = seg_model(**inputs)
     logits = outputs.logits
 
     logits = tf.transpose(logits, [0, 2, 3, 1])
@@ -105,11 +107,12 @@ def sepia(input_img):
     return fig
 
 
-def with_labels(input_img):
+def segment_and_caption(input_img):
     input_img = Image.fromarray(input_img)
 
+    # 세그멘테이션 수행
     inputs = feature_extractor(images=input_img, return_tensors="tf")
-    outputs = model(**inputs)
+    outputs = seg_model(**inputs)
     logits = outputs.logits
 
     logits = tf.transpose(logits, [0, 2, 3, 1])
@@ -118,21 +121,25 @@ def with_labels(input_img):
     )
     seg = tf.math.argmax(logits, axis=-1)[0]
 
-    color_seg = np.zeros(
-        (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
-    )
-    for label, color in enumerate(colormap):
-        color_seg[seg.numpy() == label, :] = color
+    # 세그멘테이션 결과를 텍스트로 변환
+    seg_text = ""
+    for label, label_name in enumerate(labels_list):
+        count = np.sum(seg.numpy() == label)
+        seg_text += f"{label_name}: {count} pixels\n"
+
+    # 이미지 캡션 생성
+    caption = caption_model.generate(input_img, max_length=20, num_return_sequences=1, return_dict_in_generate=True)
+    caption_text = caption[0]['text']
+
+    # 세그멘테이션 결과와 캡션을 반환
+    return input_img, seg_text, caption_text
 
-    pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
-    pred_img = pred_img.astype(np.uint8)
 
-    return input_img, pred_img, labels_list
 
 
-demo = gr.Interface(fn=with_labels,
+demo = gr.Interface(fn=segment_and_caption,
                     inputs=gr.Image(shape=(1024, 1024)),
-                    outputs=["image", "image", "text"],
+                    outputs=["image","text", "text"],
                     examples=["city-1.jpg", "city-2.jpg", "city-3.jpg", "city-4.jpg", "city-5.jpg"],
                     allow_flagging='never')