app.py

"""
Gradio Web UI for Fish Disease Detection
Enhanced with Gemini-powered treatment and Grad-CAM explainability
"""

# Load environment variables from .env file (for local development)
try:
    from dotenv import load_dotenv
    load_dotenv()
    print("✅ Environment variables loaded from .env")
except ImportError:
    print("⚠️  python-dotenv not installed (OK for production deployment)")

import gradio as gr
import google.generativeai as genai
from PIL import Image
from backend import config as cfg
from backend.predictor import FishDiseasePredictor
from backend.treatment import TreatmentGenerator
from backend.gradcam import generate_gradcam_visualization

# ==================== GEMINI SETUP FOR TREATMENT ====================
treatment_gemini_model = None

if cfg.GEMINI_API_KEY:
    try:
        genai.configure(api_key=cfg.GEMINI_API_KEY)
        treatment_gemini_model = genai.GenerativeModel(cfg.GEMINI_MODEL_NAME)
        print("✅ Gemini AI enabled for treatment recommendations")
    except Exception as e:
        print(f"❌ Gemini setup failed: {e}")
        treatment_gemini_model = None
else:
    print("⚠️  Gemini API key not found. Treatment recommendations will be limited.")

# Initialize treatment generator
treatment_generator = TreatmentGenerator(treatment_gemini_model)

# Initialize predictor
config_dict = {
    'CLASSES': cfg.CLASSES,
    'MODEL_PATH': cfg.MODEL_PATH,
    'DEVICE': cfg.DEVICE,
    'CONFIDENCE_THRESHOLD': cfg.CONFIDENCE_THRESHOLD,
    'MAX_FILE_SIZE_MB': cfg.MAX_FILE_SIZE_MB,
    'MIN_IMAGE_SIZE_PX': cfg.MIN_IMAGE_SIZE_PX,
    'VALID_EXTENSIONS': cfg.VALID_EXTENSIONS
}

predictor = FishDiseasePredictor(config_dict, gemini_model=None)

def predict_fish_disease(image):
    """
    Main prediction function with Grad-CAM visualization
    
    Args:
        image: PIL Image from Gradio
        
    Returns:
        tuple: (result_text, probability_chart, treatment_text, gradcam_image)
    """
    if image is None:
        return "⚠️ Please upload an image", None, "", None
    
    # Run prediction
    result = predictor.predict_from_image(image)
    
    if not result['success']:
        return f"❌ {result['error']}", None, "", None
    
    pred = result['prediction']
    disease = pred['disease'].replace('_', ' ')
    confidence = pred['confidence']
    display_confidence = min(confidence, 100.0)
    
    # Get sorted probabilities
    sorted_probs = sorted(pred['probabilities'].items(), key=lambda x: x[1], reverse=True)
    
    # GENERATE GRAD-CAM VISUALIZATION
    gradcam_image = None
    try:
        predicted_class_idx = cfg.CLASSES.index(pred['disease'])
        model = predictor.model_loader.model
        transform = predictor.model_loader.transform
        
        gradcam_image = generate_gradcam_visualization(
            model, image, predicted_class_idx, transform
        )
    except Exception as e:
        print(f"⚠️ Grad-CAM generation failed: {e}")
        gradcam_image = image  # Fallback to original
    
    # ENHANCED: Better handling for low confidence cases
    if pred['below_threshold']:
        result_text = f"""
## 🐟 Prediction Results

**Status:** ❓ **Uncertain - Low Confidence Detection**

⚠️ **Below confidence threshold ({cfg.CONFIDENCE_THRESHOLD}%)**

The model detected possible disease signs but cannot confidently identify the specific disease.

### 📊 Most Likely Candidates:
1. **{sorted_probs[0][0].replace('_', ' ')}**: {sorted_probs[0][1]:.1f}%
2. **{sorted_probs[1][0].replace('_', ' ')}**: {sorted_probs[1][1]:.1f}%
3. **{sorted_probs[2][0].replace('_', ' ')}**: {sorted_probs[2][1]:.1f}%

### 💡 Recommended Actions:
- Upload a **clearer, well-lit** image if available
- Capture the fish from **different angles**
- **Consult a fish health professional** for accurate diagnosis
- **Monitor the fish** closely for symptom changes
"""
        
        ai_treatment = treatment_generator.get_recommendations(pred['disease'], display_confidence)
        
        treatment_text = f"""### ⚠️ Low Confidence Warning ({display_confidence:.1f}%)

Due to uncertain diagnosis, these are **general guidelines** for the top candidate:

---

{ai_treatment}

---

### 🔴 CRITICAL REMINDER:
This diagnosis has **low confidence** and should **NOT** be used for treatment decisions without professional confirmation.

**Next Steps:**
1. Get a professional veterinary assessment
2. Document symptoms with photos/video
3. Monitor water quality parameters
4. Isolate affected fish if condition worsens"""
        
    else:
        status_emoji = "✅" if pred['disease'] == 'Healthy_Fish' else "⚠️"
        status_text = "Healthy" if pred['disease'] == 'Healthy_Fish' else "Diseased"
        
        result_text = f"""
## 🐟 Prediction Results

**Disease:** {disease}  
**Confidence:** {display_confidence:.2f}%  
**Status:** {status_emoji} {status_text}

{"✅ High confidence detection - Results are reliable" if confidence >= 80 else ""}
"""
        
        treatment_text = treatment_generator.get_recommendations(pred['disease'], display_confidence)
    
    # Convert probabilities for Gradio
    prob_data = {
        disease_name: prob / 100.0
        for disease_name, prob in pred['probabilities'].items()
    }
    
    return result_text, prob_data, treatment_text, gradcam_image

# ==================== GRADIO INTERFACE ====================
with gr.Blocks(title="Fish Disease Detection", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # 🐟 Fish Disease Detection System
    ### AI-Powered Fish Health Diagnosis using VGG16 CNN with Explainable AI
    
    Upload a fish image to detect diseases and see **how the AI made its decision** with heatmap visualization.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="pil", label="Upload Fish Image", height=400)
            predict_btn = gr.Button("🔬 Analyze Fish", variant="primary", size="lg")
            
            gr.Markdown("""
            ### 📋 Instructions:
            1. Upload a **clear fish image**
            2. Click **'Analyze Fish'**
            3. View **results, AI explanation, and treatment**
            
            **Supported:** JPG, PNG (Max 10MB)
            
            ### 💡 Tips for Best Results:
            - Use **well-lit** images
            - Show the **whole fish** clearly
            - Avoid **blurry** or **obstructed** shots
            """)
        
        with gr.Column(scale=1):
            result_output = gr.Markdown(label="Results")
            prob_output = gr.Label(label="Disease Probabilities", num_top_classes=8)
    
    # Grad-CAM Visualization Row
    with gr.Row():
        gradcam_output = gr.Image(
            label="🔍 AI Decision Heatmap - Shows which areas the model focused on for diagnosis (Red = High importance)",
            type="pil",
            height=400
        )
    
    with gr.Row():
        treatment_output = gr.Textbox(
            label="💊 Treatment Recommendations",
            lines=15,
            max_lines=25
        )
    
    # Example images
    gr.Examples(
        examples=[
            ["data/merged_dataset-all/test/Healthy_Fish/Healthy_Fish_1__1.jpg"],
            ["data/merged_dataset-all/test/Bacterial_gill_disease/Bacterial_gill_disease_1__1.jpg"],
        ],
        inputs=image_input,
        label="📸 Example Images"
    )
    
    gr.Markdown("""
    ---
    ### 📊 Model Information
    - **Architecture:** VGG16 CNN with Transfer Learning
    - **Test Accuracy:** 98.65%
    - **Training Dataset:** 5,000+ annotated images
    - **Disease Classes:** 8 diseases + Healthy
    - **Confidence Threshold:** 70% (for reliable diagnosis)
    - **Inference Time:** ~2-3 seconds
    - **AI Treatment:** Powered by Google Gemini 2.0
    - **Explainability:** Grad-CAM visualization
    
    ### 🎯 Confidence Levels:
    - **≥ 80%**: High confidence - Very reliable
    - **70-79%**: Good confidence - Reliable
    - **< 70%**: Low confidence - Requires verification
    
    ### 🔬 Understanding the Heatmap:
    - **🔴 Red areas**: Model focused here (disease symptoms, lesions, abnormalities)
    - **🟡 Yellow areas**: Moderate importance
    - **🟢 Green/Blue areas**: Less important for diagnosis
    
    The heatmap shows the AI is making decisions based on actual disease features, not random patterns.
    
    ---
    
    ### ⚠️ Medical Disclaimer
    This is an **AI diagnostic tool** for preliminary screening only. 
    
    **Always consult a qualified aquaculture veterinarian for:**
    - Professional diagnosis confirmation
    - Treatment plan approval
    - Medication dosage recommendations
    - Emergency health situations
    
    This tool is intended to **assist**, not **replace** professional veterinary care.
    """)
    
    # Connect button with 4 outputs
    predict_btn.click(
        fn=predict_fish_disease,
        inputs=image_input,
        outputs=[result_output, prob_output, treatment_output, gradcam_output]
    )

# Launch
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860
    )