fix pre_process function

ui/app.py

@@ -28,57 +28,72 @@ class PlantDiseaseApp:
         self.flagged_predictions = []
 
     def predict(self, image, modelName, confidence_threshold):
+        """
+        Predict plant disease from a single image.
+
+        Args:
+            image: PIL Image or numpy array from Gradio upload
+            modelName: Name of the model to use
+            confidence_threshold: float (0-100), only show predictions above this confidence
+
+        Returns:
+            display_predictions: dict, class_name -> probability
+            result_text: str, formatted top prediction info
+            raw_predictions: str, JSON-formatted top predictions
+        """
         if image is None:
             return None, "Please upload an image", ""
 
         try:
+            # Load model if needed
             if modelName != self.current_modelName:
-                self.model = self.model_loader.loadModel(modelName)
+                self.model, self.class_names = self.model_loader.loadModel(modelName)
                 self.current_modelName = modelName
 
             # Preprocess image
-            tensor = preprocess_image(image)
-            tensor = tensor.to(self.model_loader.device)
+            tensor = preprocess_image(image).to(self.model_loader.device)
 
-            # Get prediction
+            # Model inference
             with torch.no_grad():
                 logits = self.model(tensor)
 
-            # Postprocess
-            top_predictions, all_predictions = postprocess_predictions(
-                logits, config.CLASS_NAMES, config.TOP_K_PREDICTIONS
-            )
+            # Convert logits to probabilities
+            probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy()[0]
+
+            # Map to class names
+            predictions = {name: float(prob) for name, prob in zip(self.class_names, probs)}
 
             # Filter by confidence threshold
-            filtered_predictions = {
-                k: v for k, v in top_predictions.items() if v >= confidence_threshold / 100
-            }
+            filtered_predictions = {k: v for k, v in predictions.items() if v >= confidence_threshold / 100.0}
 
-            # Get top prediction info
+            # Top prediction info
             if filtered_predictions:
                 top_class = max(filtered_predictions.items(), key=lambda x: x[1])[0]
                 top_prob = filtered_predictions[top_class]
                 disease_info = get_disease_info(top_class)
 
                 result_text = f"""
-                **Top Prediction:** {disease_info['formatted_name']}
-                **Confidence:** {top_prob*100:.2f}%
-                **Plant:** {disease_info['plant']}
-                **Status:** {'Healthy' if disease_info['is_healthy'] else 'Disease Detected'}
-                """
+                    **Top Prediction:** {disease_info['formatted_name']}
+                    **Confidence:** {top_prob*100:.2f}%
+                    **Plant:** {disease_info['plant']}
+                    **Status:** {'Healthy' if disease_info['is_healthy'] else 'Disease Detected'}
+                    """
             else:
                 result_text = "No predictions above confidence threshold"
 
             # Format for Gradio Label component
-            display_predictions = {
-                format_class_name(k): v for k, v in filtered_predictions.items()
-            }
+            display_predictions = {format_class_name(k): v for k, v in filtered_predictions.items()}
+
+            # Raw JSON output
+            import json
+            raw_predictions = json.dumps(filtered_predictions, indent=2)
 
-            return display_predictions, result_text, json.dumps(filtered_predictions, indent=2)
+            return display_predictions, result_text, raw_predictions
 
         except Exception as e:
             return None, f"Error during prediction: {str(e)}", ""
 
+
     def flag_prediction(self, image, result_info, feedback_text):
         if image is None:
             return "No image uploaded."

ui/utils.py

@@ -6,98 +6,61 @@ import torch
 import numpy as np
 from PIL import Image
 import torchvision.transforms as transforms
-import config
 
+IMAGE_SIZE = (224, 224)
 
-def preprocess_image(image, image_size=config.IMAGE_SIZE):
-    """
-    Preprocess image for model input
+NORMALIZE_MEAN = [0.485, 0.456, 0.406]
+NORMALIZE_STD = [0.229, 0.224, 0.225]
 
-    Args:
-        image: PIL Image or numpy array
-        image_size: Target size (height, width)
+CLASS_NAMES = []
+TOP_K_PREDICTIONS = 5
+CONFIDENCE_THRESHOLD = 0.01
 
-    Returns:
-        Preprocessed tensor ready for model
+
+def preprocess_image(image):
+    """
+    Preprocess image for model input
     """
-    # Convert to PIL Image if numpy array
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image.astype('uint8'))
 
-    # Convert RGBA to RGB if necessary
     if image.mode == 'RGBA':
         image = image.convert('RGB')
 
-    # Define preprocessing transforms
     transform = transforms.Compose([
-        transforms.Resize(image_size),
+        transforms.Resize(IMAGE_SIZE),
         transforms.ToTensor(),
-        transforms.Normalize(mean=config.NORMALIZE_MEAN, std=config.NORMALIZE_STD)
+        transforms.Normalize(NORMALIZE_MEAN, NORMALIZE_STD)
     ])
 
-    # Apply transforms
     tensor = transform(image)
-
-    # Add batch dimension
-    tensor = tensor.unsqueeze(0)
-
-    return tensor
+    return tensor.unsqueeze(0)
 
 
-def postprocess_predictions(logits, class_names=config.CLASS_NAMES, top_k=config.TOP_K_PREDICTIONS):
+def postprocess_predictions(logits, class_names=CLASS_NAMES, top_k=TOP_K_PREDICTIONS):
     """
-    Convert model logits to human-readable predictions
-
-    Args:
-        logits: Raw model output
-        class_names: List of class names
-        top_k: Number of top predictions to return
-
-    Returns:
-        Dictionary of predictions with confidences
+    Convert logits to formatted predictions
     """
-    # Convert logits to probabilities using softmax
     probs = torch.nn.functional.softmax(logits, dim=1)
-
-    # Convert to numpy
     probs = probs.cpu().detach().numpy()[0]
 
-    # Create predictions dictionary
     predictions = {name: float(prob) for name, prob in zip(class_names, probs)}
-
-    # Get top-k predictions
     top_predictions = sorted(predictions.items(), key=lambda x: x[1], reverse=True)[:top_k]
 
     return dict(top_predictions), predictions
 
 
-def format_prediction_for_display(predictions):
+def format_prediction_for_display(predictions, confidence_threshold=CONFIDENCE_THRESHOLD):
     """
-    Format predictions for Gradio display
-
-    Args:
-        predictions: Dictionary of class names and probabilities
-
-    Returns:
-        Dictionary formatted for Gradio Label component
+    Filter predictions for Gradio display
     """
-    # Filter out very low confidence predictions
-    filtered = {k: v for k, v in predictions.items() if v >= config.CONFIDENCE_THRESHOLD}
-
-    return filtered
+    return {k: v for k, v in predictions.items() if v >= confidence_threshold}
 
 
 def format_class_name(class_name):
     """
-    Format class name for better readability
-
-    Args:
-        class_name: Original class name (e.g., "Tomato___Late_blight")
-
-    Returns:
-        Formatted class name (e.g., "Tomato - Late blight")
+    Format class name into readable form
     """
-    # Replace underscores with spaces and split on ___
     parts = class_name.split("___")
 
     if len(parts) == 2:
@@ -105,74 +68,48 @@ def format_class_name(class_name):
         plant = plant.replace("_", " ")
         disease = disease.replace("_", " ")
         return f"{plant} - {disease}"
-    else:
-        return class_name.replace("_", " ")
+
+    return class_name.replace("_", " ")
 
 
 def get_disease_info(class_name):
     """
-    Get information about a disease (for future enhancement)
-
-    Args:
-        class_name: Disease class name
-
-    Returns:
-        Dictionary with disease information
+    Extract structured disease info from class name
     """
-    # This is a placeholder - you could expand this with actual disease information
     parts = class_name.split("___")
 
-    info = {
+    return {
         "plant": parts[0].replace("_", " ") if len(parts) > 0 else "Unknown",
         "disease": parts[1].replace("_", " ") if len(parts) > 1 else "Unknown",
         "is_healthy": "healthy" in class_name.lower(),
         "formatted_name": format_class_name(class_name)
     }
 
-    return info
-
 
 def batch_preprocess_images(images):
     """
-    Preprocess multiple images for batch prediction
-
-    Args:
-        images: List of PIL Images or numpy arrays
-
-    Returns:
-        Batched tensor ready for model
+    Preprocess a list of images into a batch tensor
     """
     tensors = [preprocess_image(img) for img in images]
-    batch = torch.cat(tensors, dim=0)
-    return batch
+    return torch.cat(tensors, dim=0)
 
 
 def create_confidence_label(predictions, top_k=5):
     """
-    Create a formatted string showing top predictions
-
-    Args:
-        predictions: Dictionary of predictions
-        top_k: Number of top predictions to show
-
-    Returns:
-        Formatted string
+    Render a formatted multiline prediction list
     """
     top_preds = sorted(predictions.items(), key=lambda x: x[1], reverse=True)[:top_k]
 
-    lines = []
-    for i, (class_name, prob) in enumerate(top_preds, 1):
-        formatted_name = format_class_name(class_name)
-        lines.append(f"{i}. {formatted_name}: {prob*100:.2f}%")
-
+    lines = [
+        f"{i}. {format_class_name(name)}: {prob*100:.2f}%"
+        for i, (name, prob) in enumerate(top_preds, 1)
+    ]
     return "\n".join(lines)
 
 
 if __name__ == "__main__":
-    # Test utilities
     print("Testing utility functions...")
 
-    # Test class name formatting
     test_names = [
         "Tomato___Late_blight",
         "Apple___healthy",
@@ -183,7 +120,6 @@ if __name__ == "__main__":
     for name in test_names:
         print(f"  {name} -> {format_class_name(name)}")
 
-    # Test disease info
     print("\nDisease info:")
     for name in test_names:
         info = get_disease_info(name)
@@ -192,19 +128,8 @@ if __name__ == "__main__":
         print(f"    Disease: {info['disease']}")
         print(f"    Healthy: {info['is_healthy']}")
 
-    # Test image preprocessing
     print("\nImage preprocessing:")
     dummy_image = Image.new('RGB', (512, 512), color='red')
     tensor = preprocess_image(dummy_image)
     print(f"  Input size: {dummy_image.size}")
     print(f"  Output tensor shape: {tensor.shape}")
-
-    # Test mock predictions
-    print("\nMock predictions:")
-    from models.mock_model import create_mock_predictions
-    preds = create_mock_predictions(config.CLASS_NAMES)
-    top_preds, all_preds = postprocess_predictions(
-        torch.tensor([list(preds.values())]),
-        config.CLASS_NAMES
-    )
-    print(create_confidence_label(top_preds))