Update app.py

app.py

@@ -87,6 +87,8 @@ def sepia(input_img):
     ) 
     seg = tf.math.argmax(logits, axis=-1)[0]
 
+    probabilities = tf.nn.softmax(logits, axis=-1)[0]  # 확률값 추출
+
     color_seg = np.zeros(
         (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
     )
@@ -97,14 +99,19 @@ def sepia(input_img):
     pred_img = pred_img.astype(np.uint8)
 
     fig = draw_plot(pred_img, seg)
-    
+
+    # 각 클래스의 확률을 출력
+    class_probabilities = {
+        labels_list[i]: probabilities[:, :, i].numpy() for i in range(len(labels_list))
+    }
+    print(class_probabilities)
+
     return fig
 
 demo = gr.Interface(fn=sepia,
                     inputs=gr.Image(shape=(400, 600)),
-                    outputs=['plot'],
+                    outputs=["plot", "text"],
                     examples=["citiscapes-1.jpeg", "citiscapes-2.jpeg"],
                     allow_flagging='never')
 
-
 demo.launch()