Update app.py

app.py

@@ -138,9 +138,7 @@ def get_segmenter_transform():
 def process_segmentation_output(output: torch.Tensor) -> Tuple[np.ndarray, np.ndarray]:
     probs = torch.softmax(output, dim=1).cpu().numpy().squeeze()
     pred_class = np.argmax(probs, axis=0)
-    final_mask = np.zeros_like(pred_class, dtype=np.uint8)
-    final_mask[(pred_class == 1) | (pred_class == 4)] = 1
-    final_mask[(pred_class == 2) | (pred_class == 3)] = 2
+    final_mask = pred_class.astype(np.uint8)  # Already 0=bg, 1=bright, 2=red 
     return final_mask, probs
 
 # ====================== VISUALIZATION ======================
@@ -219,10 +217,8 @@ def main():
             with st.spinner("Detecting lesions..."):
                 seg_results = segment_image(original_image, segmenter)
                 overlay = create_lesion_overlay(original_image, seg_results['mask'])
-                heat_bright = create_heatmap(seg_results['probs'][1] + seg_results['probs'][4], 
-                                           original_image.size)
-                heat_red = create_heatmap(seg_results['probs'][2] + seg_results['probs'][3], 
-                                        original_image.size)
+                heat_bright = create_heatmap(seg_results['probs'][1], original_image.size) 
+                heat_red = create_heatmap(seg_results['probs'][2], original_image.size)    # Only class 2
 
                 with col2:
                     st.image(overlay, caption="Lesion Overlay", use_container_width=True)