app.py

app.py

@@ -1,28 +1,63 @@
 import gradio as gr
 from PIL import Image, ImageChops
+from transformers import BlipProcessor, BlipForQuestionAnswering
+import torch
 
-# === 加载本地 base 图 ===
-BASE_IMAGE_PATH = "base.jpg"
-base_img = Image.open(BASE_IMAGE_PATH).convert("RGB")
+# Load BLIP-VQA model
+model_name = "Salesforce/blip-vqa-base"
+processor = BlipProcessor.from_pretrained(model_name)
+model = BlipForQuestionAnswering.from_pretrained(model_name)
 
-def detect_difference(uploaded_img):
-    uploaded_img = uploaded_img.convert("RGB")
-    diff = ImageChops.difference(uploaded_img, base_img)
+valid_classes = ["plastic", "metal", "paper", "cardboard", "glass", "trash"]
+base_img = None  # Global variable to store base image
 
-    # 检查差异程度
-    bbox = diff.getbbox()
-    if bbox is None:
-        return "No difference detected (same as base)."
-    else:
-        return "Trash detected! Differences found."
+# Function to compute difference image
+def get_difference_image(base: Image.Image, trash: Image.Image) -> Image.Image:
+    diff = ImageChops.difference(base, trash).convert("RGB")
+    # Optional: enhance contrast to highlight difference
+    return diff
 
-demo = gr.Interface(
-    fn=detect_difference,
+# Set base image
+def set_base(image):
+    global base_img
+    base_img = image.convert("RGB")
+    return "Base image saved successfully."
+
+# Detect trash material
+def detect_material(trash_image):
+    global base_img
+    if base_img is None:
+        return "Please set base image first."
+
+    trash_image = trash_image.convert("RGB")
+    diff_image = get_difference_image(base_img, trash_image)
+
+    question = "What material is this object? Choose one of: plastic, metal, paper, cardboard, glass, trash."
+
+    inputs = processor(diff_image, question, return_tensors="pt")
+    out = model.generate(**inputs)
+    answer = processor.decode(out[0], skip_special_tokens=True).lower()
+
+    # Ensure answer is one of the valid classes
+    material = next((c for c in valid_classes if c in answer), "trash")
+    return material.capitalize()
+
+# Build Gradio UI
+set_base_ui = gr.Interface(
+    fn=set_base,
+    inputs=gr.Image(type="pil", label="Upload Base Image (Empty Bin)"),
+    outputs=gr.Textbox(label="Result"),
+    title="Set Base Image",
+    api_name="/set_base"
+)
+
+detect_ui = gr.Interface(
+    fn=detect_material,
     inputs=gr.Image(type="pil", label="Upload Trash Image"),
-    outputs=gr.Textbox(label="Detection Result"),
-    title="♻️ Trash Detector (Static Base)",
-    description="Base image (empty bin) is preloaded from base.jpg during build."
+    outputs=gr.Textbox(label="Detected Material"),
+    title="Trash Material Detector",
+    api_name="/detect_material"
 )
 
-if __name__ == "__main__":
-    demo.launch()
+demo = gr.TabbedInterface([set_base_ui, detect_ui], ["Set Base", "Detect Trash"])
+demo.launch()