Update model_utils.py

model_utils.py

@@ -9,10 +9,14 @@ 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.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 (these would be replaced with your actual trained classes)
+        # Define class labels
         self.labels = [
             "Ladybug", "Butterfly", "Ant", "Beetle", "Spider",
             "Grasshopper", "Moth", "Dragonfly", "Bee", "Wasp"
@@ -32,96 +36,78 @@ class BugClassifier:
             """,
             # Add more species information as needed
         }
+        
+        # Set model to evaluation mode
+        self.model.eval()
 
     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")
+        try:
+            # 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)
+            # Make prediction
+            with torch.no_grad():
+                outputs = self.model(image_tensor)
+                probs = F.softmax(outputs.logits, dim=-1).numpy()[0]
+                pred_idx = np.argmax(probs)
+                
+                # Ensure index is within bounds
+                if pred_idx >= len(self.labels):
+                    pred_idx = 0  # Default to first class if out of bounds
+                
+            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
         """
-        # 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
+        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
 
     def get_species_info(self, species):
         """
         Return information about a species
         """
-        return self.species_info.get(species, "Information not available for this species.")
+        return self.species_info.get(species, 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.
+        """)
 
     def compare_species(self, species1, species2):
         """
         Generate comparison information between two species
         """
-        # This would be expanded with actual comparison logic
+        info1 = self.get_species_info(species1)
+        info2 = self.get_species_info(species2)
+        
         return f"""
         **Comparing {species1} and {species2}:**
         
-        These species have different characteristics and roles in the ecosystem.
-        {self.get_species_info(species1)}
+        {species1}:
+        {info1}
         
-        {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")
+        {species2}:
+        {info2}
+        
+        Both species contribute to their ecosystems in unique ways.
+        """