Update model_utils.py

model_utils.py

@@ -8,50 +8,73 @@ 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",
-            num_labels=10,  # Match number of classes
-            ignore_mismatched_sizes=True  # Add this to handle size mismatch
-        )
-        self.feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
-        
-        # Define class labels
-        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
-        }
-        
-        # Set model to evaluation mode
-        self.model.eval()
+        try:
+            # Initialize model and feature extractor
+            self.model = ViTForImageClassification.from_pretrained(
+                "google/vit-base-patch16-224",
+                num_labels=10,
+                ignore_mismatched_sizes=True
+            )
+            self.feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
+            
+            # Set model to evaluation mode
+            self.model.eval()
+            
+            # Define class labels
+            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.
+                """,
+                "Ant": """
+                Ants are social insects that live in colonies. They are incredibly strong
+                for their size and play vital roles in soil health and ecosystem maintenance.
+                """,
+                # Add more species information for other classes...
+            }
+            
+        except Exception as e:
+            raise RuntimeError(f"Error initializing BugClassifier: {str(e)}")
+
+    def preprocess_image(self, image):
+        """Preprocess image for model input"""
+        try:
+            # Convert RGBA to RGB if necessary
+            if image.mode == 'RGBA':
+                image = image.convert('RGB')
+            
+            # Resize image if needed
+            if image.size != (224, 224):
+                image = image.resize((224, 224), Image.Resampling.LANCZOS)
+            
+            # Process image using feature extractor
+            inputs = self.feature_extractor(images=image, return_tensors="pt")
+            return inputs.pixel_values
+            
+        except Exception as e:
+            raise ValueError(f"Error preprocessing image: {str(e)}")
 
     def predict(self, image):
-        """
-        Make a prediction on the input image
-        Returns predicted class and confidence score
-        """
+        """Make a prediction on the input image"""
         try:
-            # Preprocess image
-            if isinstance(image, Image.Image):
-                image_tensor = self.preprocess_image(image)
-            else:
+            if not isinstance(image, Image.Image):
                 raise ValueError("Input must be a PIL Image")
 
+            # Preprocess image
+            image_tensor = self.preprocess_image(image)
+
             # Make prediction
             with torch.no_grad():
                 outputs = self.model(image_tensor)
@@ -60,43 +83,25 @@ class BugClassifier:
                 
                 # Ensure index is within bounds
                 if pred_idx >= len(self.labels):
-                    pred_idx = 0  # Default to first class if out of bounds
+                    pred_idx = 0
                 
             return self.labels[pred_idx], float(probs[pred_idx] * 100)
-        except Exception as e:
-            print(f"Prediction error: {str(e)}")
-            return self.labels[0], 0.0  # Return default prediction in case of error
-
-    def preprocess_image(self, image):
-        """
-        Preprocess image for model input
-        """
-        try:
-            # Convert RGBA to RGB if necessary
-            if image.mode == 'RGBA':
-                image = image.convert('RGB')
             
-            # Process image using feature extractor
-            inputs = self.feature_extractor(images=image, return_tensors="pt")
-            return inputs.pixel_values
         except Exception as e:
-            print(f"Preprocessing error: {str(e)}")
-            raise
+            print(f"Prediction error: {str(e)}")
+            return self.labels[0], 0.0
 
     def get_species_info(self, species):
-        """
-        Return information about a species
-        """
-        return self.species_info.get(species, f"""
+        """Return information about a species"""
+        default_info = f"""
         Information about {species}:
         This species is part of our insect database. While detailed information
         is still being compiled, all insects play important roles in their ecosystems.
-        """)
+        """
+        return self.species_info.get(species, default_info)
 
     def compare_species(self, species1, species2):
-        """
-        Generate comparison information between two species
-        """
+        """Generate comparison information between two species"""
         info1 = self.get_species_info(species1)
         info2 = self.get_species_info(species2)
         
@@ -110,4 +115,57 @@ class BugClassifier:
         {info2}
         
         Both species contribute to their ecosystems in unique ways.
-        """
+        """
+
+    def get_gradcam(self, image):
+        """Generate Grad-CAM visualization for the image"""
+        try:
+            # Preprocess image
+            image_tensor = self.preprocess_image(image)
+            
+            # Get model attention weights (using last layer's attention)
+            with torch.no_grad():
+                outputs = self.model(image_tensor, output_attentions=True)
+                attention = outputs.attentions[-1]  # Get last layer's attention
+                
+            # Convert attention to heatmap
+            attention_map = attention.mean(dim=1).mean(dim=1).numpy()[0]
+            
+            # Resize attention map to image size
+            attention_map = cv2.resize(attention_map, (224, 224))
+            
+            # Normalize attention map
+            attention_map = (attention_map - attention_map.min()) / (attention_map.max() - attention_map.min())
+            
+            # Convert to heatmap
+            heatmap = cv2.applyColorMap(np.uint8(255 * attention_map), cv2.COLORMAP_JET)
+            
+            # Convert original image to RGB numpy array
+            original_image = np.array(image.resize((224, 224)))
+            if len(original_image.shape) == 2:  # Convert grayscale to RGB
+                original_image = cv2.cvtColor(original_image, cv2.COLOR_GRAY2RGB)
+            
+            # Overlay heatmap on original image
+            overlay = cv2.addWeighted(original_image, 0.7, heatmap, 0.3, 0)
+            
+            return Image.fromarray(overlay)
+            
+        except Exception as e:
+            print(f"Error generating Grad-CAM: {str(e)}")
+            return image  # Return original image if Grad-CAM fails
+
+def get_severity_prediction(species):
+    """Predict ecological severity/impact based on species"""
+    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")