Update app.py

app.py

@@ -10,13 +10,21 @@ from typing import Tuple
 # Device configuration
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Constants
+# ====================== CONSTANTS ======================
 CLASS_NAMES = ["Mild", "Moderate", "No_DR", "Proliferate_DR", "Severe"]
 LESION_COLORS = {
     0: [0, 0, 0],      # Background (black)
     1: [255, 255, 0],  # Bright lesions (yellow)
     2: [255, 0, 0]     # Red lesions (red)
 }
+UK_GRADES = {
+    "No_DR": "R0 – No retinopathy",
+    "Mild": "R1 – Background DR",
+    "Moderate": "R1 – Background DR",
+    "Severe": "R2 – Pre-proliferative DR",
+    "Proliferate_DR": "R3 – Proliferative DR"
+}
+
 
 # ====================== UNET ARCHITECTURE ======================
 class UNet(nn.Module):
@@ -199,11 +207,21 @@ def main():
             probabilities = ps[0].cpu().numpy() * 100
 
         st.subheader("Classification Results")
-        if CLASS_NAMES[pred_class] == "No_DR":  # Changed to compare class name
-            st.success(f"**Prediction:** {CLASS_NAMES[pred_class]}")
+        predicted_class_name = CLASS_NAMES[pred_class]
+        uk_grade = UK_GRADES[predicted_class_name]
+        
+        if predicted_class_name == "No_DR":
+            st.success(f"""
+            **Prediction:** {predicted_class_name}  
+            **UK Grade:** {uk_grade}
+            """)
             st.write("No diabetic retinopathy detected - no segmentation needed.")
         else:
-            st.error(f"**Prediction:** {CLASS_NAMES[pred_class]}")
+            st.error(f"""
+            **Prediction:** {predicted_class_name}  
+            **UK Grade:** {uk_grade}
+            """)
+            
             st.write("**Confidence Levels:**")
             for name, prob in zip(CLASS_NAMES, probabilities):
                 st.progress(int(prob))