Update app.py

app.py

@@ -97,14 +97,28 @@ def sepia(input_img):
     pred_img = pred_img.astype(np.uint8)
 
     fig = draw_plot(pred_img, seg)
-    
-    return fig
+
+    # 각 물체에 대한 예측 클래스와 확률 얻기
+    unique_labels = np.unique(seg.numpy().astype("uint8"))
+    class_probabilities = {}
+    for label in unique_labels:
+        mask = (seg.numpy() == label)
+        class_name = LABEL_NAMES[label]
+        class_prob = np.mean(outputs.logits.numpy()[0][mask])
+        class_probabilities[class_name] = class_prob
+
+    # 정확성이 가장 높은 물체 정보 얻기
+    max_prob_class = max(class_probabilities, key=class_probabilities.get)
+    max_prob_value = class_probabilities[max_prob_class]
+
+    # 출력 및 반환
+    print(f"Predicted class with highest probability: {max_prob_class}, Probability: {max_prob_value:.4f}")
+    return fig, f"Predicted class with highest probability: {max_prob_class}, Probability: {max_prob_value:.4f}"
 
 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()