app.py

app.py

@@ -1,63 +1,99 @@
 import cv2
 import numpy as np
 from PIL import Image
-import gradio as gr
-from transformers import CLIPProcessor, CLIPModel
 import torch
+import gradio as gr
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+# ---- Step 1: 模型加载 ----
+MODEL_ID = "prithivMLmods/Trash-Net"
+processor = AutoImageProcessor.from_pretrained(MODEL_ID)
+model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
+
+# ---- Step 2: 专业裁剪（基于OpenCV结构化差异检测） ----
+def smart_crop(base_img: Image.Image, new_img: Image.Image):
+    if base_img is None or new_img is None:
+        return None, "Missing image input."
+
+    # 转为OpenCV格式
+    base = np.array(base_img.convert("RGB"))
+    new = np.array(new_img.convert("RGB"))
+    base = cv2.resize(base, (224, 224))
+    new = cv2.resize(new, (224, 224))
 
-# Load CLIP model for material classification
-MODEL_ID = "openai/clip-vit-base-patch32"
-model = CLIPModel.from_pretrained(MODEL_ID)
-processor = CLIPProcessor.from_pretrained(MODEL_ID)
-LABELS = ["plastic", "metal", "paper", "cardboard", "glass", "trash"]
+    # 转灰度 + 高斯模糊减少噪声
+    base_gray = cv2.GaussianBlur(cv2.cvtColor(base, cv2.COLOR_RGB2GRAY), (5,5), 0)
+    new_gray = cv2.GaussianBlur(cv2.cvtColor(new, cv2.COLOR_RGB2GRAY), (5,5), 0)
+
+    # 差异检测
+    diff = cv2.absdiff(base_gray, new_gray)
+    _, thresh = cv2.threshold(diff, 40, 255, cv2.THRESH_BINARY)
+
+    # 形态学操作：去噪 + 连通区域扩大
+    kernel = np.ones((5,5), np.uint8)
+    thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
+    thresh = cv2.dilate(thresh, kernel, iterations=2)
+
+    # 查找轮廓
+    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return None, "No significant object difference detected."
 
-def get_clip_prediction(crop_img):
-    inputs = processor(text=LABELS, images=crop_img, return_tensors="pt", padding=True)
+    # 找最大差异区域
+    contour = max(contours, key=cv2.contourArea)
+    x, y, w, h = cv2.boundingRect(contour)
+
+    # 裁剪区域并略微扩大边缘
+    margin = 10
+    x1 = max(0, x - margin)
+    y1 = max(0, y - margin)
+    x2 = min(new.shape[1], x + w + margin)
+    y2 = min(new.shape[0], y + h + margin)
+
+    cropped = new[y1:y2, x1:x2]
+    cropped_pil = Image.fromarray(cropped)
+    return cropped_pil, None
+
+# ---- Step 3: TrashNet 分类 ----
+def classify_trash(image: Image.Image):
+    if image is None:
+        return "No image to classify.", 0.0
+    inputs = processor(images=image, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
-        logits_per_image = outputs.logits_per_image
-        probs = logits_per_image.softmax(dim=1).cpu().numpy()
-    best_idx = np.argmax(probs)
-    return LABELS[best_idx], float(probs[0][best_idx])
+        preds = torch.nn.functional.softmax(outputs.logits, dim=-1)
+        label = model.config.id2label[preds.argmax().item()]
+        confidence = preds.max().item()
+    return label, confidence
 
-def detect_diff_and_classify(base_img, target_img):
-    if base_img is None or target_img is None:
-        return "Please upload both images.", None
+# ---- Step 4: 主逻辑 ----
+def detect_and_classify(base_img, new_img):
+    if base_img is None or new_img is None:
+        return "⚠️ Please upload both base and new images.", None, None
 
-    base_np = np.array(base_img.convert("RGB"))
-    target_np = np.array(target_img.convert("RGB"))
+    cropped, error = smart_crop(base_img, new_img)
+    if error:
+        return f"⚠️ {error}", None, None
 
-    diff = cv2.absdiff(base_np, target_np)
-    gray = cv2.cvtColor(diff, cv2.COLOR_RGB2GRAY)
-    _, thresh = cv2.threshold(gray, 30, 255, cv2.THRESH_BINARY)
-    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    label, conf = classify_trash(cropped)
+    return f"✅ Object classified as: {label} ({conf*100:.2f}% confidence)", cropped, label
 
-    if not contours:
-        return "No significant difference detected.", None
-
-    # Largest contour → likely new object
-    c = max(contours, key=cv2.contourArea)
-    x, y, w, h = cv2.boundingRect(c)
-    crop = target_np[y:y+h, x:x+w]
-    crop_img = Image.fromarray(crop)
-
-    # Run CLIP classification
-    label, prob = get_clip_prediction(crop_img)
-    return f"Detected material: {label} (confidence {prob:.2f})", crop_img
-
-demo = gr.Interface(
-    fn=detect_diff_and_classify,
-    inputs=[
-        gr.Image(type="pil", label="Base Image"),
-        gr.Image(type="pil", label="Target Image")
-    ],
-    outputs=[
-        gr.Textbox(label="Result"),
-        gr.Image(type="pil", label="Detected Object Region")
-    ],
-    title="Automatic Object Detection + Material Classification",
-    description="Detect differences between base and target images, crop the changed region, and classify the material (plastic, metal, paper, cardboard, glass, trash)."
-)
+# ---- Step 5: Gradio 界面 ----
+with gr.Blocks(title="Smart Trash Detector") as demo:
+    gr.Markdown("# ♻️ Smart Trash Detection (with professional background subtraction)\nUpload a **base image**, then a **new image**.")
+
+    with gr.Row():
+        base_img = gr.Image(label="Base Image", type="pil")
+        new_img = gr.Image(label="New Image", type="pil")
+
+    run_btn = gr.Button("Analyze and Classify")
+
+    with gr.Row():
+        result_text = gr.Textbox(label="Result")
+        result_crop = gr.Image(label="Detected Object (Cropped)")
+        result_label = gr.Textbox(label="Class")
+
+    run_btn.click(detect_and_classify, [base_img, new_img], [result_text, result_crop, result_label])
 
 if __name__ == "__main__":
     demo.launch()