feat: add confidence score feature to model predictions

app.py

@@ -14,7 +14,6 @@ import os
 from classifier import GarbageClassifier
 from config import Config
 
-
 # Initialize classifier
 config = Config()
 classifier = GarbageClassifier(config)
@@ -30,14 +29,14 @@ def classify_garbage_impl(image):
     Actual classification implementation
     """
     if image is None:
-        return "Please upload an image", "No image provided"
+        return "Please upload an image", "No image provided", "N/A"
 
     try:
-        classification, full_response = classifier.classify_image(image)
-        return classification, full_response
+        classification, full_response, confidence_score = classifier.classify_image(image)
+        confidence_display = f"{confidence_score}/10"
+        return classification, full_response, confidence_display
     except Exception as e:
-        return "Error", f"Classification failed: {str(e)}"
-
+        return "Error", f"Classification failed: {str(e)}", "0/10"
 
 # Apply GPU decorator based on environment
 if HF_SPACES:
@@ -78,6 +77,11 @@ with gr.Blocks(title="Garbage Classification System") as demo:
                 placeholder="Upload an image and click classify",
             )
 
+            confidence_output = gr.Textbox(
+                label="Confidence Score",
+                placeholder="Confidence score will appear here",
+            )
+
             full_response_output = gr.Textbox(
                 label="Detailed Analysis",
                 placeholder="Detailed reasoning will appear here",
@@ -102,15 +106,15 @@ with gr.Blocks(title="Garbage Classification System") as demo:
     classify_btn.click(
         fn=classify_garbage,
         inputs=image_input,
-        outputs=[classification_output, full_response_output],
+        outputs=[classification_output, full_response_output, confidence_output]
     )
 
     # Auto-classify on image upload
     image_input.change(
         fn=classify_garbage,
         inputs=image_input,
-        outputs=[classification_output, full_response_output],
+        outputs=[classification_output, full_response_output, confidence_output]
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

classifier.py

@@ -5,7 +5,7 @@ import logging
 from typing import Union, Tuple
 from config import Config
 from knowledge_base import GarbageClassificationKnowledge
-
+import re
 
 class GarbageClassifier:
     def __init__(self, config: Config = None):
@@ -86,7 +86,7 @@ class GarbageClassifier:
 
         return processed_image
 
-    def classify_image(self, image: Union[str, Image.Image]) -> Tuple[str, str]:
+    def classify_image(self, image: Union[str, Image.Image]) -> Tuple[str, str, int]:
         """
         Classify garbage in the image
 
@@ -94,7 +94,7 @@ class GarbageClassifier:
             image: PIL Image or path to image file
 
         Returns:
-            Tuple of (classification_result, detailed_analysis)
+            Tuple of (classification_result, detailed_analysis, confidence_score)
         """
         if self.model is None or self.processor is None:
             raise RuntimeError("Model not loaded. Call load_model() first.")
@@ -126,7 +126,7 @@ class GarbageClassifier:
                         {"type": "image", "image": processed_image},
                         {
                             "type": "text",
-                            "text": "Please classify what you see in this image. If it shows garbage/waste items, classify them according to the garbage classification standards. If it shows people, living things, or other non-waste items, classify it as 'Unable to classify' and explain why it's not garbage.",
+                            "text": "Please classify what you see in this image. If it shows garbage/waste items, classify them according to the garbage classification standards. If it shows people, living things, or other non-waste items, classify it as 'Unable to classify' and explain why it's not garbage. Also provide a confidence score from 1-10 indicating how certain you are about your classification.",
                         },
                     ],
                 },
@@ -158,14 +158,87 @@ class GarbageClassifier:
             # Extract reasoning from response
             reasoning = self._extract_reasoning(response)
 
-            return classification, reasoning
+            # Extract confidence score from response
+            confidence_score = self._extract_confidence_score(response, classification)
+
+            return classification, reasoning, confidence_score
 
         except Exception as e:
             self.logger.error(f"Error during classification: {str(e)}")
             import traceback
 
             traceback.print_exc()
-            return "Error", f"Classification failed: {str(e)}"
+            return "Error", f"Classification failed: {str(e)}", 0
+
+
+    def _calculate_confidence_heuristic(self, response_lower: str, classification: str) -> int:
+        """Calculate confidence based on response content and classification type"""
+        base_confidence = 5
+
+        # Confidence indicators (increase confidence)
+        high_confidence_words = ["clearly", "obviously", "definitely", "certainly", "exactly"]
+        medium_confidence_words = ["appears", "seems", "likely", "probably"]
+
+        # Uncertainty indicators (decrease confidence)
+        uncertainty_words = ["might", "could", "possibly", "maybe", "unclear", "difficult"]
+
+        # Adjust based on confidence words
+        for word in high_confidence_words:
+            if word in response_lower:
+                base_confidence += 2
+                break
+
+        for word in medium_confidence_words:
+            if word in response_lower:
+                base_confidence += 1
+                break
+
+        for word in uncertainty_words:
+            if word in response_lower:
+                base_confidence -= 2
+                break
+
+        # Classification-specific adjustments
+        if classification == "Unable to classify":
+            if any(indicator in response_lower for indicator in ["person", "people", "human", "living"]):
+                base_confidence += 1  # High confidence when clearly not waste
+            else:
+                base_confidence -= 1  # Lower confidence for unclear items
+
+        elif classification == "Error":
+            base_confidence = 1
+
+        else:
+            # Check for specific material mentions (increases confidence)
+            specific_materials = ["aluminum", "plastic", "glass", "metal", "cardboard", "paper"]
+            if any(material in response_lower for material in specific_materials):
+                base_confidence += 1
+
+        return min(max(base_confidence, 1), 10)
+
+    def _extract_confidence_score(self, response: str, classification: str) -> int:
+        """Extract confidence score from response or calculate based on classification"""
+        response_lower = response.lower()
+
+        # Look for explicit confidence scores in the response
+        confidence_patterns = [
+            r'confidence[:\s]*(\d+)',
+            r'confident[:\s]*(\d+)',
+            r'certainty[:\s]*(\d+)',
+            r'score[:\s]*(\d+)',
+            r'(\d+)/10',
+            r'(\d+)\s*out\s*of\s*10'
+        ]
+
+        for pattern in confidence_patterns:
+            match = re.search(pattern, response_lower)
+            if match:
+                score = int(match.group(1))
+                return min(max(score, 1), 10)  # Clamp between 1-10
+
+        # If no explicit score found, calculate based on classification indicators
+        return self._calculate_confidence_heuristic(response_lower, classification)
+
 
     def _extract_classification(self, response: str) -> str:
         """Extract the main classification from the response"""
@@ -268,43 +341,43 @@ class GarbageClassifier:
     def _extract_reasoning(self, response: str) -> str:
         """Extract only the reasoning content, removing all formatting markers and classification info"""
         import re
-        
+
         # Remove all formatting markers
         cleaned_response = response.replace("**Classification**:", "")
         cleaned_response = cleaned_response.replace("**Reasoning**:", "")
         cleaned_response = re.sub(r'\*\*.*?\*\*:', '', cleaned_response)  # Remove any **text**: patterns
         cleaned_response = cleaned_response.replace("**", "")  # Remove remaining ** markers
-        
+
         # Remove category names that might appear at the beginning
         categories = self.knowledge.get_categories()
         for category in categories:
             if cleaned_response.strip().startswith(category):
                 cleaned_response = cleaned_response.replace(category, "", 1)
                 break
-        
+
         # Remove common material names that might appear at the beginning
         material_names = [
-            "Glass", "Plastic", "Metal", "Paper", "Cardboard", "Aluminum", 
+            "Glass", "Plastic", "Metal", "Paper", "Cardboard", "Aluminum",
             "Steel", "Iron", "Tin", "Foil", "Wood", "Ceramic", "Fabric",
             "Recyclable Waste", "Food/Kitchen Waste", "Hazardous Waste", "Other Waste"
         ]
-        
+
         # Clean the response
         cleaned_response = cleaned_response.strip()
-        
+
         # Remove material names at the beginning
         for material in material_names:
             if cleaned_response.startswith(material):
                 # Remove the material name and any following punctuation/whitespace
                 cleaned_response = cleaned_response[len(material):].lstrip(" .,;:")
                 break
-        
+
         # Split into sentences and clean up
         sentences = []
-        
+
         # Split by common sentence endings, but keep the endings
         parts = re.split(r'([.!?])\s+', cleaned_response)
-        
+
         # Rejoin parts to maintain sentence structure
         reconstructed_parts = []
         for i in range(0, len(parts), 2):
@@ -313,49 +386,49 @@ class GarbageClassifier:
                 if i + 1 < len(parts):
                     sentence += parts[i + 1]  # Add the punctuation back
                 reconstructed_parts.append(sentence)
-        
+
         for part in reconstructed_parts:
             part = part.strip()
             if not part:
                 continue
-                
+
             # Skip parts that are just category names or material names
             if part in categories or part.rstrip(".,;:") in material_names:
                 continue
-                
+
             # Skip parts that start with category names or material names
             is_category_line = False
             for item in categories + material_names:
                 if part.startswith(item):
                     is_category_line = True
                     break
-            
+
             if is_category_line:
                 continue
-                
+
             # Clean up the sentence
             part = re.sub(r'^[A-Za-z\s]+:', '', part).strip()  # Remove "Category:" type prefixes
-            
+
             if part and len(part) > 3:  # Only keep meaningful content
                 sentences.append(part)
-        
+
         # Join sentences
         reasoning = ' '.join(sentences)
-        
+
         # Final cleanup - remove any remaining standalone material words at the beginning
         reasoning_words = reasoning.split()
         if reasoning_words and reasoning_words[0] in [m.lower() for m in material_names]:
             reasoning_words = reasoning_words[1:]
             reasoning = ' '.join(reasoning_words)
-        
+
         # Ensure proper capitalization
         if reasoning:
             reasoning = reasoning[0].upper() + reasoning[1:] if len(reasoning) > 1 else reasoning.upper()
-            
+
             # Ensure proper punctuation
             if not reasoning.endswith(('.', '!', '?')):
                 reasoning += '.'
-                
+
         return reasoning if reasoning else "Analysis not available"
 
     def get_categories_info(self):

requirements.txt

@@ -5,4 +5,4 @@ torchvision
 transformers >= 4.53
 accelerate
 timm
-gradio
+gradio