Create model_utils.py

model_utils.py

@@ -0,0 +1,127 @@
+import torch
+import torch.nn.functional as F
+from transformers import ViTForImageClassification, AutoFeatureExtractor
+import numpy as np
+from PIL import Image
+import cv2
+from scipy.special import softmax
+
+class BugClassifier:
+    def __init__(self):
+        # Initialize model and feature extractor
+        self.model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")
+        self.feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
+        
+        # Define class labels (these would be replaced with your actual trained classes)
+        self.labels = [
+            "Ladybug", "Butterfly", "Ant", "Beetle", "Spider",
+            "Grasshopper", "Moth", "Dragonfly", "Bee", "Wasp"
+        ]
+        
+        # Species information database
+        self.species_info = {
+            "Ladybug": """
+            Ladybugs are small, round beetles known for their distinctive spotted patterns.
+            They are beneficial insects that feed on plant-damaging pests like aphids.
+            Fun fact: The number of spots on a ladybug can indicate its species!
+            """,
+            "Butterfly": """
+            Butterflies are beautiful insects known for their large, colorful wings.
+            They play a crucial role in pollination and are indicators of ecosystem health.
+            They undergo complete metamorphosis from caterpillar to adult.
+            """,
+            # Add more species information as needed
+        }
+
+    def predict(self, image):
+        """
+        Make a prediction on the input image
+        Returns predicted class and confidence score
+        """
+        # Preprocess image
+        if isinstance(image, Image.Image):
+            image_tensor = self.preprocess_image(image)
+        else:
+            raise ValueError("Input must be a PIL Image")
+
+        # Make prediction
+        with torch.no_grad():
+            outputs = self.model(image_tensor)
+            probs = softmax(outputs.logits.numpy()[0])
+            pred_idx = np.argmax(probs)
+            
+        return self.labels[pred_idx], float(probs[pred_idx] * 100)
+
+    def preprocess_image(self, image):
+        """
+        Preprocess image for model input
+        """
+        # Resize image if needed
+        if image.size != (224, 224):
+            image = image.resize((224, 224))
+        
+        # Convert to tensor using feature extractor
+        inputs = self.feature_extractor(images=image, return_tensors="pt")
+        return inputs.pixel_values
+
+    def get_species_info(self, species):
+        """
+        Return information about a species
+        """
+        return self.species_info.get(species, "Information not available for this species.")
+
+    def compare_species(self, species1, species2):
+        """
+        Generate comparison information between two species
+        """
+        # This would be expanded with actual comparison logic
+        return f"""
+        **Comparing {species1} and {species2}:**
+        
+        These species have different characteristics and roles in the ecosystem.
+        {self.get_species_info(species1)}
+        
+        {self.get_species_info(species2)}
+        """
+
+def generate_gradcam(image, model):
+    """
+    Generate Grad-CAM visualization for the image
+    """
+    # This is a simplified version - you would need to implement the actual Grad-CAM logic
+    # For now, we'll return a simple heatmap overlay
+    img_array = np.array(image)
+    heatmap = cv2.applyColorMap(
+        cv2.resize(np.random.rand(7,7) * 255, (224, 224)).astype(np.uint8),
+        cv2.COLORMAP_JET
+    )
+    
+    # Overlay heatmap on original image
+    overlay = cv2.addWeighted(
+        cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR),
+        0.7,
+        heatmap,
+        0.3,
+        0
+    )
+    
+    return Image.fromarray(cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB))
+
+def get_severity_prediction(species):
+    """
+    Predict ecological severity/impact based on species
+    """
+    # This would be replaced with actual severity prediction logic
+    severity_map = {
+        "Ladybug": "Low",
+        "Butterfly": "Low",
+        "Ant": "Medium",
+        "Beetle": "Medium",
+        "Spider": "Low",
+        "Grasshopper": "Medium",
+        "Moth": "Low",
+        "Dragonfly": "Low",
+        "Bee": "Low",
+        "Wasp": "Medium"
+    }
+    return severity_map.get(species, "Medium")