app.py

from ultralytics import YOLO
import gradio as gr
import numpy as np
from PIL import Image
import os

# Load the trained YOLO classification model
model_path = "best.pt"  # Make sure this is your classification model
model = YOLO(model_path)

# Define class names (update these based on your model's classes)
class_names = {
    0: "bacterial_leaf_blight",
    1: "brown_spot", 
    2: "healthy",
    3: "leaf_blast",
    4: "leaf_scald",
    5: "narrow_brown_spot"
}

# Disease descriptions
disease_info = {
    "bacterial_leaf_blight": "Bacterial Leaf Blight is caused by Xanthomonas oryzae. Symptoms include yellowish-green to brown lesions with wavy margins on leaves, often starting from leaf tips and edges.",
    "brown_spot": "Brown Spot is caused by Cochliobolus miyabeanus. Appears as small, circular to oval brown spots with gray or light-colored centers surrounded by yellow halos.",
    "healthy": "This leaf shows no signs of disease and appears to be healthy. Continue monitoring and maintaining proper crop management practices.",
    "leaf_blast": "Leaf Blast is caused by Magnaporthe oryzae. Characterized by diamond-shaped lesions with gray centers and brown borders, often with yellow halos.",
    "leaf_scald": "Leaf Scald is caused by Microdochium oryzae. Shows elongated lesions with irregular margins, typically appearing on leaf sheaths and blades.",
    "narrow_brown_spot": "Narrow Brown Spot is caused by Cercospora janseana. Features narrow, brown lesions that run parallel to the leaf veins."
}

# Treatment recommendations
treatment_recommendations = {
    "bacterial_leaf_blight": """
Treatment Recommendations:
• Use disease-free certified seeds
• Apply copper-based bactericides (copper sulfate, copper hydroxide)
• Improve field drainage to reduce moisture
• Remove and destroy infected plant debris
• Avoid overhead irrigation
• Maintain balanced fertilization (avoid excess nitrogen)
• Plant resistant varieties when available
""",
    "brown_spot": """
Treatment Recommendations:
• Apply fungicides containing azoxystrobin or propiconazole
• Use foliar sprays of mancozeb or chlorothalonil
• Avoid excessive nitrogen fertilization
• Ensure proper field drainage
• Maintain optimal plant spacing for air circulation
• Remove infected leaves and debris
• Monitor silicon levels in soil
""",
    "healthy": """
Maintenance Recommendations:
• Continue regular field monitoring
• Maintain proper plant nutrition and irrigation
• Ensure good field sanitation
• Monitor for early disease symptoms
• Apply preventive fungicide if conditions are favorable for disease
• Maintain proper plant spacing
• Remove weeds that may harbor pathogens
""",
    "leaf_blast": """
Treatment Recommendations:
• Apply fungicides containing tricyclazole, azoxystrobin, or propiconazole
• Use systemic fungicides like carbendazim for severe infections
• Avoid excessive nitrogen application (use balanced fertilization)
• Improve air circulation through proper spacing
• Remove infected plant debris
• Use resistant varieties when available
• Apply silicon fertilizers to strengthen plant defense
""",
    "leaf_scald": """
Treatment Recommendations:
• Apply fungicides containing propiconazole or tebuconazole
• Use copper-based fungicides for early prevention
• Improve air circulation around plants
• Avoid overhead irrigation and minimize leaf wetness
• Remove infected plant materials
• Maintain proper plant nutrition
• Consider using resistant varieties
""",
    "narrow_brown_spot": """
Treatment Recommendations:
• Apply fungicides containing propiconazole, azoxystrobin, or mancozeb
• Use preventive copper-based sprays
• Maintain proper plant spacing for air circulation
• Remove infected plant debris promptly
• Ensure balanced nutrition (avoid nitrogen excess)
• Improve drainage to reduce humidity
• Monitor and control alternate hosts
"""
}

def predict_rice_disease(image):
    """
    Predict rice leaf disease using YOLO classification model
    """
    if image is None:
        return {}, "Please upload an image", "", ""
    
    try:
        # Convert numpy array to PIL Image if needed
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        # Ensure RGB format
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # Run inference using YOLO classification
        results = model(image, verbose=False)
        
        # Extract prediction results
        for r in results:
            # Get prediction probabilities
            probs = r.probs.data.tolist()
            top_class = r.probs.top1
            confidence = r.probs.top1conf.item()
            
            # Create probability dictionary for all classes
            class_probabilities = {}
            for i, prob in enumerate(probs):
                class_name = class_names.get(i, f"class_{i}")
                class_probabilities[class_name.replace('_', ' ').title()] = float(prob)
            
            # Get predicted class name
            predicted_class = class_names.get(top_class, f"class_{top_class}")
            
            # Format result text
            result_text = f"""
Prediction: {predicted_class.replace('_', ' ').title()}
Confidence: {confidence:.2%}
Class Index: {top_class}
            """.strip()
            
            # Get disease information and treatment
            disease_desc = disease_info.get(predicted_class, "No information available")
            treatment = treatment_recommendations.get(predicted_class, "No treatment information available")
            
            return class_probabilities, result_text, disease_desc, treatment
    
    except Exception as e:
        error_msg = f"Error during prediction: {str(e)}"
        print(error_msg)  # For debugging
        return {}, error_msg, "", ""

def batch_predict_demo():
    """
    Demonstrate batch prediction capability
    """
    # This would be used if you have example images
    # For demo purposes, we'll just show how it would work
    demo_text = """
    Batch Prediction Example:
    
    # For batch prediction on multiple images:
    results = model([
        "path/to/image1.jpg",
        "path/to/image2.jpg", 
        "path/to/image3.jpg"
    ])
    
    for i, r in enumerate(results):
        probs = r.probs.data.tolist()
        top_class = r.probs.top1
        confidence = r.probs.top1conf.item()
        print(f"Image {i+1}: Class {top_class}, Confidence: {confidence:.4f}")
    """
    return demo_text

# Create the Gradio interface
with gr.Blocks(
    theme=gr.themes.Soft(),
    title="Rice Leaf Disease Classification",
    css="""
    .gradio-container {
        background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
    }
    .main-header {
        text-align: center;
        color: #2c3e50;
        margin-bottom: 20px;
    }
    """
) as app:
    
    # Header
    gr.Markdown(
        """
        # 🌾 Rice Leaf Disease Classification System
        ## AI-Powered Disease Detection for Rice Crops
        
        Upload an image of a rice leaf to get instant disease classification with confidence scores, 
        detailed information, and treatment recommendations.
        """,
        elem_classes=["main-header"]
    )
    
    # Main interface
    with gr.Row():
        # Left column - Input
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Upload Rice Leaf Image")
            input_image = gr.Image(
                label="Select Image",
                type="numpy",
                height=400,
                elem_id="input-image"
            )
            
            predict_btn = gr.Button(
                "🔍 Analyze Disease",
                variant="primary",
                size="lg",
                elem_id="predict-button"
            )
            
            gr.Markdown("""
            ### 📝 Image Guidelines:
            • Use clear, well-lit photos
            • Focus on diseased areas
            • Avoid blurry images
            • Include full leaf when possible
            """)
        
        # Right column - Results
        with gr.Column(scale=1):
            gr.Markdown("### 📊 Analysis Results")
            
            # Prediction result
            result_output = gr.Textbox(
                label="🎯 Prediction Summary",
                lines=4,
                elem_id="result-text"
            )
            
            # Probability scores
            probability_output = gr.Label(
                label="📈 Confidence Scores",
                num_top_classes=6,
                elem_id="probability-scores"
            )
    
    # Disease information and treatment
    with gr.Row():
        with gr.Column():
            disease_info_output = gr.Textbox(
                label="🦠 Disease Information",
                lines=5,
                elem_id="disease-info"
            )
        
        with gr.Column():
            treatment_output = gr.Textbox(
                label="💊 Treatment Recommendations",
                lines=8,
                elem_id="treatment-info"
            )
    
    # Information sections
    with gr.Accordion("🔬 Detectable Diseases", open=False):
        gr.Markdown("""
        This model can identify the following rice leaf conditions:
        
        1. **Bacterial Leaf Blight** - Bacterial infection causing wavy-margin lesions
        2. **Brown Spot** - Fungal disease with circular brown spots
        3. **Healthy** - Disease-free rice leaves
        4. **Leaf Blast** - Fungal infection causing diamond-shaped lesions  
        5. **Leaf Scald** - Fungal disease with elongated irregular lesions
        6. **Narrow Brown Spot** - Fungal infection with narrow brown streaks
        """)
    
    with gr.Accordion("🔧 Technical Details", open=False):
        gr.Markdown("""
        ### Model Information:
        - **Architecture**: YOLO Classification Model
        - **Input Size**: Automatically resized by model
        - **Output**: Probability scores for each disease class
        - **Confidence Threshold**: Shows confidence percentage for predictions
        
        ### Usage Notes:
        - Higher confidence scores indicate more reliable predictions
        - For critical decisions, always consult agricultural experts
        - Model performance depends on image quality and lighting conditions
        """)
        
        batch_demo = gr.Textbox(
            label="Batch Prediction Code Example",
            value=batch_predict_demo(),
            lines=12,
            max_lines=12,
            interactive=False
        )
    
    # Event handlers
    predict_btn.click(
        fn=predict_rice_disease,
        inputs=[input_image],
        outputs=[probability_output, result_output, disease_info_output, treatment_output]
    )
    
    # Auto-predict when image is uploaded
    input_image.change(
        fn=predict_rice_disease,
        inputs=[input_image],
        outputs=[probability_output, result_output, disease_info_output, treatment_output]
    )
    
    # Footer
    gr.Markdown("""
    ---
    ### ⚠️ Important Disclaimer
    This AI model is designed to assist in rice disease identification. For accurate diagnosis and treatment decisions, 
    please consult with agricultural experts or plant pathologists. The model's predictions should be used as a 
    preliminary assessment tool only.
    """)

# Launch the application
if __name__ == "__main__":
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True,
        show_error=True
    )