app.py

app.py

@@ -1,5 +1,62 @@
-# Use a pipeline as a high-level helper
-from transformers import pipeline
+# app.py
+import gradio as gr
+from transformers import AutoModelForImageClassification, AutoImageProcessor
+from PIL import Image
+import torch
 
-pipe = pipeline("image-classification", model="yangy50/garbage-classification")
-pipe("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/hub/parrots.png")
+# ---------------- 配置模型 ----------------
+MODEL_MAIN = "prithivMLmods/Trash-Net"            # 主模型
+MODEL_SECOND = "yangy50/garbage-classification"  # 验证模型
+
+# ---------------- 加载模型 ----------------
+print("🔹 正在加载主模型...")
+processor1 = AutoImageProcessor.from_pretrained(MODEL_MAIN)
+model1 = AutoModelForImageClassification.from_pretrained(MODEL_MAIN)
+model1.eval()
+print("✅ 主模型加载完成")
+
+print("🔹 正在加载验证模型...")
+processor2 = AutoImageProcessor.from_pretrained(MODEL_SECOND)
+model2 = AutoModelForImageClassification.from_pretrained(MODEL_SECOND)
+model2.eval()
+print("✅ 验证模型加载完成")
+
+# ---------------- 推理函数 ----------------
+def classify_image(image: Image.Image):
+    # 主模型预测
+    inputs1 = processor1(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs1 = model1(**inputs1)
+        pred1 = outputs1.logits.argmax(-1).item()
+    label1 = model1.config.id2label[pred1]
+
+    # 验证模型预测
+    inputs2 = processor2(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs2 = model2(**inputs2)
+        pred2 = outputs2.logits.argmax(-1).item()
+    label2 = model2.config.id2label[pred2]
+
+    # 双重验证
+    if label1 == label2:
+        result = f"✅ 双重验证一致，最终判定为：{label1.upper()}"
+    else:
+        result = f"⚠️ 双重验证不一致！\n主模型：{label1}\n验证模型：{label2}"
+
+    return label1, label2, result
+
+# ---------------- Gradio 界面 ----------------
+iface = gr.Interface(
+    fn=classify_image,
+    inputs=gr.Image(type="pil", label="上传图片"),
+    outputs=[
+        gr.Textbox(label="主模型结果"),
+        gr.Textbox(label="验证模型结果"),
+        gr.Textbox(label="双重验证判定")
+    ],
+    title="垃圾分类双模型检测",
+    description="使用 Trash-Net 和 Garbage-ViT 两个模型进行垃圾分类，结果一致才判定最终类别。"
+)
+
+if __name__ == "__main__":
+    iface.launch()