Update app.py

app.py

@@ -1,117 +1,117 @@
-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 gradio as gr
-
-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_path):
-        """
-        Predict the tomato leaf disease from the given image.
-
-        Args:
-            image_path (str): Path to the 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
-            ])
-            image = Image.open(image_path).convert("RGB")
-            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
-        except Exception as e:
-            return f"Error: {str(e)}", 0.0
-
-    def gradio_interface(self):
-        """
-        Launch the Gradio interface for tomato leaf disease detection.
-        """
-        def classify_image(image_path):
-            disease_name, confidence = self.predict_disease(image_path)
-            return disease_name, f"Confidence: {confidence:.2f}"
-
-        with gr.Blocks(theme=gr.themes.Soft()) as demo:
-            gr.HTML("<center><h1>Tomato Leaf Disease Detection</h1></center>")
-            
-            with gr.Row():
-                input_image = gr.Image(type="filepath", label="Upload Leaf Image", source="upload")
-                with gr.Column():
-                    output_label = gr.Label(label="Predicted Disease")
-                    confidence_text = gr.Textbox(label="Confidence Score")
-            
-            with gr.Row():
-                button = gr.Button(value="Detect Disease")
-            
-            button.click(
-                classify_image,
-                inputs=[input_image],
-                outputs=[output_label, confidence_text]
-            )
-
-            gr.Examples(
-                examples=[
-                    "tomato_earlt_blight.jpg",  # Replace with your example paths
-                    "yellow_leaf_curl.jpg",
-                ],
-                inputs=[input_image],
-                outputs=[output_label, confidence_text],
-                label="Example Images"
-            )
-
-            demo.launch(debug=True)
-
-if __name__ == "__main__":
-    app = TomatoLeafDiseaseDetectionApp()
-    app.gradio_interface()
+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 gradio as gr
+
+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_path):
+        """
+        Predict the tomato leaf disease from the given image.
+
+        Args:
+            image_path (str): Path to the 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
+            ])
+            image = Image.open(image_path).convert("RGB")
+            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
+        except Exception as e:
+            return f"Error: {str(e)}", 0.0
+
+    def gradio_interface(self):
+        """
+        Launch the Gradio interface for tomato leaf disease detection.
+        """
+        def classify_image(image_path):
+            disease_name, confidence = self.predict_disease(image_path)
+            return disease_name, f"Confidence: {confidence:.2f}"
+
+        with gr.Blocks(theme=gr.themes.Soft()) as demo:
+            gr.HTML("<center><h1>Tomato Leaf Disease Detection</h1></center>")
+            
+            with gr.Row():
+                input_image = gr.Image(type="filepath", label="Upload Leaf Image")
+                with gr.Column():
+                    output_label = gr.Label(label="Predicted Disease")
+                    confidence_text = gr.Textbox(label="Confidence Score")
+            
+            with gr.Row():
+                button = gr.Button(value="Detect Disease")
+            
+            button.click(
+                classify_image,
+                inputs=[input_image],
+                outputs=[output_label, confidence_text]
+            )
+
+            gr.Examples(
+                examples=[
+                    "tomato_earlt_blight.jpg",  # Replace with your example paths
+                    "yellow_leaf_curl.jpg",
+                ],
+                inputs=[input_image],
+                outputs=[output_label, confidence_text],
+                label="Example Images"
+            )
+
+            demo.launch(debug=True)
+
+if __name__ == "__main__":
+    app = TomatoLeafDiseaseDetectionApp()
+    app.gradio_interface()