feat: add mixed garbage classification rules

README.md

@@ -1,25 +1,76 @@
----
-title: Gemma3n Challenge
-emoji: 🌍
-colorFrom: green
-colorTo: green
-sdk: gradio
-sdk_version: 5.37.0
-app_file: app.py
-pinned: false
-license: mit
----
-
-# Garbage Classification System
-
-A Gradio app for classifying garbage into 4 categories using Gemma-3n model.
-
-## Usage
-1. Upload an image.
-2. Click "Classify Garbage" or wait for auto-classification.
-
-## Categories
-- Recyclable Waste
-- Food/Kitchen Waste
-- Hazardous Waste
-- Other Waste
+# Waste Classification Application
+
+This project is a web-based application that classifies waste materials from user-uploaded images. It identifies the type of waste (e.g., cardboard, glass, metal) and provides information on how to properly dispose of it.
+
+## 🚀 Live Demo
+
+Try the application live on Hugging Face Spaces\!
+
+**➡️ [Waste Classification Demo](https://huggingface.co/spaces/HMWCS/Gemma3n-challenge-demo)**
+
+-----
+
+## ✨ Features
+
+  * **Image-based classification:** Upload an image of a waste item to have it automatically classified.
+  * **Multiple waste categories:** The application can identify a variety of waste materials.
+  * **Disposal information:** After classification, the app provides guidance on how to dispose of the identified waste material.
+  * **Web interface:** A user-friendly web interface built with Gradio makes the application easy to use.
+
+-----
+
+## 💡 How it works
+
+The application uses a pre-trained Gemma3n model to perform the image classification. The model has been fine-tuned on a dataset of waste images to accurately identify different materials. The disposal information is retrieved from a knowledge base within the application.
+
+-----
+
+## 📓 Kaggle Notebook
+
+Explore the model fine-tuning process and the underlying code in our detailed Kaggle Notebook.
+
+**➡️ [Gemma3n Challenge Notebook](https://www.kaggle.com/code/yichuanhuang/gemma3n-garbage-classification)**
+
+-----
+
+## 🛠️ Getting Started
+
+### Prerequisites
+
+  * Python 3.9+
+  * Pip
+  * Cuda (optional)
+
+### Installation
+
+1.  Clone the repository:
+    ```bash
+    git clone https://github.com/yichuan-huang/gemma3n-challenge
+    ```
+2.  Navigate to the project directory:
+    ```bash
+    cd gemma3n-challenge
+    ```
+3.  Install the required dependencies:
+    ```bash
+    pip install -r requirements.txt
+    ```
+
+### Running the application
+
+To start the application, run the following command:
+
+```bash
+python app.py
+```
+
+This will launch a Gradio web server. You can access the application by opening the provided URL in your web browser.
+
+## 📂 Project Structure
+
+  * `app.py`: The main application file, containing the Gradio interface and the classification logic.
+  * `classifier.py`: Handles the image classification using the pre-trained model.
+  * `config.py`: Contains configuration settings for the application, such as the model name and labels.
+  * `knowledge_base.py`: A simple knowledge base containing disposal information for different waste materials.
+  * `requirements.txt`: A list of the Python dependencies required to run the application.
+  * `test_images/`: A directory containing sample images for testing the application.

classifier.py

@@ -7,13 +7,55 @@ from config import Config
 from knowledge_base import GarbageClassificationKnowledge
 import re
 
+
+def preprocess_image(image: Image.Image) -> Image.Image:
+    """
+    Preprocess image to meet Gemma3n requirements (512x512)
+    """
+    # Convert to RGB if necessary
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    # Resize to 512x512 as required by Gemma3n
+    target_size = (512, 512)
+
+    # Calculate aspect ratio preserving resize
+    original_width, original_height = image.size
+    aspect_ratio = original_width / original_height
+
+    if aspect_ratio > 1:
+        # Width is larger
+        new_width = target_size[0]
+        new_height = int(target_size[0] / aspect_ratio)
+    else:
+        # Height is larger or equal
+        new_height = target_size[1]
+        new_width = int(target_size[1] * aspect_ratio)
+
+    # Resize image maintaining aspect ratio
+    image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+    # Create a new image with target size and paste the resized image
+    processed_image = Image.new(
+        "RGB", target_size, (255, 255, 255)
+    )  # White background
+
+    # Calculate position to center the image
+    x_offset = (target_size[0] - new_width) // 2
+    y_offset = (target_size[1] - new_height) // 2
+
+    processed_image.paste(image, (x_offset, y_offset))
+
+    return processed_image
+
+
 class GarbageClassifier:
     def __init__(self, config: Config = None):
         self.config = config or Config()
         self.knowledge = GarbageClassificationKnowledge()
         self.processor = None
         self.model = None
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        # self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
         # Setup logging
         logging.basicConfig(level=logging.INFO)
@@ -46,46 +88,6 @@ class GarbageClassifier:
             self.logger.error(f"Error loading model: {str(e)}")
             raise
 
-    def preprocess_image(self, image: Image.Image) -> Image.Image:
-        """
-        Preprocess image to meet Gemma3n requirements (512x512)
-        """
-        # Convert to RGB if necessary
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-
-        # Resize to 512x512 as required by Gemma3n
-        target_size = (512, 512)
-
-        # Calculate aspect ratio preserving resize
-        original_width, original_height = image.size
-        aspect_ratio = original_width / original_height
-
-        if aspect_ratio > 1:
-            # Width is larger
-            new_width = target_size[0]
-            new_height = int(target_size[0] / aspect_ratio)
-        else:
-            # Height is larger or equal
-            new_height = target_size[1]
-            new_width = int(target_size[1] * aspect_ratio)
-
-        # Resize image maintaining aspect ratio
-        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
-
-        # Create a new image with target size and paste the resized image
-        processed_image = Image.new(
-            "RGB", target_size, (255, 255, 255)
-        )  # White background
-
-        # Calculate position to center the image
-        x_offset = (target_size[0] - new_width) // 2
-        y_offset = (target_size[1] - new_height) // 2
-
-        processed_image.paste(image, (x_offset, y_offset))
-
-        return processed_image
-
     def classify_image(self, image: Union[str, Image.Image]) -> Tuple[str, str, int]:
         """
         Classify garbage in the image
@@ -107,7 +109,7 @@ class GarbageClassifier:
                 raise ValueError("Image must be a PIL Image or file path")
 
             # Preprocess image to meet Gemma3n requirements
-            processed_image = self.preprocess_image(image)
+            processed_image = preprocess_image(image)
 
             # Prepare messages with system prompt and user query
             messages = [
@@ -239,71 +241,94 @@ class GarbageClassifier:
         # 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"""
         response_lower = response.lower()
-        
+
+        # Check for mixed garbage warnings first
+        mixed_garbage_indicators = [
+            "multiple garbage types",
+            "separate items",
+            "mixed together",
+            "different types of garbage"
+        ]
+
+        if any(indicator in response_lower for indicator in mixed_garbage_indicators):
+            return "Unable to classify"
+
+        # Check for contaminated containers that should go to Food/Kitchen Waste
+        contamination_indicators = [
+            "food residue", "contaminated", "not empty", "not rinsed",
+            "tip: empty and rinse", "empty and rinse this container"
+        ]
+
+        if any(indicator in response_lower for indicator in contamination_indicators):
+            # If it mentions recycling tip but has contamination, it's Food/Kitchen Waste
+            return "Food/Kitchen Waste"
+
         # First, look for positive waste category indicators
         # Check exact category matches first
         categories = self.knowledge.get_categories()
         waste_categories = [cat for cat in categories if cat != "Unable to classify"]
-        
+
         for category in waste_categories:
             if category.lower() in response_lower:
                 # Make sure it's not in a negative context
                 category_index = response_lower.find(category.lower())
-                context_before = response_lower[max(0, category_index-30):category_index]
-                
+                context_before = response_lower[max(0, category_index - 30):category_index]
+
                 # Only skip if there's a clear negation right before
                 if not any(neg in context_before[-10:] for neg in ["not", "cannot", "isn't", "doesn't"]):
                     return category
-        
-        # Look for strong recyclable indicators
+
+        # Look for strong recyclable indicators (only if clean/empty)
         recyclable_indicators = [
             "recyclable", "recycle", "aluminum", "plastic", "glass", "metal",
             "foil", "can", "bottle", "cardboard", "paper", "tin", "steel", "iron"
         ]
-        
+
         if any(indicator in response_lower for indicator in recyclable_indicators):
-            # Check if it's explicitly said to be recyclable
-            recyclable_phrases = [
-                "recyclable", "can be recycled", "made of recyclable",
-                "recyclable material", "recyclable aluminum", "recyclable plastic"
-            ]
-            if any(phrase in response_lower for phrase in recyclable_phrases):
-                return "Recyclable Waste"
-            
-            # Check for specific materials
-            if any(material in response_lower for material in ["aluminum", "foil", "metal"]):
-                return "Recyclable Waste"
-            if any(material in response_lower for material in ["plastic", "bottle"]):
-                return "Recyclable Waste"
-            if any(material in response_lower for material in ["glass", "cardboard", "paper"]):
-                return "Recyclable Waste"
-        
+            # Check if it's contaminated or has food content
+            if not any(cont in response_lower for cont in contamination_indicators):
+                # Check if it's explicitly said to be recyclable
+                recyclable_phrases = [
+                    "recyclable", "can be recycled", "made of recyclable",
+                    "recyclable material", "recyclable aluminum", "recyclable plastic",
+                    "clean", "empty", "rinsed"
+                ]
+                if any(phrase in response_lower for phrase in recyclable_phrases):
+                    return "Recyclable Waste"
+
+                # Check for specific materials
+                if any(material in response_lower for material in ["aluminum", "foil", "metal"]):
+                    return "Recyclable Waste"
+                if any(material in response_lower for material in ["plastic", "bottle"]):
+                    return "Recyclable Waste"
+                if any(material in response_lower for material in ["glass", "cardboard", "paper"]):
+                    return "Recyclable Waste"
+
         # Look for food waste indicators
         food_indicators = [
             "food", "fruit", "vegetable", "organic", "kitchen waste",
-            "peel", "core", "scraps", "leftovers"
+            "peel", "core", "scraps", "leftovers", "food content", "food residue"
         ]
         if any(indicator in response_lower for indicator in food_indicators):
             return "Food/Kitchen Waste"
-        
+
         # Look for hazardous waste indicators
         hazardous_indicators = [
             "battery", "chemical", "medicine", "paint", "toxic", "hazardous"
         ]
         if any(indicator in response_lower for indicator in hazardous_indicators):
             return "Hazardous Waste"
-        
+
         # Look for other waste indicators
         other_waste_indicators = [
             "cigarette", "ceramic", "dust", "diaper", "tissue", "other waste"
         ]
         if any(indicator in response_lower for indicator in other_waste_indicators):
             return "Other Waste"
-        
+
         # Only classify as "Unable to classify" if there are explicit indicators
         unable_phrases = [
             "unable to classify",
@@ -311,29 +336,29 @@ class GarbageClassifier:
             "cannot be classified as waste",
             "not garbage", "not waste", "not trash"
         ]
-        
+
         if any(phrase in response_lower for phrase in unable_phrases):
             return "Unable to classify"
-        
+
         # Check for non-garbage items (people, living things, etc.)
         non_garbage_indicators = [
             "person", "people", "human", "face", "man", "woman",
             "living", "alive", "animal", "pet",
             "portrait", "photo of a person"
         ]
-        
+
         if any(indicator in response_lower for indicator in non_garbage_indicators):
             return "Unable to classify"
-        
+
         # If we found waste-related content but no clear category, try to infer
         waste_related = any(word in response_lower for word in [
             "waste", "trash", "garbage", "discard", "throw", "bin"
         ])
-        
+
         if waste_related:
             # Default to Other Waste if it's clearly waste but unclear category
             return "Other Waste"
-        
+
         # If no clear classification found and no clear non-waste indicators, 
         # default to "Unable to classify"
         return "Unable to classify"

knowledge_base.py

@@ -5,6 +5,17 @@ class GarbageClassificationKnowledge:
 
 IMPORTANT: You should ONLY classify items that are actually garbage/waste. If the image contains people, living things, furniture, electronics in use, or other non-waste items, you should classify it as "Unable to classify" and explain that it's not garbage.
 
+**MIXED GARBAGE HANDLING RULES:**
+
+1. **Containers with Food Content**: For any container (bottles, cans, boxes, wrappers) that contains visible food residue or content:
+   - Classify as "Food/Kitchen Waste" due to contamination risk
+   - Always include this warning: "⚠️ Tip: Empty and rinse this container first, then it can be recycled!"
+   - Only completely empty and rinsed containers qualify as "Recyclable Waste"
+
+2. **Multiple Different Garbage Types**: If the image shows multiple different types of garbage mixed together (e.g., electronics with food, batteries with organic waste):
+   - Classify as "Unable to classify"
+   - Include warning: "⚠️ Warning: Multiple garbage types detected. Please separate items for proper classification."
+
 Garbage classification standards:
 
 **Recyclable Waste**:
@@ -13,12 +24,14 @@ Garbage classification standards:
 - Metals: aluminum cans, tin cans, toothpaste tubes, metal toys, metal stationery, nails, metal sheets, aluminum foil, etc.
 - Glass: glass bottles, broken glass pieces, mirrors, light bulbs, vacuum flasks, etc.
 - Textiles: old clothing, textile products, shoes, curtains, towels, bags, etc.
+- NOTE: Only clean, empty containers qualify. Contaminated containers go to Food/Kitchen Waste.
 
 **Food/Kitchen Waste**:
 - Food scraps: rice, noodles, bread, meat, fish, shrimp shells, crab shells, bones, etc.
 - Fruit peels and cores: watermelon rinds, apple cores, orange peels, banana peels, nut shells, etc.
 - Plants: withered branches and leaves, flowers, traditional Chinese medicine residue, etc.
 - Expired food: expired canned food, cookies, candy, etc.
+- Contaminated containers: any container with visible food residue or content
 
 **Hazardous Waste**:
 - Batteries: dry batteries, rechargeable batteries, button batteries, and all types of batteries
@@ -39,13 +52,17 @@ Garbage classification standards:
 - Furniture, appliances, electronics in normal use
 - Buildings, landscapes, vehicles
 - Any item that is not intended to be discarded as waste
+- Multiple different garbage types mixed together
 
 Please observe the items in the image carefully according to the above classification standards. If the image shows garbage/waste items, provide accurate garbage classification results. If the image does NOT show garbage/waste (e.g., people, living things, functioning items), classify it as "Unable to classify" and explain why it's not garbage.
 
-Format your response as:
+For mixed garbage situations, apply the special handling rules above and include appropriate warnings.
+
+Format your response EXACTLY as follows:
 
 **Classification**: [Category Name or "Unable to classify"]
-**Reasoning**: [Brief explanation of why this item belongs to this category, or why it cannot be classified as garbage]"""
+**Reasoning**: [Brief explanation of why this item belongs to this category, or why it cannot be classified as garbage]
+**Confidence Score**: [Number from 1-10]"""
 
     @staticmethod
     def get_categories():
@@ -60,9 +77,9 @@ Format your response as:
     @staticmethod
     def get_category_descriptions():
         return {
-            "Recyclable Waste": "Items that can be processed and reused, including paper, plastic, metal, glass, and textiles",
-            "Food/Kitchen Waste": "Organic waste from food preparation and consumption",
+            "Recyclable Waste": "Items that can be processed and reused, including paper, plastic, metal, glass, and textiles (must be clean and empty)",
+            "Food/Kitchen Waste": "Organic waste from food preparation and consumption, including contaminated containers",
             "Hazardous Waste": "Items containing harmful substances that require special disposal",
             "Other Waste": "Items that don't fit into other categories and go to general waste",
-            "Unable to classify": "Items that are not garbage/waste, such as people, living things, or functioning objects",
+            "Unable to classify": "Items that are not garbage/waste, such as people, living things, functioning objects, or mixed garbage types",
         }