import os
import torch
import torch.nn as nn
from torchvision import transforms, models
from PIL import Image
import torch.nn.functional as F
import streamlit as st

class TomatoLeafDiseaseDetectionApp:
    def __init__(self):
        self.class_names = [
            'Pepper__bell___Bacterial_spot', 'Pepper__bell___healthy', 'Potato___Early_blight',
            'Potato___Late_blight', 'Potato___healthy', 'Tomato_Bacterial_spot',
            'Tomato_Early_blight', 'Tomato_Late_blight', 'Tomato_Leaf_Mold',
            'Tomato_Septoria_leaf_spot', 'Tomato_Spider_mites_Two_spotted_spider_mite',
            'Tomato__Target_Spot', 'Tomato__Tomato_YellowLeaf__Curl_Virus',
            'Tomato__Tomato_mosaic_virus', 'Tomato_healthy'
        ]

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = self.load_model()

    def load_model(self):
        """
        Load the trained EfficientNet model with the weights for tomato leaf disease detection.
        """
        # Define the model structure
        base_model = models.efficientnet_b0(weights=None)  # No pretrained weights
        base_model.classifier = nn.Identity()  # Remove the original classifier
        feature_size = 1280  # EfficientNetB0 output feature size

        model = nn.Sequential(
            base_model,
            nn.Dropout(0.3),
            nn.Linear(feature_size, len(self.class_names))
        )

        # Load the model weights
        model_path = "tomato_leaf_disease_model.pth"  # Update this path
        model.load_state_dict(torch.load(model_path, map_location=self.device))
        model.to(self.device)
        model.eval()  # Set the model to evaluation mode
        return model

    def predict_disease(self, image):
        """
        Predict the tomato leaf disease from the given image.

        Args:
            image (PIL.Image): Input image.

        Returns:
            tuple: Predicted disease name and confidence score.
        """
        try:
            # Image preprocessing
            transform = transforms.Compose([
                transforms.Resize((224, 224)),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # Normalize for EfficientNet
            ])
            input_tensor = transform(image).unsqueeze(0).to(self.device)

            # Perform prediction
            with torch.no_grad():
                outputs = self.model(input_tensor)
                probabilities = F.softmax(outputs, dim=1)
                predicted_class = probabilities.argmax(1)
                confidence_score = probabilities[0, predicted_class.item()].item()

            predicted_class_name = self.class_names[predicted_class.item()]
            return predicted_class_name, confidence_score*100
        except Exception as e:
            return f"Error: {str(e)}", 0.0

def main():
    st.title("Tomato Leaf Disease Detection")

    # Upload image
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
    if uploaded_file is not None:
        # Open the image
        image = Image.open(uploaded_file).convert("RGB")
        st.image(image, caption='Uploaded Image.', use_container_width =True)

        # Initialize the app
        app = TomatoLeafDiseaseDetectionApp()

        # Predict disease
        disease_name, confidence = app.predict_disease(image)
        st.write(f"Predicted Disease: {disease_name}")
        st.write(f"Confidence Score: {confidence:.2f}%")

if __name__ == "__main__":
    main()