app.py

app.py

@@ -1,67 +1,93 @@
 import gradio as gr
-from PIL import Image, ImageChops
-from transformers import BlipProcessor, BlipForQuestionAnswering
+from PIL import Image
+import numpy as np
 import torch
+from transformers import (
+    AutoModelForImageSegmentation,
+    AutoProcessor,
+    AutoFeatureExtractor,
+    AutoModelForImageClassification,
+)
+
+# === ① Load SAM model for segmentation ===
+sam_model_id = "facebook/sam-vit-base"
+processor_sam = AutoProcessor.from_pretrained(sam_model_id)
+model_sam = AutoModelForImageSegmentation.from_pretrained(sam_model_id)
+
+# === ② Load garbage classification model ===
+cls_model_id = "yangy50/garbage-classification"
+extractor = AutoFeatureExtractor.from_pretrained(cls_model_id)
+cls_model = AutoModelForImageClassification.from_pretrained(cls_model_id)
+
+base_img = None  # Global memory for base image
+
+# === Step 1: Set base ===
+def set_base(image):
+    global base_img
+    if image is None:
+        return "Please upload an empty bin image."
+    base_img = image.convert("RGB")
+    return "✅ Base image saved successfully."
+
+# === Step 2: Detect and classify trash ===
+def detect_trash(image):
+    global base_img
+    if base_img is None:
+        return "Please set a base image first."
+
+    current_img = image.convert("RGB")
+
+    # Convert to numpy
+    base_np = np.array(base_img).astype(np.float32)
+    current_np = np.array(current_img).astype(np.float32)
+
+    # Difference mask
+    diff = np.abs(current_np - base_np).mean(axis=2)
+    mask = (diff > 40).astype(np.uint8) * 255  # threshold
+    mask_img = Image.fromarray(mask).convert("RGB")
+
+    # Use SAM to refine the mask
+    inputs = processor_sam(images=current_img, segmentation_maps=mask_img, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model_sam(**inputs)
+    seg = outputs.pred_masks[0].cpu().numpy()
+
+    # Crop bounding box of detected trash
+    ys, xs = np.where(seg > 0.5)
+    if len(xs) == 0 or len(ys) == 0:
+        return "No significant object detected."
+
+    x1, x2, y1, y2 = xs.min(), xs.max(), ys.min(), ys.max()
+    cropped = current_img.crop((x1, y1, x2, y2))
+
+    # Classify the cropped object
+    cls_inputs = extractor(images=cropped, return_tensors="pt")
+    with torch.no_grad():
+        cls_out = cls_model(**cls_inputs)
+        probs = torch.nn.functional.softmax(cls_out.logits, dim=-1)
+        pred_idx = torch.argmax(probs, dim=-1).item()
+        pred_class = cls_model.config.id2label[pred_idx]
+
+    return f"🧩 Detected Material: {pred_class}"
+
+# === Build UI ===
+set_base_ui = gr.Interface(
+    fn=set_base,
+    inputs=gr.Image(type="pil", label="Upload Empty Bin (Base)"),
+    outputs=gr.Textbox(label="Status"),
+    title="🧩 Set Base",
+)
+
+detect_trash_ui = gr.Interface(
+    fn=detect_trash,
+    inputs=gr.Image(type="pil", label="Upload Trash Image"),
+    outputs=gr.Textbox(label="Detection Result"),
+    title="♻️ Detect & Classify Trash",
+)
 
-# Load BLIP model
-model_name = "Salesforce/blip-vqa-base"
-processor = BlipProcessor.from_pretrained(model_name)
-model = BlipForQuestionAnswering.from_pretrained(model_name)
-
-# Ensure device
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
-
-def crop_difference(base_img: Image.Image, trash_img: Image.Image) -> Image.Image:
-    # Convert to same mode
-    base_img = base_img.convert("RGB")
-    trash_img = trash_img.convert("RGB")
-    
-    # Compute difference
-    diff = ImageChops.difference(trash_img, base_img)
-    # Crop to non-zero bbox
-    bbox = diff.getbbox()
-    if bbox:
-        cropped = trash_img.crop(bbox)
-        return cropped
-    else:
-        return trash_img  # fallback if no difference
-
-def identify_material(base_img: Image.Image, trash_img: Image.Image) -> str:
-    if base_img is None or trash_img is None:
-        return "Please upload both base and trash images."
-    
-    cropped = crop_difference(base_img, trash_img)
-    
-    question = "What material is this? Choose from: plastic, metal, paper, cardboard, glass, trash."
-    inputs = processor(cropped, question, return_tensors="pt").to(device)
-    out = model.generate(**inputs)
-    answer = processor.decode(out[0], skip_special_tokens=True)
-    
-    valid_classes = ["plastic", "metal", "paper", "cardboard", "glass", "trash"]
-    result = next((c for c in valid_classes if c in answer.lower()), "trash")
-    return result.capitalize()
-
-title = "Smart Waste Material Detector"
-description = """
-Upload two images:
-1. Base image (empty background)
-2. Trash image (object placed on background)
-
-The AI will detect the difference and classify the material:
-plastic, metal, paper, cardboard, glass, or trash.
-"""
-
-demo = gr.Interface(
-    fn=identify_material,
-    inputs=[
-        gr.Image(type="pil", label="Upload Base Image (Empty)"),
-        gr.Image(type="pil", label="Upload Trash Image")
-    ],
-    outputs=gr.Textbox(label="Detected Material"),
-    title=title,
-    description=description,
-    allow_flagging="never"
+demo = gr.TabbedInterface(
+    [set_base_ui, detect_trash_ui],
+    ["Set Base", "Detect Trash"]
 )
 
 if __name__ == "__main__":