Update app.py

app.py

@@ -1,290 +1,331 @@
-import gradio as gr
-import torch
-import torchvision.models as models
-import torchvision.transforms as transforms
-from PIL import Image
-
-# Define class labels for Diabetic Retinopathy
-class_labels = ['No DR', 'Mild', 'Moderate', 'Severe', 'Proliferative DR']
-
-# Define image preprocessing function
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-])
-
-def load_model():
-    model = models.vit_b_16(weights=None)
-    if isinstance(model.heads, torch.nn.Sequential):
-        in_features = model.heads[0].in_features
-    else:
-        in_features = model.heads.in_features
-    model.heads = torch.nn.Linear(in_features, len(class_labels))
-    state_dict = torch.load("best_model.pt", map_location=torch.device('cpu'))
-    model.load_state_dict(state_dict)
-    model.eval()
-    return model
-
-model = load_model()
-
-def predict_retinopathy(image):
-    if image is None:
-        return "Please upload an image for analysis."
-    
-    image = transform(image).unsqueeze(0)
-    with torch.no_grad():
-        output = model(image)
-    probabilities = torch.nn.functional.softmax(output[0], dim=0)
-    prediction = torch.argmax(probabilities).item()
-    
-    result_text = "<div style='text-align: center; padding: 20px; background-color: #46a7f7; border-radius: 10px; margin: 20px 0;'>"
-    result_text += "<h2>👁️ Analysis Results</h2>"
-    result_text += f"<h3>Primary Classification: {class_labels[prediction]}</h3>"
-    result_text += "</div>"
-    
-    result_text += "<table style='width: 100%; border-collapse: collapse; margin: 20px 0;'>"
-    result_text += "<tr><th style='padding: 10px; text-align: left; background-color: #46a7f7;'>Classification</th>"
-    result_text += "<th style='padding: 10px; text-align: right; background-color: #46a7f7;'>Confidence</th></tr>"
-    
-    for i, label in enumerate(class_labels):
-        prob = probabilities[i] * 100
-        color = f"hsl({120 - prob}, 70%, 50%)"
-        result_text += f"<tr style='border-bottom: 1px solid #dee2e6;'>"
-        result_text += f"<td style='padding: 10px;'>{label}</td>"
-        result_text += f"<td style='padding: 10px; text-align: right;'>"
-        result_text += f"<span style='color: {color}; font-weight: bold;'>{prob:.2f}%</span></td></tr>"
-    
-    result_text += "</table>"
-    return result_text
-
-# Custom CSS for enhanced styling
-custom_css = """
-footer {display: none !important;}
-.container {max-width: 1000px; margin: auto; padding: 20px;}
-.header {text-align: center; margin-bottom: 40px;}
-.main-title {color: #2c3e50; font-size: 2.5em; margin-bottom: 20px;}
-.subtitle {color: #34495e; font-size: 1.2em; line-height: 1.6;}
-.input-section {background: #f8f9fa; padding: 30px; border-radius: 15px; margin: 20px 0;}
-.disclaimer {background: #fff3cd; color: #856404; padding: 15px; border-radius: 8px; margin-top: 30px;}
-.info-box {background: #e3f2fd; padding: 20px; border-radius: 10px; margin: 15px 0;}
-.warning-box {background: #ffebee; padding: 20px; border-radius: 10px; margin: 15px 0;}
-"""
-
-# Create Gradio UI with enhanced styling
-with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
-    # Header Section
-    gr.HTML("""
-        <div class="header">
-            <h1 class="main-title">👁️ Welcome to EarlyMed Diabetic Retinopathy Scanner</h1>
-        </div>
-    """)
-    
-    # Introduction
-    gr.Markdown("""
-        Welcome to EarlyMed—an initiative by our team at VIT-AP University dedicated to empowering you with early health insights. 
-        Leveraging AI for early detection, our mission is simple: "Early Detection, Smarter Decision." 
-        Our Diabetic Retinopathy detection project is one of our key efforts to help you stay informed before visiting a doctor.
-    """)
-    
-    # What is Diabetic Retinopathy Section
-    gr.Markdown("""
-        ## 🔍 Understanding Diabetic Retinopathy (DR)
-        
-        Diabetic Retinopathy is a diabetes complication that affects the eyes. It's caused by damage to the blood vessels 
-        in the tissue at the back of the eye (retina). Poorly controlled blood sugar is a risk factor. 
-        Early detection can prevent vision loss.
-        
-        ### Primary Symptoms of DR:
-        
-        * **Blurred or fluctuating vision**
-        * **Impaired color vision**
-        * **Dark or empty areas in your vision**
-        * **Floaters (spots or dark strings floating in your vision)**
-        * **Vision loss**
-        * **Difficulty seeing at night**
-        
-        ### Risk Factors:
-        
-        * Duration of diabetes
-        * Poor control of blood sugar level
-        * High blood pressure
-        * High cholesterol
-        * Pregnancy
-        * Tobacco use
-        * Being of Hispanic, Native American, or African heritage
-    """)
-    
-    # Image Upload Section
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("### 📤 Upload Retinal Image")
-            image_input = gr.Image(
-                type="pil",
-                label="Upload a Fundus (Retinal) Image",
-                elem_classes="input-section"
-            )
-            submit_btn = gr.Button(
-                "🔍 Analyze Image",
-                variant="primary",
-                size="lg"
-            )
-    
-    # Results Section
-    diagnosis_output = gr.HTML()
-    
-    submit_btn.click(
-        fn=predict_retinopathy,
-        inputs=image_input,
-        outputs=diagnosis_output
-    )
-    
-    # Understanding Stages
-    gr.Markdown("""
-        ## 📊 Understanding DR Stages and Required Actions
-        
-        ### No DR Stage
-        **What it means:**
-        * No visible signs of diabetic retinopathy
-        * Normal retinal appearance
-        * Blood vessels appear healthy
-        
-        **Required Actions:**
-        * Continue regular eye check-ups (yearly for diabetic patients)
-        * Maintain blood sugar control
-        * Follow healthy lifestyle recommendations
-        
-        ### Mild Stage
-        **What it means:**
-        * Small areas of balloon-like swelling in the retina's tiny blood vessels
-        * Microaneurysms detected
-        * Early changes that don't typically affect vision
-        
-        **Required Actions:**
-        * Schedule follow-up with an ophthalmologist within 6-12 months
-        * Improve blood sugar control
-        * Monitor blood pressure and cholesterol
-        * Consider more frequent eye screenings
-        
-        ### Moderate Stage
-        **What it means:**
-        * More extensive damage to blood vessels
-        * Blood vessels may begin to leak
-        * Some retinal swelling might be present
-        
-        **Required Actions:**
-        * Consult with an ophthalmologist within 3-6 months
-        * Possible referral to a retina specialist
-        * Strict blood sugar, blood pressure, and cholesterol management
-        * Discuss potential early interventions
-        
-        ### Severe Stage
-        **What it means:**
-        * Significant blockage of blood vessels
-        * Retina being deprived of proper blood supply
-        * Increased risk for vision-threatening complications
-        
-        **Required Actions:**
-        * Urgent consultation with a retina specialist
-        * Possible laser treatment consideration
-        * Intensive monitoring of diabetes management
-        * Prepare for possible treatments
-        
-        ### Proliferative DR Stage
-        **What it means:**
-        * Advanced disease with new abnormal blood vessels growing
-        * High risk of vitreous hemorrhage
-        * Potential retinal detachment
-        * Significant threat to vision
-        
-        **Required Actions:**
-        * Immediate specialist intervention
-        * Likely need for laser treatment or surgery
-        * Aggressive management of all contributing factors
-        * Regular monitoring post-treatment
-    """)
-    
-    # Technical Details
-    gr.Markdown("""
-        ## 🔬 How Our AI Works
-        
-        * **Image Processing:** Fundus (retinal) images are resized and normalized to match the model's input specifications.
-        * **Feature Extraction:** Our advanced Vision Transformer scans the image to detect subtle changes in blood vessels, microaneurysms, exudates, and other retinal features.
-        * **Classification:** Based on the extracted features, the AI classifies the image into one of five categories:
-           * **No DR**
-           * **Mild**
-           * **Moderate**
-           * **Severe**
-           * **Proliferative DR**
-           
-        ## 🎯 Why It Is Reliable
-        
-        * **Robust Training:** Trained on thousands of diverse retinal images, our model has learned to recognize even the faintest indicators of DR.
-        * **Rigorous Validation:** Extensive testing and validation protocols ensure the model's predictions are accurate and trustworthy.
-        * **Continuous Improvement:** Regular updates and refinements keep the model up-to-date with the latest ophthalmological insights and technological advancements.
-        * **Cutting-Edge Technology:** Utilizing state-of-the-art Vision Transformer architecture, our AI leverages modern deep learning techniques to provide early and reliable health insights.
-    """)
-    
-    # Treatment and Support Information
-    gr.Markdown("""
-        ## 💉 Treatment Approaches
-        
-        ### Standard Treatment Options:
-        * **Laser Treatment (Photocoagulation):** Seals or destroys leaking blood vessels
-        * **Anti-VEGF Injections:** Reduces growth of abnormal blood vessels
-        * **Vitrectomy:** Surgical removal of vitreous gel and blood
-        * **Corticosteroid Injections:** Reduces inflammation and swelling
-        
-        ### Supportive Care:
-        * Blood sugar control
-        * Blood pressure management
-        * Regular eye examinations
-        * Vision rehabilitation if needed
-        * Psychological support
-        
-        ### Long-term Considerations:
-        * Lifestyle modifications
-        * Regular monitoring of diabetes
-        * Adherence to medication regimens
-        * Regular ophthalmologic follow-ups
-    """)
-    
-    # Prevention and Early Detection
-    gr.Markdown("""
-        ## 🛡️ Prevention and Early Detection
-        
-        ### Risk Reduction Strategies:
-        * Maintaining optimal blood sugar levels
-        * Controlling blood pressure
-        * Managing cholesterol levels
-        * Regular exercise
-        * Healthy diet low in refined carbohydrates
-        * Avoiding tobacco use
-        
-        ### When to Seek Medical Attention:
-        * Sudden vision changes
-        * Eye pain or redness
-        * Floaters or spots in vision
-        * Blurred vision
-        * Any noticeable change in visual acuity
-    """)
-   # Disclaimer
-    gr.Markdown("""
-        ---
-        ### ⚠️ Important Medical Disclaimer
-        
-        This AI-powered tool is designed to assist in the early detection of Diabetic Retinopathy (DR) through 
-        retinal image analysis. We strongly urge users to consult an ophthalmologist for appropriate medical 
-        guidance after getting the diagnosis. 
-        
-        **Please Note:**
-        * Results should be verified by healthcare professionals
-        * Consult a qualified ophthalmologist for proper treatment
-        
-        This initiative is developed by our team at VIT-AP University with the goal of empowering individuals to be more 
-        aware of their eye health before visiting a doctor. Our mission is to leverage AI for early detection and better 
-        healthcare awareness.
-    """)
-
-if __name__ == "__main__":
+import gradio as gr
+import torch
+import torchvision.models as models
+import torchvision.transforms as transforms
+from PIL import Image
+
+# Define class labels for Diabetic Retinopathy
+class_labels = ['No DR', 'Mild', 'Moderate', 'Severe', 'Proliferative DR']
+
+# Define image preprocessing function with augmentation
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+# Define augmented transform for training (not used in inference but shown for completeness)
+train_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.RandomHorizontalFlip(),
+    transforms.RandomRotation(20),
+    transforms.ColorJitter(brightness=0.1, contrast=0.1),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+def load_model():
+    # Load a pretrained ResNet-50 model
+    model = models.resnet50(pretrained=False)
+    # Modify the final layer for 5 class classification
+    in_features = model.fc.in_features
+    model.fc = torch.nn.Linear(in_features, len(class_labels))
+    # Load the saved model weights
+    state_dict = torch.load("best_model.pt", map_location=torch.device('cpu'))
+    model.load_state_dict(state_dict)
+    model.eval()
+    return model
+
+model = load_model()
+
+def predict_retinopathy(image):
+    if image is None:
+        return "Please upload an image for analysis."
+    
+    # Preprocess the image
+    image = transform(image).unsqueeze(0)
+    
+    # Make prediction
+    with torch.no_grad():
+        output = model(image)
+    probabilities = torch.nn.functional.softmax(output[0], dim=0)
+    prediction = torch.argmax(probabilities).item()
+    
+    # Create result HTML
+    result_text = "<div style='text-align: center; padding: 20px; background-color: #46a7f7; border-radius: 10px; margin: 20px 0;'>"
+    result_text += "<h2>👁️ Analysis Results</h2>"
+    result_text += f"<h3>Primary Classification: {class_labels[prediction]} (Stage {prediction})</h3>"
+    result_text += "</div>"
+    
+    result_text += "<table style='width: 100%; border-collapse: collapse; margin: 20px 0;'>"
+    result_text += "<tr><th style='padding: 10px; text-align: left; background-color: #46a7f7;'>Classification</th>"
+    result_text += "<th style='padding: 10px; text-align: right; background-color: #46a7f7;'>Confidence</th></tr>"
+    
+    for i, label in enumerate(class_labels):
+        prob = probabilities[i] * 100
+        color = f"hsl({120 - prob}, 70%, 50%)"
+        result_text += f"<tr style='border-bottom: 1px solid #dee2e6;'>"
+        result_text += f"<td style='padding: 10px;'>{label} (Stage {i})</td>"
+        result_text += f"<td style='padding: 10px; text-align: right;'>"
+        result_text += f"<span style='color: {color}; font-weight: bold;'>{prob:.2f}%</span></td></tr>"
+    
+    result_text += "</table>"
+    return result_text
+
+# Custom CSS for enhanced styling
+custom_css = """
+footer {display: none !important;}
+.container {max-width: 1000px; margin: auto; padding: 20px;}
+.header {text-align: center; margin-bottom: 40px;}
+.main-title {color: #2c3e50; font-size: 2.5em; margin-bottom: 20px;}
+.subtitle {color: #34495e; font-size: 1.2em; line-height: 1.6;}
+.input-section {background: #f8f9fa; padding: 30px; border-radius: 15px; margin: 20px 0;}
+.disclaimer {background: #fff3cd; color: #856404; padding: 15px; border-radius: 8px; margin-top: 30px;}
+.info-box {background: #e3f2fd; padding: 20px; border-radius: 10px; margin: 15px 0;}
+.warning-box {background: #ffebee; padding: 20px; border-radius: 10px; margin: 15px 0;}
+"""
+
+# Create Gradio UI with enhanced styling
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    # Header Section
+    gr.HTML("""
+        <div class="header">
+            <h1 class="main-title">👁️ Welcome to EarlyMed Diabetic Retinopathy Scanner</h1>
+        </div>
+    """)
+    
+    # Introduction
+    gr.Markdown("""
+        Welcome to EarlyMed—an initiative by our team at VIT-AP University dedicated to empowering you with early health insights. 
+        Leveraging AI for early detection, our mission is simple: "Early Detection, Smarter Decision." 
+        Our Diabetic Retinopathy detection project uses a powerful ResNet-50 deep learning model to help you stay informed before visiting a doctor.
+    """)
+    
+    # What is Diabetic Retinopathy Section
+    gr.Markdown("""
+        ## 🔍 Understanding Diabetic Retinopathy (DR)
+        
+        Diabetic Retinopathy is a diabetes complication that affects the eyes. It's caused by damage to the blood vessels 
+        in the tissue at the back of the eye (retina). Poorly controlled blood sugar is a risk factor. 
+        Early detection can prevent vision loss.
+        
+        ### Primary Symptoms of DR:
+        
+        * **Blurred or fluctuating vision**
+        * **Impaired color vision**
+        * **Dark or empty areas in your vision**
+        * **Floaters (spots or dark strings floating in your vision)**
+        * **Vision loss**
+        * **Difficulty seeing at night**
+        
+        ### Risk Factors:
+        
+        * Duration of diabetes
+        * Poor control of blood sugar level
+        * High blood pressure
+        * High cholesterol
+        * Pregnancy
+        * Tobacco use
+        * Being of Hispanic, Native American, or African heritage
+    """)
+    
+    # Image Upload Section
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### 📤 Upload Retinal Image")
+            image_input = gr.Image(
+                type="pil",
+                label="Upload a Fundus (Retinal) Image",
+                elem_classes="input-section"
+            )
+            submit_btn = gr.Button(
+                "🔍 Analyze Image",
+                variant="primary",
+                size="lg"
+            )
+    
+    # Results Section
+    diagnosis_output = gr.HTML()
+    
+    submit_btn.click(
+        fn=predict_retinopathy,
+        inputs=image_input,
+        outputs=diagnosis_output
+    )
+    
+    # Understanding Stages
+    gr.Markdown("""
+        ## 📊 Understanding DR Stages and Required Actions
+        
+        ### No DR (Stage 0)
+        **What it means:**
+        * No visible signs of diabetic retinopathy
+        * Normal retinal appearance
+        * Blood vessels appear healthy
+        
+        **Required Actions:**
+        * Continue regular eye check-ups (yearly for diabetic patients)
+        * Maintain blood sugar control
+        * Follow healthy lifestyle recommendations
+        
+        ### Mild DR (Stage 1)
+        **What it means:**
+        * Small areas of balloon-like swelling in the retina's tiny blood vessels
+        * Microaneurysms detected
+        * Early changes that don't typically affect vision
+        
+        **Required Actions:**
+        * Schedule follow-up with an ophthalmologist within 6-12 months
+        * Improve blood sugar control
+        * Monitor blood pressure and cholesterol
+        * Consider more frequent eye screenings
+        
+        ### Moderate DR (Stage 2)
+        **What it means:**
+        * More extensive damage to blood vessels
+        * Blood vessels may begin to leak
+        * Some retinal swelling might be present
+        
+        **Required Actions:**
+        * Consult with an ophthalmologist within 3-6 months
+        * Possible referral to a retina specialist
+        * Strict blood sugar, blood pressure, and cholesterol management
+        * Discuss potential early interventions
+        
+        ### Severe DR (Stage 3)
+        **What it means:**
+        * Significant blockage of blood vessels
+        * Retina being deprived of proper blood supply
+        * Increased risk for vision-threatening complications
+        
+        **Required Actions:**
+        * Urgent consultation with a retina specialist
+        * Possible laser treatment consideration
+        * Intensive monitoring of diabetes management
+        * Prepare for possible treatments
+        
+        ### Proliferative DR (Stage 4)
+        **What it means:**
+        * Advanced disease with new abnormal blood vessels growing
+        * High risk of vitreous hemorrhage
+        * Potential retinal detachment
+        * Significant threat to vision
+        
+        **Required Actions:**
+        * Immediate specialist intervention
+        * Likely need for laser treatment or surgery
+        * Aggressive management of all contributing factors
+        * Regular monitoring post-treatment
+    """)
+    
+    # Technical Details
+    gr.Markdown("""
+        ## 🔬 How Our AI Works
+        
+        * **Image Preprocessing:**
+          * Resizing to 224×224 resolution
+          * Normalization for better CNN performance
+        
+        * **Deep Learning Model:**
+          * Uses ResNet-50, a powerful convolutional neural network
+          * Pretrained on ImageNet for feature extraction
+          * Final layer modified to classify 5 DR severity levels
+        
+        * **Training Process:**
+          * Trained with image augmentation (random flips, rotations, color adjustments)
+          * Optimized using CrossEntropy Loss & Adam Optimizer
+          * Hyperparameter tuning with grid search for optimal learning rate, batch size, and epochs
+        
+        * **Classification:**
+          Based on the extracted features, the AI classifies the image into one of five categories:
+          * **No DR (Stage 0)**
+          * **Mild (Stage 1)**
+          * **Moderate (Stage 2)**
+          * **Severe (Stage 3)**
+          * **Proliferative DR (Stage 4)**
+           
+        ## 🎯 Why It Is Reliable
+        
+        * **Robust Training:** Trained on thousands of diverse retinal images, our model has learned to recognize even the faintest indicators of DR.
+        * **Evaluation Metrics:** Model performance evaluated using accuracy, precision, recall, F1-score, and ROC AUC.
+        * **Rigorous Validation:** Extensive testing and validation protocols ensure the model's predictions are accurate and trustworthy.
+        * **Cutting-Edge Technology:** Utilizing state-of-the-art ResNet-50 architecture, our AI leverages modern deep learning techniques to provide early and reliable health insights.
+    """)
+    
+    # Treatment and Support Information
+    gr.Markdown("""
+        ## 💉 Treatment Approaches
+        
+        ### Standard Treatment Options:
+        * **Laser Treatment (Photocoagulation):** Seals or destroys leaking blood vessels
+        * **Anti-VEGF Injections:** Reduces growth of abnormal blood vessels
+        * **Vitrectomy:** Surgical removal of vitreous gel and blood
+        * **Corticosteroid Injections:** Reduces inflammation and swelling
+        
+        ### Supportive Care:
+        * Blood sugar control
+        * Blood pressure management
+        * Regular eye examinations
+        * Vision rehabilitation if needed
+        * Psychological support
+        
+        ### Long-term Considerations:
+        * Lifestyle modifications
+        * Regular monitoring of diabetes
+        * Adherence to medication regimens
+        * Regular ophthalmologic follow-ups
+    """)
+    
+    # Prevention and Early Detection
+    gr.Markdown("""
+        ## 🛡️ Prevention and Early Detection
+        
+        ### Risk Reduction Strategies:
+        * Maintaining optimal blood sugar levels
+        * Controlling blood pressure
+        * Managing cholesterol levels
+        * Regular exercise
+        * Healthy diet low in refined carbohydrates
+        * Avoiding tobacco use
+        
+        ### When to Seek Medical Attention:
+        * Sudden vision changes
+        * Eye pain or redness
+        * Floaters or spots in vision
+        * Blurred vision
+        * Any noticeable change in visual acuity
+    """)
+    
+    # Future Improvements Section
+    gr.Markdown("""
+        ## 🚀 Future Improvements
+        
+        Our team is continuously working to improve the Retinopathy Scanner with:
+        
+        * **Larger Dataset Integration:** Expanding our training data for better generalization
+        * **Advanced Architectures:** Exploring EfficientNet and Vision Transformers alongside ResNet-50
+        * **Mobile Deployment:** Making the scanner accessible on smartphones for wider reach
+        * **Multi-modal Integration:** Combining fundus images with OCT scans for more comprehensive diagnosis
+        * **Explainable AI Features:** Providing visual explanations for why certain predictions are made
+    """)
+   
+   # Disclaimer
+    gr.Markdown("""
+        ---
+        ### ⚠️ Important Medical Disclaimer
+        
+        This AI-powered tool is designed to assist in the early detection of Diabetic Retinopathy (DR) through 
+        retinal image analysis. We strongly urge users to consult an ophthalmologist for appropriate medical 
+        guidance after getting the diagnosis. 
+        
+        **Please Note:**
+        * Results should be verified by healthcare professionals
+        * Consult a qualified ophthalmologist for proper treatment
+        
+        This initiative is developed by our team at VIT-AP University with the goal of empowering individuals to be more 
+        aware of their eye health before visiting a doctor. Our mission is to leverage AI for early detection and better 
+        healthcare awareness.
+    """)
+
+if __name__ == "__main__":
     demo.launch()