Update app.py

app.py

@@ -82,14 +82,10 @@ def sepia(input_img):
     logits = outputs.logits
 
     logits = tf.transpose(logits, [0, 2, 3, 1])
-    logits = tf.image.resize(
-        logits, input_img.size[::-1]
-    ) 
+    logits = tf.image.resize(logits, input_img.size[::-1]) 
     seg = tf.math.argmax(logits, axis=-1)[0]
 
-    color_seg = np.zeros(
-        (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
-    )
+    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
 
@@ -103,8 +99,8 @@ def sepia(input_img):
     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_name = labels_list[label]
+        class_prob = np.mean(logits.numpy()[0][mask])
         class_probabilities[class_name] = class_prob
 
     # 정확성이 가장 높은 물체 정보 얻기
@@ -115,6 +111,7 @@ def sepia(input_img):
     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', 'text'],