update

README.md

@@ -4,40 +4,48 @@ This project is a web-based application that classifies waste materials from use
 
 ## 🚀 Live Demo
 
-Try the application live on Hugging Face Spaces!
+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.
+  * **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 (E2B) 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-challenge)**
+
+-----
 
 ## 🛠️ Getting Started
 
 ### Prerequisites
 
-* Python 3.9+
-* Pip
-* Cuda (optional)
+  * Python 3.9+
+  * Pip
+  * Cuda (optional)
 
 ### Installation
 
 1.  Clone the repository:
     ```bash
-    git clone [https://github.com/yichuan-huang/gemma3n-challenge](https://github.com/yichuan-huang/gemma3n-challenge)
+    git clone https://github.com/yichuan-huang/gemma3n-challenge
     ```
 2.  Navigate to the project directory:
     ```bash
@@ -60,9 +68,9 @@ This will launch a Gradio web server. You can access the application by opening
 
 ## 📂 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.
+  * `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

@@ -94,7 +94,7 @@ class GarbageClassifier:
             image: PIL Image or path to image file
 
         Returns:
-            Tuple of (classification_result, full_response)
+            Tuple of (classification_result, detailed_analysis)
         """
         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 the garbage in this image and explain your reasoning.",
+                            "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.",
                         },
                     ],
                 },
@@ -155,10 +155,10 @@ class GarbageClassifier:
             # Extract classification from response
             classification = self._extract_classification(response)
 
-            # Create formatted response
-            formatted_response = self._format_response(classification, response)
+            # Extract reasoning from response
+            reasoning = self._extract_reasoning(response)
 
-            return classification, formatted_response
+            return classification, reasoning
 
         except Exception as e:
             self.logger.error(f"Error during classification: {str(e)}")
@@ -169,81 +169,162 @@ class GarbageClassifier:
 
     def _extract_classification(self, response: str) -> str:
         """Extract the main classification from the response"""
-        categories = self.knowledge.get_categories()
-
-        # Convert response to lowercase for matching
         response_lower = response.lower()
+        
+        # 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]
+                
+                # 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
+        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"
+        
+        # Look for food waste indicators
+        food_indicators = [
+            "food", "fruit", "vegetable", "organic", "kitchen waste",
+            "peel", "core", "scraps", "leftovers"
+        ]
+        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",
+            "cannot classify",
+            "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"
 
-        # Look for exact category matches first
+    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 category.lower() in response_lower:
-                return category
-
-        # Look for key terms if no exact match
-        category_keywords = {
-            "Recyclable Waste": [
-                "recyclable",
-                "recycle",
-                "plastic",
-                "paper",
-                "metal",
-                "glass",
-                "bottle",
-                "can",
-                "aluminum",
-                "cardboard",
-            ],
-            "Food/Kitchen Waste": [
-                "food",
-                "kitchen",
-                "organic",
-                "fruit",
-                "vegetable",
-                "leftovers",
-                "scraps",
-                "peel",
-                "core",
-                "bone",
-            ],
-            "Hazardous Waste": [
-                "hazardous",
-                "dangerous",
-                "toxic",
-                "battery",
-                "chemical",
-                "medicine",
-                "paint",
-                "pharmaceutical",
-            ],
-            "Other Waste": [
-                "other",
-                "general",
-                "trash",
-                "garbage",
-                "waste",
-                "cigarette",
-                "ceramic",
-                "dust",
-            ],
-        }
-
-        for category, keywords in category_keywords.items():
-            if any(keyword in response_lower for keyword in keywords):
-                return category
+            if cleaned_response.strip().startswith(category):
+                cleaned_response = cleaned_response.replace(category, "", 1)
+                break
 
-        return "Unable to classify"
+        # Split into sentences and clean up
+        sentences = []
+
+        # Split by common sentence endings
+        parts = re.split(r"[.!?]\s+", cleaned_response)
+
+        for part in parts:
+            part = part.strip()
+            if not part:
+                continue
+
+            # Skip parts that are just category names
+            if part in categories:
+                continue
+
+            # Skip parts that start with category names
+            is_category_line = False
+            for category in categories:
+                if part.startswith(category):
+                    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
 
-    def _format_response(self, classification: str, full_response: str) -> str:
-        """Format the response with classification and reasoning"""
-        if not full_response.strip():
-            return f"**Classification**: {classification}\n**Reasoning**: No detailed analysis available."
+            if part and len(part) > 3:  # Only keep meaningful content
+                sentences.append(part)
 
-        # If response already contains structured format, return as is
-        if "**Classification**" in full_response and "**Reasoning**" in full_response:
-            return full_response
+        # Join sentences and ensure proper punctuation
+        reasoning = ". ".join(sentences)
+        if reasoning and not reasoning.endswith((".", "!", "?")):
+            reasoning += "."
 
-        # Otherwise, format it
-        return f"**Classification**: {classification}\n\n**Reasoning**: {full_response}"
+        return reasoning if reasoning else "Analysis not available"
 
     def get_categories_info(self):
         """Get information about all categories"""

knowledge_base.py

@@ -3,6 +3,8 @@ class GarbageClassificationKnowledge:
     def get_system_prompt():
         return """You are a professional garbage classification expert. You need to carefully observe the items in the picture, analyze their materials, properties and uses, and then make accurate judgments according to garbage classification standards.
 
+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.
+
 Garbage classification standards:
 
 **Recyclable Waste**:
@@ -31,10 +33,19 @@ Garbage classification standards:
 - Large bones, hard shells, hard fruit pits (coconut shells, durian shells, walnut shells, corn cobs, etc.)
 - Hair, pet waste, cat litter, etc.
 
-Please observe the items in the image carefully according to the above classification standards, provide accurate garbage classification results, and briefly explain the classification reasoning. Format your response as:
+**Unable to classify**:
+- People, human faces, human body parts
+- Living animals, pets
+- Furniture, appliances, electronics in normal use
+- Buildings, landscapes, vehicles
+- Any item that is not intended to be discarded as waste
+
+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:
 
-**Classification**: [Category Name]
-**Reasoning**: [Brief explanation of why this item belongs to this category]"""
+**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]"""
 
     @staticmethod
     def get_categories():
@@ -43,6 +54,7 @@ Please observe the items in the image carefully according to the above classific
             "Food/Kitchen Waste",
             "Hazardous Waste",
             "Other Waste",
+            "Unable to classify",
         ]
 
     @staticmethod
@@ -52,4 +64,5 @@ Please observe the items in the image carefully according to the above classific
             "Food/Kitchen Waste": "Organic waste from food preparation and consumption",
             "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",
         }