app.py

app.py

@@ -1,61 +1,54 @@
 import gradio as gr
 from transformers import BlipProcessor, BlipForQuestionAnswering
 from PIL import Image
-import os
-import difflib
 
-# Model
+# Load BLIP model
 model_name = "Salesforce/blip-vqa-base"
 processor = BlipProcessor.from_pretrained(model_name)
 model = BlipForQuestionAnswering.from_pretrained(model_name)
 
-# Categories
-CATEGORIES = ["plastic", "metal", "paper", "cardboard", "glass", "trash"]
-BASE_IMAGE_PATH = "base_reference.jpg"
-
-# --- Functions ---
+# Global variable to store base image
+base_image = None
 
 def set_base(image):
-    if image is None:
-        return "⚠️ Please upload an image."
-    image.save(BASE_IMAGE_PATH)
+    global base_image
+    base_image = image
     return "✅ Base image saved successfully."
 
-def identify_material(image):
+def detect_trash(image):
     if image is None:
-        return "⚠️ Please upload an image."
+        return "Please upload an image."
+
+    if base_image is None:
+        return "Please set base image first."
+
+    # Ask question
+    question = "Ignore the base, what material is this? Choose from: plastic, metal, paper, cardboard, glass, trash."
 
-    question = "Ignore the base, what material is this trash made of?"
     inputs = processor(image, question, return_tensors="pt")
     out = model.generate(**inputs)
-    answer = processor.decode(out[0], skip_special_tokens=True).lower()
+    answer = processor.decode(out[0], skip_special_tokens=True)
 
-    # Match to one of the six categories
-    match = difflib.get_close_matches(answer, CATEGORIES, n=1, cutoff=0.3)
-    if match:
-        return f"🧠 Detected: **{match[0]}**"
-    else:
-        return f"🤔 Unclear, model said: {answer}"
+    # Keep only defined categories
+    classes = ["plastic", "metal", "paper", "cardboard", "glass", "trash"]
+    matched = next((c for c in classes if c in answer.lower()), "trash")
 
-# --- UI ---
+    return matched.capitalize()
 
+# Two interfaces
 set_base_ui = gr.Interface(
     fn=set_base,
-    inputs=gr.Image(type="pil", label="Upload Empty Trash Bin Base"),
-    outputs="text",
-    title="🧱 Set Base",
-    description="Upload an image of the empty trash bin (no object)."
+    inputs=gr.Image(type="pil", label="Upload Empty Base Image"),
+    outputs=gr.Textbox(label="Result"),
+    title="🧩 Set Base"
 )
 
-identify_ui = gr.Interface(
-    fn=identify_material,
-    inputs=gr.Image(type="pil", label="Upload Trash Image to Identify"),
-    outputs="markdown",
-    title="🧠 Waste Material Classifier",
-    description="Upload a trash image. Model will predict one of: plastic, metal, paper, cardboard, glass, or trash."
+detect_trash_ui = gr.Interface(
+    fn=detect_trash,
+    inputs=gr.Image(type="pil", label="Upload Trash Image"),
+    outputs=gr.Textbox(label="Detected Material"),
+    title="♻️ Trash Material Detector"
 )
 
-demo = gr.TabbedInterface([set_base_ui, identify_ui], ["Set Base", "Detect Trash"])
-
-if __name__ == "__main__":
-    demo.launch()
+demo = gr.TabbedInterface([set_base_ui, detect_trash_ui], ["Set Base", "Detect Trash"])
+demo.launch()