File size: 3,582 Bytes
fdff8ca 1957ca3 fdff8ca 1957ca3 fdff8ca 1957ca3 fdff8ca 8e93315 fdff8ca 1957ca3 fdff8ca 1957ca3 7911342 fdff8ca 1957ca3 fdff8ca 1957ca3 8e93315 fdff8ca 1957ca3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
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()
|