Fix model input

plant_disease_classifier.py

@@ -94,12 +94,16 @@ class PlantDiseaseClassifier:
         accuracy = (correct / total) * 100 if total > 0 else 0.0
         return accuracy
 
-    def predict_image(self, image_path):
-        # Load and transform the image
-        image = Image.open(image_path).convert('RGB')
-        transformed_image = self.data_transforms(image).unsqueeze(0).to(self.device)
-
-        # Predict
+    def predict(self, image):
+    # Ensure the image is in RGB format if not already
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+        
+        # Transform the image to match the model's input requirements
+        transformed_image = self.data_transforms(image).unsqueeze(0)
+        transformed_image = transformed_image.to(self.device)
+        
+        # Make prediction
         with torch.no_grad():
             outputs = self.model(transformed_image)
             logits = outputs.logits if self.model_type in ["levit"] else outputs