Update app.py

app.py

@@ -11,7 +11,7 @@ from typing import Tuple
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Constants
-CLASS_NAMES = ["No_DR", "Mild", "Moderate", "Severe", "Proliferate_DR"]
+CLASS_NAMES = ["Mild", "Moderate", "No_DR", "Proliferate_DR", "Severe"]
 LESION_COLORS = {
     0: [0, 0, 0],      # Background (black)
     1: [255, 255, 0],  # Bright lesions (yellow)
@@ -152,10 +152,6 @@ def create_lesion_overlay(original: Image.Image, mask: np.ndarray) -> Image.Imag
         overlay[mask_resized == class_idx] = color
     return Image.fromarray(cv2.addWeighted(overlay, 0.4, original_np, 0.6, 0))
 
-def create_heatmap(prob_map: np.ndarray, original_size: Tuple[int, int]) -> np.ndarray:
-    resized = cv2.resize(prob_map, original_size)
-    return cv2.applyColorMap((resized * 255).astype(np.uint8), cv2.COLORMAP_JET)
-
 def segment_image(image: Image.Image, model: nn.Module) -> dict:
     processed_img = preprocess_segmenter(image)
     img_pil = Image.fromarray(processed_img)
@@ -217,13 +213,9 @@ 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], 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)
-                    st.image(heat_bright, caption="Bright Lesion Probability", use_container_width=True)
-                    st.image(heat_red, caption="Red Lesion Probability", use_container_width=True)
 
                 # Metrics
                 st.write("**Lesion Analysis:**")