EUNSEO56 commited on
Commit
66da7ef
1 Parent(s): 132aec0

Create aapp.py

Browse files
Files changed (1) hide show
  1. aapp.py +39 -0
aapp.py ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import numpy as np
3
+ from PIL import Image
4
+ import tensorflow as tf
5
+ from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
6
+
7
+ # Segformer 모델과 관련 객체를 초기화
8
+ feature_extractor = SegformerFeatureExtractor.from_pretrained("nickmuchi/segformer-b4-finetuned-segments-sidewalk")
9
+ model = TFSegformerForSemanticSegmentation.from_pretrained("nickmuchi/segformer-b4-finetuned-segments-sidewalk", from_pt=True)
10
+
11
+ def perform_semantic_segmentation(input_img):
12
+ input_img = Image.fromarray(input_img)
13
+
14
+ # 이미지를 처리하고 모델에 전달
15
+ inputs = feature_extractor(images=input_img, return_tensors="tf")
16
+ outputs = model(**inputs)
17
+ logits = outputs.logits
18
+
19
+ # 모델 출력을 처리하여 시맨틱 분할 결과를 얻음
20
+ logits = tf.transpose(logits, [0, 2, 3, 1])
21
+ logits = tf.image.resize(logits, input_img.size[::-1])
22
+ seg = tf.math.argmax(logits, axis=-1)[0]
23
+
24
+ return input_img, seg.numpy()
25
+
26
+ def segformer_interface(input_image):
27
+ original_image, segmentation_map = perform_semantic_segmentation(input_image)
28
+ return original_image, segmentation_map
29
+
30
+ # Gradio 데모 구성
31
+ demo = gr.Interface(
32
+ fn=segformer_interface,
33
+ inputs=gr.Image(shape=(400, 600)),
34
+ outputs=[gr.Image(type="plot"), gr.Image(type="plot")], # 원본 이미지 및 시맨틱 분할 맵을 출력
35
+ examples=["side-1.jpg", "side-2.jpg", "side-3.jpg"],
36
+ allow_flagging='never'
37
+ )
38
+
39
+ demo.launch()