Update model_utils.py

model_utils.py

@@ -11,43 +11,120 @@ class BugClassifier:
         try:
             # Initialize model and feature extractor
             self.model = ViTForImageClassification.from_pretrained(
-                "google/vit-base-patch16-224",
+                "microsoft/beit-base-patch16-224-pt22k-ft22k",
                 num_labels=10,
                 ignore_mismatched_sizes=True
             )
-            self.feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
+            
+            # Add custom classification head
+            self.model.classifier = torch.nn.Sequential(
+                torch.nn.Linear(768, 512),
+                torch.nn.ReLU(),
+                torch.nn.Dropout(0.2),
+                torch.nn.Linear(512, 10)  # 10 classes
+            )
+            
+            self.feature_extractor = AutoFeatureExtractor.from_pretrained(
+                "microsoft/beit-base-patch16-224-pt22k-ft22k"
+            )
             
             # Set model to evaluation mode
             self.model.eval()
             
-            # Define class labels
+            # Define detailed class labels
             self.labels = [
-                "Ladybug", "Butterfly", "Ant", "Beetle", "Spider",
-                "Grasshopper", "Moth", "Dragonfly", "Bee", "Wasp"
+                "Seven-spotted Ladybug", "Monarch Butterfly", "Carpenter Ant", 
+                "Japanese Beetle", "Garden Spider", "Green Grasshopper",
+                "Luna Moth", "Common Dragonfly", "Honey Bee", "Paper Wasp"
             ]
             
-            # Species information database
+            # Create a mapping of general categories for better classification
+            self.category_mapping = {
+                "Seven-spotted Ladybug": ["ladybug", "ladybird", "coccinellidae"],
+                "Monarch Butterfly": ["butterfly", "lepidoptera"],
+                "Carpenter Ant": ["ant", "formicidae"],
+                "Japanese Beetle": ["beetle", "coleoptera"],
+                "Garden Spider": ["spider", "arachnid"],
+                "Green Grasshopper": ["grasshopper", "orthoptera"],
+                "Luna Moth": ["moth", "lepidoptera"],
+                "Common Dragonfly": ["dragonfly", "odonata"],
+                "Honey Bee": ["bee", "apidae"],
+                "Paper Wasp": ["wasp", "vespidae"]
+            }
+            
+            # Detailed 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!
+                "Seven-spotted Ladybug": """
+                The Seven-spotted Ladybug (Coccinella septempunctata) is one of the most common ladybug species.
+                These beneficial insects are natural predators of garden pests like aphids and scale insects.
+                Each ladybug can eat up to 5,000 aphids during its lifetime, making them excellent natural pest controllers.
+                Their distinct red coloring with seven black spots serves as a warning to predators.
+                """,
+                "Monarch Butterfly": """
+                The Monarch Butterfly (Danaus plexippus) is known for its spectacular annual migration.
+                These butterflies play a crucial role in pollination and are indicators of ecosystem health.
+                They have a unique relationship with milkweed plants, which their caterpillars exclusively feed on.
+                Their orange and black wings serve as warning colors to predators about their toxicity.
                 """,
-                "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.
+                "Carpenter Ant": """
+                Carpenter Ants (Camponotus spp.) are large ants that build nests in wood.
+                While they don't eat wood like termites, they can cause structural damage to buildings.
+                These social insects live in colonies and play important roles in forest ecosystems,
+                helping to break down dead wood and maintain soil health.
                 """,
-                "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.
+                "Japanese Beetle": """
+                The Japanese Beetle (Popillia japonica) is recognized by its metallic green body.
+                While beautiful, these beetles can be significant garden pests, feeding on many plant species.
+                They are most active in summer months and can be managed through various natural control methods.
+                Their presence often indicates a healthy soil ecosystem, though their feeding can damage plants.
                 """,
-                # Add more species information for other classes...
+                # Add other species info here...
             }
             
         except Exception as e:
             raise RuntimeError(f"Error initializing BugClassifier: {str(e)}")
 
+    def predict(self, image):
+        """Make a prediction on the input image with improved confidence handling"""
+        try:
+            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)
+                probs = F.softmax(outputs.logits, dim=-1).numpy()[0]
+                
+                # Get top 3 predictions
+                top3_idx = np.argsort(probs)[-3:][::-1]
+                top3_probs = probs[top3_idx]
+                
+                # Use confidence threshold
+                CONFIDENCE_THRESHOLD = 0.4  # 40% confidence threshold
+                
+                if top3_probs[0] < CONFIDENCE_THRESHOLD:
+                    # If confidence is too low, return "Unknown"
+                    return "Unknown Insect", float(top3_probs[0] * 100)
+                    
+                # Check if there's a clear winner (significantly higher than second best)
+                if (top3_probs[0] - top3_probs[1]) > 0.2:  # 20% margin
+                    pred_idx = top3_idx[0]
+                else:
+                    # If it's close, consider image quality and features
+                    image_quality = self.assess_image_quality(image)
+                    if image_quality < 0.5:
+                        return "Image Unclear", 0.0
+                    pred_idx = top3_idx[0]
+                
+                return self.labels[pred_idx], float(probs[pred_idx] * 100)
+            
+        except Exception as e:
+            print(f"Prediction error: {str(e)}")
+            return "Error Processing Image", 0.0
+
     def preprocess_image(self, image):
         """Preprocess image for model input"""
         try:
@@ -66,30 +143,35 @@ class BugClassifier:
         except Exception as e:
             raise ValueError(f"Error preprocessing image: {str(e)}")
 
-    def predict(self, image):
-        """Make a prediction on the input image"""
+    def assess_image_quality(self, image):
+        """Assess the quality of the input image"""
         try:
-            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)
-                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
-                
-            return self.labels[pred_idx], float(probs[pred_idx] * 100)
+            # Convert to numpy array
+            img_array = np.array(image)
+            
+            # Check brightness
+            brightness = np.mean(img_array)
+            
+            # Check contrast
+            contrast = np.std(img_array)
+            
+            # Check blur
+            if len(img_array.shape) == 3:
+                gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
+            else:
+                gray = img_array
+            blur_score = cv2.Laplacian(gray, cv2.CV_64F).var()
+            
+            # Normalize and combine scores
+            brightness_score = 1 - abs(brightness - 128) / 128
+            contrast_score = min(contrast / 50, 1)
+            blur_score = min(blur_score / 1000, 1)
+            
+            return (brightness_score + contrast_score + blur_score) / 3
             
         except Exception as e:
-            print(f"Prediction error: {str(e)}")
-            return self.labels[0], 0.0
+            print(f"Error assessing image quality: {str(e)}")
+            return 0.5  # Return middle value if assessment fails
 
     def get_species_info(self, species):
         """Return information about a species"""
@@ -100,33 +182,16 @@ class BugClassifier:
         """
         return self.species_info.get(species, default_info)
 
-    def compare_species(self, species1, species2):
-        """Generate comparison information between two species"""
-        info1 = self.get_species_info(species1)
-        info2 = self.get_species_info(species2)
-        
-        return f"""
-        **Comparing {species1} and {species2}:**
-        
-        {species1}:
-        {info1}
-        
-        {species2}:
-        {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)
+            # Get model attention weights
             with torch.no_grad():
                 outputs = self.model(image_tensor, output_attentions=True)
-                attention = outputs.attentions[-1]  # Get last layer's attention
+                attention = outputs.attentions[-1]
                 
             # Convert attention to heatmap
             attention_map = attention.mean(dim=1).mean(dim=1).numpy()[0]
@@ -142,7 +207,7 @@ class BugClassifier:
             
             # 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
+            if len(original_image.shape) == 2:
                 original_image = cv2.cvtColor(original_image, cv2.COLOR_GRAY2RGB)
             
             # Overlay heatmap on original image
@@ -154,18 +219,37 @@ class BugClassifier:
             print(f"Error generating Grad-CAM: {str(e)}")
             return image  # Return original image if Grad-CAM fails
 
+    def compare_species(self, species1, species2):
+        """Generate comparison information between two species"""
+        info1 = self.get_species_info(species1)
+        info2 = self.get_species_info(species2)
+        
+        return f"""
+        **Comparing {species1} and {species2}:**
+        
+        {species1}:
+        {info1}
+        
+        {species2}:
+        {info2}
+        
+        Both species contribute to their ecosystems in unique ways.
+        """
+
 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"
+        "Seven-spotted Ladybug": "Low",
+        "Monarch Butterfly": "Low",
+        "Carpenter Ant": "Medium",
+        "Japanese Beetle": "High",
+        "Garden Spider": "Low",
+        "Green Grasshopper": "Medium",
+        "Luna Moth": "Low",
+        "Common Dragonfly": "Low",
+        "Honey Bee": "Low",
+        "Paper Wasp": "Medium",
+        "Unknown Insect": "Unknown",
+        "Image Unclear": "Unknown"
     }
-    return severity_map.get(species, "Medium")
+    return severity_map.get(species, "Unknown")