m

app.py

@@ -24,19 +24,22 @@ MODEL_OPTIONS = {
 class GlaucomaModel(object):
     def __init__(self, 
                  cls_model_path="pamixsun/swinv2_tiny_for_glaucoma_classification", 
-                 seg_model_path=None,  # Make this optional
+                 seg_model_path=None,
                  device=torch.device('cpu')):
         self.device = device
-        # Classification model for glaucoma (always the same)
+        self.seg_model_path = seg_model_path or MODEL_OPTIONS["Pamixsun"]
+        
+        # Classification model setup remains the same
         self.cls_extractor = AutoImageProcessor.from_pretrained(cls_model_path)
         self.cls_model = Swinv2ForImageClassification.from_pretrained(cls_model_path).to(device).eval()
         
-        # Segmentation model - use provided path or default
-        seg_path = seg_model_path or MODEL_OPTIONS["Pamixsun"]  # Default to Pamixsun if none provided
-        self.seg_extractor = AutoImageProcessor.from_pretrained(seg_path)
-        self.seg_model = SegformerForSemanticSegmentation.from_pretrained(seg_path).to(device).eval()
+        # Segmentation model setup with model type detection
+        self.seg_extractor = AutoImageProcessor.from_pretrained(self.seg_model_path)
+        self.seg_model = SegformerForSemanticSegmentation.from_pretrained(self.seg_model_path).to(device).eval()
+        
+        # Detect model type
+        self.is_ferferefer = "ferferefer" in self.seg_model_path.lower()
         
-        # Mapping for class labels
         self.cls_id2label = self.cls_model.config.id2label
 
     def glaucoma_pred(self, image):
@@ -54,19 +57,38 @@ class GlaucomaModel(object):
         with torch.no_grad():
             inputs.to(self.device)
             outputs = self.seg_model(**inputs)
+        
         logits = outputs.logits.cpu()
         upsampled_logits = nn.functional.interpolate(
             logits, size=image.shape[:2], mode="bilinear", align_corners=False
         )
-        seg_probs = F.softmax(upsampled_logits, dim=1)
-        pred_disc_cup = upsampled_logits.argmax(dim=1)[0]
         
-        # Calculate segmentation confidence based on probability distribution
-        # For each pixel classified as cup/disc, check how confident the model is
+        if self.is_ferferefer:
+            # ferferefer model specific processing
+            seg_probs = F.softmax(upsampled_logits, dim=1)
+            pred_disc_cup = upsampled_logits.argmax(dim=1)[0]
+            
+            # Map ferferefer model classes to match Pamixsun format
+            # Assuming ferferefer uses different class indices
+            class_mapping = {
+                0: 0,  # background
+                1: 1,  # disc
+                2: 2   # cup
+            }
+            
+            pred_disc_cup_mapped = torch.zeros_like(pred_disc_cup)
+            for old_class, new_class in class_mapping.items():
+                pred_disc_cup_mapped[pred_disc_cup == old_class] = new_class
+            pred_disc_cup = pred_disc_cup_mapped
+        else:
+            # Pamixsun model processing (original logic)
+            seg_probs = F.softmax(upsampled_logits, dim=1)
+            pred_disc_cup = upsampled_logits.argmax(dim=1)[0]
+        
+        # Calculate confidences
         cup_mask = pred_disc_cup == 2
         disc_mask = pred_disc_cup == 1
         
-        # Get confidence only for pixels predicted as cup/disc
         cup_confidence = seg_probs[0, 2, cup_mask].mean().item() * 100 if cup_mask.any() else 0
         disc_confidence = seg_probs[0, 1, disc_mask].mean().item() * 100 if disc_mask.any() else 0
         
@@ -225,107 +247,107 @@ def main():
     
     if uploaded_files:
         try:
-            # Enhanced model loading feedback
-            st.write("🤖 Initializing AI models...")
-            st.write(f"• Loading classification model: pamixsun/swinv2_tiny_for_glaucoma_classification")
-            st.write(f"• Loading segmentation model: {selected_model}")
+            # Model loading feedback with better visuals
+            st.info("🤖 Loading AI Models")
+            st.write("Classification: pamixsun/swinv2_tiny_for_glaucoma_classification")
+            st.write(f"Segmentation: {selected_model}")
             
-            # Initialize model with selected segmentation model
             model = GlaucomaModel(
                 device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"),
                 seg_model_path=MODEL_OPTIONS[selected_model]
             )
             
-            # Show model loading completion
-            st.write("✅ Models loaded successfully")
-            st.write(f"🖥️ Using: {'GPU' if torch.cuda.is_available() else 'CPU'} for processing")
+            st.success(f"✅ Models loaded successfully - Using {'GPU' if torch.cuda.is_available() else 'CPU'}")
             st.write("---")
             
             for file in uploaded_files:
                 try:
-                    # Process each image with enhanced feedback
-                    st.write(f"📸 Processing image: {file.name}")
+                    st.info(f"📸 Processing: {file.name}")
                     image = Image.open(file).convert('RGB')
                     image_np = np.array(image)
                     
-                    # Get predictions
                     disease_idx, disc_cup_image, vcdr, cls_conf, cup_conf, disc_conf, cropped_image = model.process(image_np)
                     
-                    # Enhanced results display
                     st.write("---")
-                    st.write(f"Results for {file.name}")
+                    st.success(f"Results for: {file.name}")
                     
-                    # Diagnosis section
-                    st.write("📊 **Diagnosis Results:**")
-                    st.write(f"• Finding: {model.cls_id2label[disease_idx]}")
-                    st.write(f"• AI Confidence: {cls_conf:.1f}% ({get_confidence_level(cls_conf)})")
+                    # Key findings with better visuals
+                    st.info("📊 Key Findings")
                     
-                    # Enhanced Segmentation confidence section with detailed explanations
-                    st.write("\n🔍 **Understanding Segmentation Confidence:**")
-                    st.write("""
-                    Segmentation confidence shows how certain the AI is about each pixel it classified:
-                    • For the Optic Cup (central depression):
-                      - Measures the AI's certainty that the red-colored pixels are truly part of the cup
-                      - Higher confidence means clearer cup boundaries and more reliable VCDR
-                      
-                    • For the Optic Disc (entire circular area):
-                      - Indicates how sure the AI is about the green-outlined disc boundary
-                      - Higher confidence suggests better disc margin visibility
-                    
-                    Confidence scores are calculated by averaging the model's certainty 
-                    for each pixel it identified as cup or disc. A score of 100% would mean 
-                    the model is absolutely certain about every pixel's classification.
-                    """)
-                    
-                    st.write("\n📊 **Current Segmentation Confidence Scores:**")
-                    st.write(f"• Optic Cup Detection: {cup_conf:.1f}% - {get_confidence_level(cup_conf)}")
-                    st.write(f"• Optic Disc Detection: {disc_conf:.1f}% - {get_confidence_level(disc_conf)}")
-                    
-                    # Add interpretation guidance
-                    if cup_conf >= 75 and disc_conf >= 75:
-                        st.write("✅ High confidence scores indicate reliable measurements")
-                    elif cup_conf < 60 or disc_conf < 60:
-                        st.write("""
-                        ⚠️ Lower confidence scores might be due to:
-                        • Image quality issues (blur, poor contrast)
-                        • Unusual anatomical variations
-                        • Pathological changes affecting visibility
-                        • Poor image centering or focus
-                        
-                        Consider retaking the image if possible.
-                        """)
+                    # Diagnosis with color-coded warning levels
+                    diagnosis = model.cls_id2label[disease_idx]
+                    if diagnosis == "Glaucoma":
+                        st.warning(f"Diagnosis: {diagnosis} ({cls_conf:.1f}% confidence)")
+                    else:
+                        st.success(f"Diagnosis: {diagnosis} ({cls_conf:.1f}% confidence)")
                     
-                    # Clinical metrics
-                    st.write("\n📏 **Clinical Measurements:**")
-                    st.write(f"• Cup-to-Disc Ratio (VCDR): {vcdr:.3f}")
+                    # VCDR with risk levels
                     if vcdr > 0.7:
-                        st.write("  ⚠️ High VCDR - Potential risk indicator")
+                        st.warning(f"VCDR: {vcdr:.3f} - ⚠️ High Risk")
                     elif vcdr > 0.5:
-                        st.write("  ℹ️ Borderline VCDR - Follow-up recommended")
+                        st.warning(f"VCDR: {vcdr:.3f} - ⚠️ Borderline")
                     else:
-                        st.write("  ✅ Normal VCDR range")
+                        st.success(f"VCDR: {vcdr:.3f} - ✅ Normal")
+                    
+                    # Segmentation confidence
+                    st.info("🔍 Segmentation Confidence")
+                    st.write("""
+                    • Optic Cup (red area): Central depression
+                    • Optic Disc (green outline): Entire nerve area
+                    """)
                     
-                    # Image display with enhanced captions
-                    st.write("\n🖼️ **Visual Analysis:**")
-                    st.image(disc_cup_image, caption="""
-                        Segmentation Overlay
-                        • Green outline: Optic Disc boundary
-                        • Red area: Optic Cup region
-                        • Transparency shows underlying retina
-                        """)
-                    st.image(cropped_image, caption="Zoomed Region of Interest")
+                    # Cup and Disc confidence with warnings
+                    if cup_conf < 60:
+                        st.warning(f"Cup Detection: {cup_conf:.1f}% - Low Confidence")
+                    else:
+                        st.write(f"Cup Detection: {cup_conf:.1f}%")
+                        
+                    if disc_conf < 60:
+                        st.warning(f"Disc Detection: {disc_conf:.1f}% - Low Confidence")
+                    else:
+                        st.write(f"Disc Detection: {disc_conf:.1f}%")
                     
-                    # Add quality note if needed
-                    if cup_conf < 60 or disc_conf < 60:
-                        st.write("\n⚠️ Note: Low segmentation confidence. Image quality might affect measurements.")
+                    # Images with clear sections
+                    st.info("🖼️ Analysis Images")
+                    st.image(disc_cup_image, caption="Green: Optic Disc | Red: Optic Cup")
+                    st.image(cropped_image, caption="Region of Interest")
                     
                 except Exception as e:
                     st.error(f"Error processing {file.name}: {str(e)}")
                     continue
             
-            # Simple summary at the end
-            st.write("---")
-            st.write("Processing complete!")
+            # Download section
+            try:
+                st.info("📥 Preparing Download")
+                results = []
+                original_images = []
+                for file in uploaded_files:
+                    image = Image.open(file).convert('RGB')
+                    image_np = np.array(image)
+                    disease_idx, disc_cup_image, vcdr, cls_conf, cup_conf, disc_conf, cropped_image = model.process(image_np)
+                    results.append({
+                        'file_name': file.name,
+                        'diagnosis': model.cls_id2label[disease_idx],
+                        'confidence': cls_conf,
+                        'vcdr': vcdr,
+                        'cup_conf': cup_conf,
+                        'disc_conf': disc_conf,
+                        'processed_image': disc_cup_image,
+                        'cropped_image': cropped_image
+                    })
+                    original_images.append(image_np)
+                
+                zip_data = save_results(results, original_images)
+                b64_zip = base64.b64encode(zip_data).decode()
+                
+                st.success("✅ Download Ready")
+                href = f'<a href="data:application/zip;base64,{b64_zip}" download="glaucoma_results.zip">📥 Download All Results (ZIP)</a>'
+                st.markdown(href, unsafe_allow_html=True)
+                
+            except Exception as e:
+                st.error(f"Error preparing download: {str(e)}")
+            
+            st.success("✅ All Processing Complete!")
             
         except Exception as e:
             st.error(f"An error occurred: {str(e)}")