Update app.py

app.py

@@ -115,7 +115,17 @@ def sepia(input_img):
         print(f"{label} probabilities:")
         print(prob_map)
     
-    return fig, str(class_probabilities)
+     class_probabilities = {
+        labels_list[i]: np.max(probabilities[:, :, i].numpy()) for i in range(len(labels_list))
+    }
+
+    # 가장 높은 확률을 가진 클래스를 출력
+    max_prob_class = max(class_probabilities, key=class_probabilities.get)
+    max_prob_value = class_probabilities[max_prob_class]
+    
+    print(f"Predicted class: {max_prob_class}, Probability: {max_prob_value:.4f}")
+
+    return fig, f"Predicted class: {max_prob_class}, Probability: {max_prob_value:.4f}"
 
 demo = gr.Interface(fn=sepia,
                     inputs=gr.Image(shape=(400, 600)),