Update app.py

app.py

@@ -335,9 +335,9 @@ if st.button("Run Inference"):
             # Ground truth overlay
             gt_overlay = optic_images[i].copy()
             if is_multiclass:
-                gt_overlay[real_masks[i][:, :, 1] == 1] = red_color
+                gt_overlay[masks[i][:, :, 1] == 1] = red_color
             else:
-                gt_overlay[real_masks[i].squeeze() == 1] = red_color
+                gt_overlay[masks[i].squeeze() == 1] = red_color
         
             ax[2].imshow(optic_images[i])
             ax[2].imshow(gt_overlay, alpha=0.4)