Spaces:
Running
Running
| """ | |
| 垃圾图像分类模型 - 基于 DETR 主体检测 + CLIP 零样本图像分类 | |
| Garbage Classification Model - DETR object detection + CLIP zero-shot classification | |
| """ | |
| import os | |
| import random | |
| import logging | |
| import time | |
| from datetime import date as date_type | |
| from PIL import Image | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| from typing import Dict, List, Tuple, Optional | |
| from knowledge_base import query_dna, get_knowledge_base | |
| # Import transformers - will be lazy-loaded | |
| transformers_available = False | |
| try: | |
| from transformers import pipeline | |
| import torch | |
| transformers_available = True | |
| except ImportError: | |
| pass | |
| logger = logging.getLogger(__name__) | |
| class GarbageClassifier: | |
| """垃圾分类分类器 - 使用 DETR 先检测主体区域，再用 CLIP 零样本分类""" | |
| # COCO 类别中与垃圾无关的（人、动物等），检测到这些时忽略 | |
| SKIP_CATEGORIES = { | |
| 'person', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', | |
| 'elephant', 'bear', 'zebra', 'giraffe', 'umbrella', | |
| 'snowboard', 'skis', 'surfboard', 'frisbee', | |
| 'baseball bat', 'baseball glove', 'tennis racket', | |
| 'car', 'truck', 'bus', 'motorcycle', 'bicycle', | |
| 'train', 'airplane', 'boat', 'traffic light', | |
| 'fire hydrant', 'stop sign', 'parking meter', 'bench', | |
| } | |
| def __init__(self, model_name: str = "openai/clip-vit-base-patch32"): | |
| self.classifier = None | |
| self.detector = None # DETR 目标检测模型（延迟加载） | |
| self.ocr = None # TrOCR 文字识别模型（延迟加载） | |
| self.model_name = model_name | |
| self.device = self._detect_device() | |
| self._executor = ThreadPoolExecutor(max_workers=2) # 并行推理线程池 | |
| self.max_image_pixels = 640 # 限制图片最长边（640px 足够识别，速度翻倍） | |
| self.enable_ocr = os.getenv('GC_ENABLE_OCR', '').lower() in ('1', 'true', 'yes') # 默认关闭 OCR 加速 | |
| # 候选标签（英文 - CLIP 基于英文训练，使用详细描述提高准确率） | |
| # 使用更具体的描述，帮助 CLIP 区分相似物品 | |
| self.candidate_labels = [ | |
| # === 常见可回收垃圾 === | |
| "plastic bottle for drinking", # 塑料瓶 | |
| "glass bottle for drinking", # 玻璃瓶 | |
| "aluminum can for soda or beer", # 易拉罐 | |
| "paper and cardboard for recycling", # 纸张纸板 | |
| "cardboard box for packaging", # 纸箱 | |
| "newspaper and magazines", # 报纸杂志 | |
| "clothing and textiles", # 衣物织物 | |
| "plastic bag for shopping", # 塑料袋 | |
| "metal tools and hardware", # 金属制品 | |
| "electronics and wires with circuit", # 电子产品 | |
| "wooden furniture or wood pieces", # 木材家具 | |
| "books and notebooks for school", # 书本 | |
| "glass jar with lid", # 玻璃罐 | |
| "iron and steel food cans", # 铁罐 | |
| "shoes and bags made of fabric", # 鞋包 | |
| "plastic container for food storage", # 塑料容器 | |
| "milk carton or tetra pak drink box", # 牛奶盒 | |
| "clean paper coffee cup", # 干净纸杯 | |
| "plastic bottle cap", # 瓶盖 | |
| # === 常见有害垃圾 === | |
| "battery cell for electronics", # 电池 | |
| "light bulb glass transparent", # 灯泡 | |
| "medicine pills and tablets", # 药品 | |
| "paint can with metal handle", # 油漆桶 | |
| "mercury thermometer glass", # 温度计 | |
| "pesticide spray bottle", # 杀虫剂 | |
| "fluorescent tube light", # 荧光灯管 | |
| "nail polish bottle small", # 化妆品 | |
| # === 常见厨余垃圾 === | |
| "food waste and leftovers on plate", # 食物残渣 | |
| "fruit and fruit peels organic waste", # 水果果皮 | |
| "banana peel yellow", # 香蕉皮 | |
| "apple core with seeds", # 苹果核 | |
| "vegetable leaves and green scraps", # 菜叶菜梗 | |
| "tea leaves and coffee grounds wet", # 茶叶咖啡渣 | |
| "egg shell broken pieces", # 蛋壳 | |
| "fish bones and chicken bones", # 鱼骨鸡骨 | |
| "nut shells hard brown", # 坚果壳 | |
| "leftover rice and noodles in bowl", # 剩饭剩面 | |
| "flowers and plants wilted", # 花草植物 | |
| "cooked meat and bones leftovers", # 骨头肉渣 | |
| "shrimp crab seafood shells", # 海鲜壳 | |
| # === 常见其他垃圾 === | |
| "styrofoam and foam packaging white", # 泡沫塑料 | |
| "disposable face mask blue", # 口罩 | |
| "ceramics and pottery plate or bowl", # 陶瓷 | |
| "tissue paper and napkins used", # 纸巾 | |
| "cigarette butt with filter", # 烟蒂 | |
| "diaper and sanitary pads", # 尿不湿 | |
| "dust and dirt pile", # 尘土 | |
| "broken glass not recyclable shards", # 碎玻璃 | |
| # === 中国常见特殊垃圾（详细描述以区分） === | |
| "instant noodle cup styrofoam cup with lid", # 泡面桶 | |
| "disposable wooden chopsticks pair", # 一次性筷子 | |
| "takeout food container plastic box", # 外卖盒 | |
| "bubble tea cup with plastic lid and straw", # 奶茶杯 | |
| "disposable bamboo chopsticks pair", # 竹筷 | |
| "wooden chopsticks reusable", # 木筷 | |
| "paper bowl disposable", # 纸碗 | |
| "wooden toothpick small", # 牙签 | |
| "plastic drinking straw", # 吸管 | |
| "wet wipes in package", # 湿巾 | |
| "cotton swab for ears", # 棉签 | |
| "bandage and adhesive tape medical", # 创可贴 | |
| "broken porcelain ceramic shards", # 碎瓷器 | |
| "disposable foam lunch box", # 泡沫饭盒 | |
| "plastic wrap and cling film", # 保鲜膜 | |
| ] | |
| # 标签 → 中国垃圾分类四分类映射 | |
| self.label_to_category: Dict[str, str] = { | |
| # ♻️ 可回收垃圾 (Recyclable) | |
| "plastic bottle for drinking": "可回收垃圾", | |
| "glass bottle for drinking": "可回收垃圾", | |
| "aluminum can for soda or beer": "可回收垃圾", | |
| "paper and cardboard for recycling": "可回收垃圾", | |
| "cardboard box for packaging": "可回收垃圾", | |
| "newspaper and magazines": "可回收垃圾", | |
| "clothing and textiles": "可回收垃圾", | |
| "plastic bag for shopping": "可回收垃圾", | |
| "metal tools and hardware": "可回收垃圾", | |
| "electronics and wires with circuit": "可回收垃圾", | |
| "wooden furniture or wood pieces": "可回收垃圾", | |
| "books and notebooks for school": "可回收垃圾", | |
| "glass jar with lid": "可回收垃圾", | |
| "iron and steel food cans": "可回收垃圾", | |
| "shoes and bags made of fabric": "可回收垃圾", | |
| "plastic container for food storage": "可回收垃圾", | |
| "milk carton or tetra pak drink box": "可回收垃圾", | |
| "clean paper coffee cup": "可回收垃圾", | |
| "plastic bottle cap": "可回收垃圾", | |
| # ☣️ 有害垃圾 (Hazardous) | |
| "battery cell for electronics": "有害垃圾", | |
| "light bulb glass transparent": "有害垃圾", | |
| "medicine pills and tablets": "有害垃圾", | |
| "paint can with metal handle": "有害垃圾", | |
| "mercury thermometer glass": "有害垃圾", | |
| "pesticide spray bottle": "有害垃圾", | |
| "fluorescent tube light": "有害垃圾", | |
| "nail polish bottle small": "有害垃圾", | |
| # 🍚 厨余垃圾 (Kitchen Waste) | |
| "food waste and leftovers on plate": "厨余垃圾", | |
| "fruit and fruit peels organic waste": "厨余垃圾", | |
| "banana peel yellow": "厨余垃圾", | |
| "apple core with seeds": "厨余垃圾", | |
| "vegetable leaves and green scraps": "厨余垃圾", | |
| "tea leaves and coffee grounds wet": "厨余垃圾", | |
| "egg shell broken pieces": "厨余垃圾", | |
| "fish bones and chicken bones": "厨余垃圾", | |
| "nut shells hard brown": "厨余垃圾", | |
| "leftover rice and noodles in bowl": "厨余垃圾", | |
| "flowers and plants wilted": "厨余垃圾", | |
| "cooked meat and bones leftovers": "厨余垃圾", | |
| "shrimp crab seafood shells": "厨余垃圾", | |
| # 🗑️ 其他垃圾 (Other/Residual) | |
| "styrofoam and foam packaging white": "其他垃圾", | |
| "disposable face mask blue": "其他垃圾", | |
| "ceramics and pottery plate or bowl": "其他垃圾", | |
| "tissue paper and napkins used": "其他垃圾", | |
| "cigarette butt with filter": "其他垃圾", | |
| "diaper and sanitary pads": "其他垃圾", | |
| "dust and dirt pile": "其他垃圾", | |
| "broken glass not recyclable shards": "其他垃圾", | |
| "instant noodle cup styrofoam cup with lid": "其他垃圾", | |
| "takeout food container plastic box": "其他垃圾", | |
| "bubble tea cup with plastic lid and straw": "其他垃圾", | |
| "disposable wooden chopsticks pair": "可回收垃圾", | |
| "disposable bamboo chopsticks pair": "可回收垃圾", | |
| "wooden chopsticks reusable": "可回收垃圾", | |
| "paper bowl disposable": "可回收垃圾", | |
| "wooden toothpick small": "其他垃圾", | |
| "plastic drinking straw": "其他垃圾", | |
| "wet wipes in package": "其他垃圾", | |
| "cotton swab for ears": "其他垃圾", | |
| "bandage and adhesive tape medical": "其他垃圾", | |
| "broken porcelain ceramic shards": "其他垃圾", | |
| "disposable foam lunch box": "其他垃圾", | |
| "plastic wrap and cling film": "其他垃圾", | |
| } | |
| # === 通用物品识别标签 === | |
| # 用于"识物"功能 - 识别照片中的是什么物品（不限于垃圾） | |
| self.general_object_labels = [ | |
| # 饮具 | |
| "water bottle", "glass cup", "coffee mug", "tea cup", | |
| "aluminum can", "thermos", "wine glass", "beer bottle", | |
| "plastic cup", "paper cup", "stainless steel bottle", | |
| # 餐具 | |
| "plate", "bowl", "chopsticks", "fork", "spoon", "knife", | |
| "ceramic bowl", "takeout box", "food container", | |
| # 食物 | |
| "apple", "banana", "orange", "bread", "cake", "cookie", | |
| "egg", "rice", "noodles", "sandwich", "pizza", | |
| "chocolate bar", "candy wrapper", "fruit", "vegetable", | |
| "watermelon", "grape", "strawberry", "corn", | |
| "potato chip bag", "instant noodle", "soup", | |
| # 电子产品 | |
| "cell phone", "laptop", "tablet", "headphones", | |
| "charging cable", "remote control", "keyboard", | |
| "computer mouse", "smart watch", "digital camera", | |
| "power bank", "earphones", | |
| # 服饰 | |
| "t-shirt", "pants", "shoe", "hat", "jacket", "sock", | |
| "backpack", "handbag", "wallet", "belt", | |
| "glove", "scarf", "sunglasses", | |
| # 文具 | |
| "book", "notebook", "pen", "pencil", "scissors", | |
| "tape dispenser", "newspaper", "magazine", "paper sheet", | |
| "envelope", "cardboard box", | |
| # 日用品 | |
| "key", "umbrella", "eyeglasses", "toothbrush", | |
| "towel", "bar of soap", "shampoo bottle", | |
| "tissue box", "trash bag", "plastic shopping bag", | |
| "mirror", "clock", "pillow", "blanket", "cushion", | |
| # 包装 | |
| "glass jar", "metal can", "milk carton", | |
| "gift box", "plastic bottle", "cardboard box", | |
| # 工具 | |
| "hammer", "screwdriver", "wrench", "pliers", | |
| "measuring tape", "flashlight", | |
| # 运动 | |
| "soccer ball", "basketball", "tennis ball", | |
| "baseball", "frisbee", "jump rope", | |
| # 特殊物品 | |
| "disposable face mask", "cigarette butt", "battery", | |
| "light bulb", "medicine pill", "disposable diaper", | |
| "wet wipe", "drinking straw", "toothpick", | |
| "cotton swab", "bandage", "flower", "potted plant", | |
| "toy", "stuffed animal", | |
| # 个护 | |
| "lipstick", "makeup brush", "nail polish", | |
| "comb", "hair brush", "razor", | |
| # 清洁 | |
| "sponge", "cleaning cloth", "broom", "dustpan", | |
| "laundry detergent bottle", | |
| ] | |
| # 通用物品标签 → 中文名 | |
| self.general_label_zh: Dict[str, str] = { | |
| "water bottle": "水瓶", "glass cup": "玻璃杯", | |
| "coffee mug": "咖啡杯", "tea cup": "茶杯", | |
| "aluminum can": "易拉罐", "thermos": "保温杯", | |
| "wine glass": "酒杯", "beer bottle": "啤酒瓶", | |
| "plastic cup": "塑料杯", "paper cup": "纸杯", | |
| "stainless steel bottle": "不锈钢瓶", | |
| "plate": "盘子", "bowl": "碗", "chopsticks": "筷子", | |
| "fork": "叉子", "spoon": "勺子", "knife": "刀", | |
| "ceramic bowl": "陶瓷碗", "takeout box": "外卖盒", | |
| "food container": "食物容器", | |
| "apple": "苹果", "banana": "香蕉", "orange": "橙子", | |
| "bread": "面包", "cake": "蛋糕", "cookie": "饼干", | |
| "egg": "鸡蛋", "rice": "米饭", "noodles": "面条", | |
| "sandwich": "三明治", "pizza": "披萨", | |
| "chocolate bar": "巧克力", "candy wrapper": "糖果包装", | |
| "fruit": "水果", "vegetable": "蔬菜", | |
| "watermelon": "西瓜", "grape": "葡萄", | |
| "strawberry": "草莓", "corn": "玉米", | |
| "potato chip bag": "薯片袋", "instant noodle": "泡面", | |
| "soup": "汤", | |
| "cell phone": "手机", "laptop": "笔记本电脑", | |
| "tablet": "平板电脑", "headphones": "耳机", | |
| "charging cable": "充电线", "remote control": "遥控器", | |
| "keyboard": "键盘", "computer mouse": "鼠标", | |
| "smart watch": "智能手表", "digital camera": "相机", | |
| "power bank": "充电宝", "earphones": "耳塞", | |
| "t-shirt": "T恤", "pants": "裤子", "shoe": "鞋子", | |
| "hat": "帽子", "jacket": "外套", "sock": "袜子", | |
| "backpack": "背包", "handbag": "手提包", | |
| "wallet": "钱包", "belt": "腰带", | |
| "glove": "手套", "scarf": "围巾", "sunglasses": "墨镜", | |
| "book": "书本", "notebook": "笔记本", "pen": "笔", | |
| "pencil": "铅笔", "scissors": "剪刀", | |
| "tape dispenser": "胶带", "newspaper": "报纸", | |
| "magazine": "杂志", "paper sheet": "纸张", | |
| "envelope": "信封", "cardboard box": "纸箱", | |
| "key": "钥匙", "umbrella": "雨伞", | |
| "eyeglasses": "眼镜", "toothbrush": "牙刷", | |
| "towel": "毛巾", "bar of soap": "肥皂", | |
| "shampoo bottle": "洗发水瓶", "tissue box": "纸巾盒", | |
| "trash bag": "垃圾袋", "plastic shopping bag": "塑料袋", | |
| "mirror": "镜子", "clock": "时钟", | |
| "pillow": "枕头", "blanket": "毯子", "cushion": "靠垫", | |
| "glass jar": "玻璃罐", "metal can": "金属罐", | |
| "milk carton": "牛奶盒", "gift box": "礼盒", | |
| "plastic bottle": "塑料瓶", | |
| "hammer": "锤子", "screwdriver": "螺丝刀", | |
| "wrench": "扳手", "pliers": "钳子", | |
| "measuring tape": "卷尺", "flashlight": "手电筒", | |
| "soccer ball": "足球", "basketball": "篮球", | |
| "tennis ball": "网球", "baseball": "棒球", | |
| "frisbee": "飞盘", "jump rope": "跳绳", | |
| "disposable face mask": "口罩", | |
| "cigarette butt": "烟蒂", "battery": "电池", | |
| "light bulb": "灯泡", "medicine pill": "药片", | |
| "disposable diaper": "尿不湿", "wet wipe": "湿巾", | |
| "drinking straw": "吸管", "toothpick": "牙签", | |
| "cotton swab": "棉签", "bandage": "创可贴", | |
| "flower": "花", "potted plant": "盆栽", | |
| "toy": "玩具", "stuffed animal": "毛绒玩具", | |
| "lipstick": "口红", "makeup brush": "化妆刷", | |
| "nail polish": "指甲油", "comb": "梳子", | |
| "hair brush": "发刷", "razor": "剃须刀", | |
| "sponge": "海绵", "cleaning cloth": "抹布", | |
| "broom": "扫把", "dustpan": "簸箕", | |
| "laundry detergent bottle": "洗衣液瓶", | |
| } | |
| # DETR COCO 类别 → 中文名（用于展示检测结果） | |
| self.coco_label_zh: Dict[str, str] = { | |
| "person": "人", "bicycle": "自行车", "car": "汽车", | |
| "motorcycle": "摩托车", "airplane": "飞机", "bus": "公交车", | |
| "train": "火车", "truck": "卡车", "boat": "船", | |
| "traffic light": "红绿灯", "fire hydrant": "消防栓", | |
| "stop sign": "停车标志", "parking meter": "停车计时器", | |
| "bench": "长椅", "bird": "鸟", "cat": "猫", "dog": "狗", | |
| "horse": "马", "sheep": "羊", "cow": "牛", "elephant": "大象", | |
| "bear": "熊", "zebra": "斑马", "giraffe": "长颈鹿", | |
| "backpack": "背包", "umbrella": "雨伞", | |
| "handbag": "手提包", "tie": "领带", "suitcase": "行李箱", | |
| "frisbee": "飞盘", "skis": "滑雪板", "snowboard": "滑雪板", | |
| "sports ball": "球", "kite": "风筝", | |
| "baseball bat": "棒球棒", "baseball glove": "棒球手套", | |
| "skateboard": "滑板", "surfboard": "冲浪板", | |
| "tennis racket": "网球拍", "bottle": "瓶子", | |
| "wine glass": "酒杯", "cup": "杯子", "fork": "叉子", | |
| "knife": "刀", "spoon": "勺子", "bowl": "碗", | |
| "banana": "香蕉", "apple": "苹果", "sandwich": "三明治", | |
| "orange": "橙子", "broccoli": "西兰花", "carrot": "胡萝卜", | |
| "hot dog": "热狗", "pizza": "披萨", "donut": "甜甜圈", | |
| "cake": "蛋糕", "chair": "椅子", "couch": "沙发", | |
| "potted plant": "盆栽", "bed": "床", | |
| "dining table": "餐桌", "toilet": "马桶", | |
| "tv": "电视", "laptop": "笔记本电脑", "mouse": "鼠标", | |
| "remote": "遥控器", "keyboard": "键盘", | |
| "cell phone": "手机", "microwave": "微波炉", | |
| "oven": "烤箱", "toaster": "烤面包机", "sink": "水槽", | |
| "refrigerator": "冰箱", "book": "书", "clock": "时钟", | |
| "vase": "花瓶", "scissors": "剪刀", | |
| "teddy bear": "泰迪熊", "hair drier": "吹风机", | |
| "toothbrush": "牙刷", | |
| } | |
| # 分类详细信息 | |
| self.category_info: Dict[str, Dict] = { | |
| "可回收垃圾": { | |
| "bin_color": "blue", | |
| "bin_name": "蓝色垃圾桶", | |
| "icon": "♻️", | |
| "description": "适宜回收利用和资源化利用的废弃物", | |
| "examples": "废纸、塑料瓶、金属罐、玻璃、织物、电子产品", | |
| "tip": "请保持清洁干燥，压扁后投放" | |
| }, | |
| "有害垃圾": { | |
| "bin_color": "red", | |
| "bin_name": "红色垃圾桶", | |
| "icon": "☣️", | |
| "description": "对人体健康或自然环境造成直接或潜在危害的废弃物", | |
| "examples": "废电池、废灯泡、过期药品、废油漆、杀虫剂", | |
| "tip": "请小心轻放，避免破损泄漏" | |
| }, | |
| "厨余垃圾": { | |
| "bin_color": "green", | |
| "bin_name": "绿色垃圾桶", | |
| "icon": "🍚", | |
| "description": "食品加工和消费过程中产生的剩菜剩饭等废弃物", | |
| "examples": "剩饭剩菜、果皮果核、茶叶渣、蛋壳、骨头", | |
| "tip": "请沥干水分，不要混入牙签、纸巾等杂物" | |
| }, | |
| "其他垃圾": { | |
| "bin_color": "gray", | |
| "bin_name": "灰色/黑色垃圾桶", | |
| "icon": "🗑️", | |
| "description": "除可回收物、有害垃圾、厨余垃圾之外的废弃物", | |
| "examples": "卫生纸、陶瓷碎片、尘土、烟蒂、使用过的口罩", | |
| "tip": "尽量沥干水分，分类投放" | |
| } | |
| } | |
| def _detect_device(self) -> str: | |
| """检测可用设备""" | |
| if transformers_available: | |
| try: | |
| if torch.cuda.is_available(): | |
| return "cuda" | |
| elif hasattr(torch, 'mps') and torch.backends.mps.is_available(): | |
| return "mps" | |
| except Exception: | |
| pass | |
| return "cpu" | |
| def _resize_if_needed(self, image: Image.Image) -> Image.Image: | |
| """限制图片最长边，减少计算量""" | |
| w, h = image.size | |
| if max(w, h) > self.max_image_pixels: | |
| ratio = self.max_image_pixels / max(w, h) | |
| new_w, new_h = int(w * ratio), int(h * ratio) | |
| logger.info(f"[优化] 图片 {w}x{h} -> {new_w}x{new_h}") | |
| return image.resize((new_w, new_h), Image.LANCZOS) | |
| return image | |
| def _load_model(self): | |
| """延迟加载 CLIP 分类模型""" | |
| if self.classifier is not None: | |
| return | |
| if not transformers_available: | |
| raise RuntimeError( | |
| "transformers/torch 未安装。请运行: pip install transformers torch" | |
| ) | |
| logger.info(f"[CLIP] 正在加载模型 {self.model_name} 到 {self.device}...") | |
| self.classifier = pipeline( | |
| "zero-shot-image-classification", | |
| model=self.model_name, | |
| device=self.device if self.device != "cpu" else -1, | |
| ) | |
| logger.info("[CLIP] 模型加载完成！") | |
| def _load_detector(self): | |
| """延迟加载 DETR 目标检测模型""" | |
| if self.detector is not None: | |
| return | |
| if not transformers_available: | |
| raise RuntimeError( | |
| "transformers/torch 未安装。请运行: pip install transformers torch" | |
| ) | |
| logger.info("[DETR] 正在加载目标检测模型 facebook/detr-resnet-50...") | |
| self.detector = pipeline( | |
| "object-detection", | |
| model="facebook/detr-resnet-50", | |
| device=self.device if self.device != "cpu" else -1, | |
| ) | |
| logger.info("[DETR] 目标检测模型加载完成！") | |
| def _load_ocr(self): | |
| """延迟加载 TrOCR 文字识别模型""" | |
| if self.ocr is not None: | |
| return | |
| if not transformers_available: | |
| raise RuntimeError( | |
| "transformers/torch 未安装。请运行: pip install transformers torch" | |
| ) | |
| logger.info("[OCR] 正在加载文字识别模型 microsoft/trocr-small-printed...") | |
| try: | |
| from transformers import TrOCRProcessor, VisionEncoderDecoderModel | |
| self._ocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-small-printed") | |
| self.ocr = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-small-printed") | |
| self.ocr.to(self.device if self.device != "cpu" else "cpu") | |
| self.ocr.eval() | |
| logger.info("[OCR] 文字识别模型加载完成！") | |
| except Exception as e: | |
| logger.warning(f"[OCR] 模型加载失败，跳过文字识别: {e}") | |
| self.ocr = None | |
| def _extract_text(self, image: Image.Image) -> Optional[str]: | |
| """ | |
| 用 TrOCR 从图片中提取文字（品牌名、标签等） | |
| Args: | |
| image: PIL Image | |
| Returns: | |
| str or None: 检测到的文字，如果没有则返回 None | |
| """ | |
| try: | |
| self._load_ocr() | |
| if self.ocr is None: | |
| return None | |
| import torch | |
| pixel_values = self._ocr_processor(images=image, return_tensors="pt").pixel_values | |
| pixel_values = pixel_values.to(self.device if self.device != "cpu" else "cpu") | |
| with torch.no_grad(): | |
| generated_ids = self.ocr.generate(pixel_values, max_new_tokens=20) | |
| text = self._ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip() | |
| # 过滤无效结果：至少2个字符且包含字母数字 | |
| if len(text) >= 2 and any(c.isalnum() for c in text): | |
| # 限制长度，只取前30个字符 | |
| text = text[:30] | |
| logger.info(f"[OCR] 检测到文字: '{text}'") | |
| return text | |
| return None | |
| except Exception as e: | |
| logger.warning(f"[OCR] 文字识别异常: {e}") | |
| return None | |
| def _detect_and_crop(self, image: Image.Image) -> Tuple[Image.Image, Optional[Dict]]: | |
| """ | |
| 用 DETR 检测图片中的主体物体，裁切到物体区域。 | |
| 过滤掉无关类别（人、动物等），选择置信度最高的物体。 | |
| Args: | |
| image: PIL Image | |
| Returns: | |
| (cropped_image, detection_info) - 裁切后的图片和检测信息 | |
| 如果检测失败，返回原图和 None | |
| """ | |
| try: | |
| self._load_detector() | |
| # 保持原图尺寸，DETR 内部会处理缩放 | |
| detections = self.detector(image) | |
| if not detections: | |
| logger.info("[DETR] 未检测到任何物体，使用原图") | |
| return image, None | |
| logger.info(f"[DETR] 检测到 {len(detections)} 个物体") | |
| for d in detections: | |
| box = d['box'] | |
| bw = box['xmax'] - box['xmin'] | |
| bh = box['ymax'] - box['ymin'] | |
| logger.info( | |
| f" - {d['label']}: {d['score']:.3f} " | |
| f"[{bw}x{bh}]" | |
| ) | |
| # 过滤掉无关类别（人、动物、交通工具等） | |
| valid = [ | |
| d for d in detections | |
| if d['label'] not in self.SKIP_CATEGORIES | |
| ] | |
| # 如果没有有效检测，放宽限制（只排除 person） | |
| if not valid: | |
| valid = [ | |
| d for d in detections | |
| if d['label'] not in {'person'} | |
| ] | |
| # 仍然没有，就用全部检测结果 | |
| if not valid: | |
| valid = detections | |
| # 选出置信度最高的 | |
| best = max(valid, key=lambda x: x['score']) | |
| # 置信度太低则放弃裁切 | |
| if best['score'] < 0.3: | |
| logger.info(f"[DETR] 最高置信度 {best['score']:.3f} 过低，使用原图") | |
| return image, None | |
| box = best['box'] | |
| w, h = image.size | |
| # 计算边界框尺寸 | |
| bw = box['xmax'] - box['xmin'] | |
| bh = box['ymax'] - box['ymin'] | |
| # 加 25% 边距，保留上下文 | |
| margin_x = bw * 0.25 | |
| margin_y = bh * 0.25 | |
| # 最小边距 20px，防止物体太小 | |
| margin_x = max(margin_x, 20) | |
| margin_y = max(margin_y, 20) | |
| x1 = max(0, int(box['xmin'] - margin_x)) | |
| y1 = max(0, int(box['ymin'] - margin_y)) | |
| x2 = min(w, int(box['xmax'] + margin_x)) | |
| y2 = min(h, int(box['ymax'] + margin_y)) | |
| cropped = image.crop((x1, y1, x2, y2)) | |
| logger.info( | |
| f"[DETR] 主体 '{best['label']}' (conf={best['score']:.3f}) " | |
| f"裁切区域 [{x1},{y1} -> {x2},{y2}]" | |
| ) | |
| detection_info = { | |
| 'label': best['label'], | |
| 'score': best['score'], | |
| 'box': box, | |
| } | |
| return cropped, detection_info | |
| except Exception as e: | |
| logger.warning(f"[DETR] 目标检测异常，使用原图: {e}") | |
| return image, None | |
| def _identify_object(self, image: Image.Image, detection_info: Optional[Dict] = None, ocr_text: Optional[str] = None) -> Dict: | |
| """ | |
| 识别图片中的是什么物品（识物功能） | |
| 策略： | |
| 1. OCR 文字（由调用方传入，已在并行中完成） | |
| 2. DETR 检测结果作为物体名称（高置信度时） | |
| 3. CLIP 通用物品标签作为后备 | |
| 4. 组合文字 + 物品名 + 置信度生成自然描述 | |
| Args: | |
| image: 裁切后的图片 | |
| detection_info: DETR 检测结果 | |
| ocr_text: 预提取的 OCR 文字（由调用方在并行线程中完成） | |
| Returns: | |
| dict: {name, name_zh, description, confidence, source, ocr_text} | |
| """ | |
| obj_name = None | |
| obj_name_zh = None | |
| obj_confidence = 0.0 | |
| obj_source = None | |
| # 第一步: OCR 提取文字（由调用方并行完成，此处直接使用结果） | |
| # ocr_text 已在 classify() 中通过 ThreadPoolExecutor 异步完成 | |
| # 第二步: DETR 检测结果 | |
| if detection_info and detection_info.get('score', 0) >= 0.5: | |
| label_en = detection_info['label'] | |
| label_zh = self.coco_label_zh.get(label_en, label_en) | |
| if label_en not in self.SKIP_CATEGORIES: | |
| obj_name = label_en | |
| obj_name_zh = label_zh | |
| obj_confidence = detection_info['score'] | |
| obj_source = "detr" | |
| # 第三步: CLIP 通用物品识别（如果 DETR 没有给出好结果） | |
| if obj_source is None: | |
| try: | |
| self._load_model() | |
| clip_results = self.classifier( | |
| image, candidate_labels=self.general_object_labels | |
| ) | |
| top = clip_results[0] | |
| obj_name = top['label'] | |
| obj_name_zh = self.general_label_zh.get(obj_name, obj_name) | |
| obj_confidence = top['score'] | |
| obj_source = "clip" | |
| logger.info( | |
| f"[识物] CLIP 检测: {obj_name_zh} ({obj_name}) " | |
| f"置信度 {obj_confidence:.2%}" | |
| ) | |
| except Exception as e: | |
| logger.warning(f"[识物] 识别失败: {e}") | |
| obj_name_zh = "未知物品" | |
| # 第四步: 根据置信度和 OCR 文字生成自然描述 | |
| description = self._build_description(obj_name_zh, ocr_text, obj_confidence) | |
| return { | |
| "name": obj_name, | |
| "name_zh": obj_name_zh or "未知物品", | |
| "description": description, | |
| "confidence": round(obj_confidence, 4), | |
| "source": obj_source, | |
| "ocr_text": ocr_text, | |
| } | |
| def _build_description(self, obj_name_zh: str, ocr_text: Optional[str], confidence: float) -> str: | |
| """ | |
| 根据物品名、OCR 文字和置信度生成自然语言描述 | |
| Examples: | |
| - "FoYes" + "塑料杯" + 高置信度 → "一个 FoYes 的塑料杯" | |
| - 无文字 + "塑料杯" + 高置信度 → "塑料杯" | |
| - 无文字 + "杯子" + 中置信度 → "一个类似杯子的物品" | |
| - 无文字 + 低置信度 → "一个物品" | |
| """ | |
| if not obj_name_zh: | |
| if ocr_text and len(ocr_text) >= 2: | |
| return f"印有「{ocr_text}」的物品" | |
| return "一个物品" | |
| # 置信度足够高时，直接用物品名 | |
| if confidence >= 0.6: | |
| if ocr_text and len(ocr_text) >= 2: | |
| # 有文字：品牌名 + 物品名 | |
| return f"一个 {ocr_text} 的{obj_name_zh}" | |
| return obj_name_zh | |
| # 中等置信度：加不确定性措辞 | |
| if confidence >= 0.3: | |
| if ocr_text and len(ocr_text) >= 2: | |
| return f"印有「{ocr_text}」的{obj_name_zh}" | |
| return f"可能是{obj_name_zh}" | |
| # 低置信度：模糊描述 | |
| if ocr_text and len(ocr_text) >= 2: | |
| return f"印有「{ocr_text}」的物品（{obj_name_zh}）" | |
| return f"一个类似{obj_name_zh}的物品" | |
| def classify(self, image_path: str) -> Dict: | |
| """ | |
| 对传入的图片进行垃圾分类 + 物品识别 | |
| 流程（并行优化版）： | |
| 1. 限制图片尺寸以加速推理 | |
| 2. 用 DETR 检测主体物体并裁切 | |
| 3. 同时启动：OCR（后台线程）+ CLIP 垃圾分类（主线程） | |
| 4. DETR/CLIP 物品识别 | |
| 5. 映射到中国四分类标准 | |
| Args: | |
| image_path: 图片文件路径 | |
| Returns: | |
| dict: 包含分类结果和物品识别结果的字典 | |
| """ | |
| t_start = time.time() | |
| self._load_model() | |
| # 打开并限制图片大小 | |
| image = Image.open(image_path).convert("RGB") | |
| image = self._resize_if_needed(image) | |
| original_size = image.size | |
| logger.info(f"输入图片: {original_size}") | |
| # 第一步：DETR 主体检测 + 裁切（串行——需要图片本身） | |
| cropped_image, detection_info = self._detect_and_crop(image) | |
| logger.info(f"裁切后: {cropped_image.size}") | |
| # 第二步：并行启动 OCR（后台线程）和 CLIP 垃圾分类（主线程） | |
| ocr_text = None | |
| if self.enable_ocr: | |
| try: | |
| self._load_ocr() | |
| except Exception: | |
| pass | |
| ocr_future = self._executor.submit(self._extract_text, cropped_image) | |
| else: | |
| ocr_future = None | |
| # 主线程：CLIP 垃圾分类 | |
| results = self.classifier( | |
| cropped_image, candidate_labels=self.candidate_labels | |
| ) | |
| # 等待 OCR 完成（非阻塞：超时则跳过 OCR） | |
| if ocr_future is not None: | |
| try: | |
| ocr_text = ocr_future.result(timeout=15) | |
| except Exception: | |
| logger.warning("[OCR] 超时或失败，跳过文字识别") | |
| # 第三步：物品识别（使用并行获取的 OCR 文本，不再额外耗时） | |
| object_id = self._identify_object(cropped_image, detection_info, ocr_text) | |
| # 取 top-3 结果 | |
| top3 = results[:3] | |
| top_label = top3[0]["label"] | |
| top_score = top3[0]["score"] | |
| # 映射到垃圾类别 | |
| category = self.label_to_category.get(top_label) | |
| if category is None: | |
| category = "其他垃圾" | |
| info = self.category_info.get(category, self.category_info["其他垃圾"]) | |
| # 构建预测详情 | |
| predictions = [] | |
| for r in top3: | |
| cat = self.label_to_category.get(r["label"], "其他垃圾") | |
| predictions.append({ | |
| "label": r["label"], | |
| "label_zh": self._get_label_zh(r["label"]), | |
| "score": round(r["score"], 4), | |
| "category": cat | |
| }) | |
| # 查询垃圾 DNA 知识库 | |
| dna_info = query_dna(top_label, self._get_label_zh(top_label)) | |
| elapsed = time.time() - t_start | |
| logger.info(f"[分类] 总耗时 {elapsed:.1f}s") | |
| return { | |
| "success": True, | |
| "item": top_label, | |
| "item_zh": self._get_label_zh(top_label), | |
| "confidence": round(top_score, 4), | |
| "category": category, | |
| "category_info": info, | |
| "predictions": predictions, | |
| "object": object_id, # 物品识别结果 | |
| "dna": dna_info, # 垃圾 DNA 信息 | |
| "detection": { | |
| "object": detection_info['label'] if detection_info else None, | |
| "confidence": detection_info['score'] if detection_info else None, | |
| "cropped": detection_info is not None, | |
| } if detection_info else None, | |
| } | |
| def _get_label_zh(self, label_en: str) -> str: | |
| """英文标签转中文""" | |
| label_map = { | |
| # 可回收 | |
| "plastic bottle for drinking": "塑料瓶", | |
| "glass bottle for drinking": "玻璃瓶", | |
| "aluminum can for soda or beer": "易拉罐", | |
| "paper and cardboard for recycling": "纸张纸板", | |
| "cardboard box for packaging": "纸箱", | |
| "newspaper and magazines": "报纸杂志", | |
| "clothing and textiles": "衣物织物", | |
| "plastic bag for shopping": "塑料袋", | |
| "metal tools and hardware": "金属制品", | |
| "electronics and wires with circuit": "电子产品", | |
| "wooden furniture or wood pieces": "木材家具", | |
| "books and notebooks for school": "书本", | |
| "glass jar with lid": "玻璃罐", | |
| "iron and steel food cans": "铁罐", | |
| "shoes and bags made of fabric": "鞋包", | |
| "plastic container for food storage": "塑料容器", | |
| "milk carton or tetra pak drink box": "牛奶盒", | |
| "clean paper coffee cup": "干净纸杯", | |
| "plastic bottle cap": "瓶盖", | |
| # 有害 | |
| "battery cell for electronics": "电池", | |
| "light bulb glass transparent": "灯泡", | |
| "medicine pills and tablets": "药品", | |
| "paint can with metal handle": "油漆桶", | |
| "mercury thermometer glass": "温度计", | |
| "pesticide spray bottle": "杀虫剂瓶", | |
| "fluorescent tube light": "荧光灯管", | |
| "nail polish bottle small": "化妆品", | |
| # 厨余 | |
| "food waste and leftovers on plate": "食物残渣", | |
| "fruit and fruit peels organic waste": "水果果皮", | |
| "banana peel yellow": "香蕉皮", | |
| "apple core with seeds": "苹果核", | |
| "vegetable leaves and green scraps": "菜叶菜梗", | |
| "tea leaves and coffee grounds wet": "茶叶咖啡渣", | |
| "egg shell broken pieces": "蛋壳", | |
| "fish bones and chicken bones": "鱼骨鸡骨", | |
| "nut shells hard brown": "坚果壳", | |
| "leftover rice and noodles in bowl": "剩饭剩面", | |
| "flowers and plants wilted": "花草植物", | |
| "cooked meat and bones leftovers": "骨头肉渣", | |
| "shrimp crab seafood shells": "海鲜壳", | |
| # 其他 | |
| "styrofoam and foam packaging white": "泡沫塑料", | |
| "disposable face mask blue": "口罩", | |
| "ceramics and pottery plate or bowl": "陶瓷", | |
| "tissue paper and napkins used": "纸巾", | |
| "cigarette butt with filter": "烟蒂", | |
| "diaper and sanitary pads": "尿不湿/卫生巾", | |
| "dust and dirt pile": "尘土", | |
| "broken glass not recyclable shards": "碎玻璃", | |
| "instant noodle cup styrofoam cup with lid": "泡面桶", | |
| "takeout food container plastic box": "外卖盒", | |
| "bubble tea cup with plastic lid and straw": "奶茶杯", | |
| "disposable wooden chopsticks pair": "一次性筷子", | |
| "disposable bamboo chopsticks pair": "一次性竹筷", | |
| "wooden chopsticks reusable": "木筷", | |
| "paper bowl disposable": "纸碗", | |
| "wooden toothpick small": "牙签", | |
| "plastic drinking straw": "吸管", | |
| "wet wipes in package": "湿巾", | |
| "cotton swab for ears": "棉签", | |
| "bandage and adhesive tape medical": "创可贴", | |
| "broken porcelain ceramic shards": "碎瓷器", | |
| "disposable foam lunch box": "泡沫饭盒", | |
| "plastic wrap and cling film": "保鲜膜", | |
| } | |
| return label_map.get(label_en, label_en) | |
| def _classify_single_crop(self, cropped_image: Image.Image, | |
| detection_info: Optional[Dict] = None, | |
| ocr_text: Optional[str] = None) -> Dict: | |
| """ | |
| 对单一裁剪区域进行垃圾分类（CLIP）+ 物品识别 | |
| Args: | |
| cropped_image: 裁剪后的 PIL Image | |
| detection_info: DETR 检测信息 | |
| ocr_text: 预提取的 OCR 文字 | |
| Returns: | |
| dict: {label, label_zh, confidence, category, category_info, | |
| object, dna, bbox} | |
| """ | |
| # CLIP 垃圾分类 | |
| results = self.classifier( | |
| cropped_image, candidate_labels=self.candidate_labels | |
| ) | |
| # 物品识别 | |
| object_id = self._identify_object(cropped_image, detection_info, ocr_text) | |
| # Top-1 垃圾分类结果 | |
| top = results[0] | |
| top_label = top["label"] | |
| top_score = top["score"] | |
| category = self.label_to_category.get(top_label, "其他垃圾") | |
| info = self.category_info.get(category, self.category_info["其他垃圾"]) | |
| # 构造预测详情 | |
| predictions = [] | |
| for r in results[:3]: | |
| cat = self.label_to_category.get(r["label"], "其他垃圾") | |
| predictions.append({ | |
| "label": r["label"], | |
| "label_zh": self._get_label_zh(r["label"]), | |
| "score": round(r["score"], 4), | |
| "category": cat, | |
| }) | |
| # DNA 知识 | |
| dna_info = query_dna(top_label, self._get_label_zh(top_label)) | |
| bbox = detection_info.get('box') if detection_info else None | |
| out = { | |
| "label": top_label, | |
| "label_zh": self._get_label_zh(top_label), | |
| "confidence": round(top_score, 4), | |
| "category": category, | |
| "category_info": info, | |
| "predictions": predictions, | |
| "object": object_id, | |
| "dna": dna_info, | |
| } | |
| if bbox: | |
| out["bbox"] = bbox | |
| if detection_info and detection_info.get('detr_label'): | |
| out["detected_by"] = detection_info['detr_label'] | |
| out["detected_confidence"] = detection_info['detr_confidence'] | |
| return out | |
| def classify_multi(self, image_path: str, max_objects: int = 4) -> Dict: | |
| """ | |
| 多物体垃圾分类识别 - 同时检测并分类图片中的多个物体 | |
| 流程： | |
| 1. DETR 检测所有物体 | |
| 2. 过滤无关类别，按置信度排序取 top-N | |
| 3. 对每个物体裁剪并逐一 CLIP 分类 | |
| 4. 汇总为数组返回 | |
| Args: | |
| image_path: 图片文件路径 | |
| max_objects: 最多处理物体数（默认 4，过多会显著增加耗时） | |
| Returns: | |
| dict: {success, multi, total, objects: [...], overall} | |
| """ | |
| t_start = time.time() | |
| self._load_model() | |
| # 打开并限制图片大小 | |
| image = Image.open(image_path).convert("RGB") | |
| image = self._resize_if_needed(image) | |
| original_size = image.size | |
| logger.info(f"[多物体] 输入图片: {original_size}") | |
| # ---- 第一步：DETR 检测所有物体 ---- | |
| all_detections = self._detect_all_objects(image) | |
| if not all_detections: | |
| logger.info("[多物体] 未检测到有效物体，降级为单物体模式") | |
| # 降级：对整图做一次分类 | |
| single = self.classify(image_path) | |
| single["multi"] = False | |
| single["total_objects"] = 0 | |
| single["objects"] = [single] | |
| return single | |
| logger.info(f"[多物体] 检测到 {len(all_detections)} 个有效物体") | |
| # ---- 第二步：裁剪 + 并行 OCR（默认关闭） ---- | |
| objects_data = [] | |
| ocr_futures = {} | |
| if self.enable_ocr: | |
| try: | |
| self._load_ocr() | |
| except Exception as e: | |
| logger.warning(f"[多物体] OCR 加载失败，跳过 OCR: {e}") | |
| self.enable_ocr = False | |
| for i, det in enumerate(all_detections[:max_objects]): | |
| box = det['box'] | |
| w, h = image.size | |
| bw, bh = box['xmax'] - box['xmin'], box['ymax'] - box['ymin'] | |
| margin_x, margin_y = max(bw * 0.2, 15), max(bh * 0.2, 15) | |
| x1 = max(0, int(box['xmin'] - margin_x)) | |
| y1 = max(0, int(box['ymin'] - margin_y)) | |
| x2 = min(w, int(box['xmax'] + margin_x)) | |
| y2 = min(h, int(box['ymax'] + margin_y)) | |
| cropped = image.crop((x1, y1, x2, y2)) | |
| # 跳过过小的裁剪区域（CLIP 推理会异常） | |
| if x2 - x1 < 10 or y2 - y1 < 10: | |
| logger.warning(f"[多物体] 物体 {i} 裁剪区域过小 ({x2-x1}x{y2-y1})，跳过") | |
| continue | |
| # 对每个裁剪区域并行 OCR（仅在启用时） | |
| if self.enable_ocr: | |
| ocr_futures[i] = self._executor.submit(self._extract_text, cropped) | |
| detection_info = { | |
| 'label': det['label'], | |
| 'score': det['score'], | |
| 'box': box, | |
| 'detr_label': det['label'], | |
| 'detr_confidence': det['score'], | |
| } | |
| objects_data.append({ | |
| 'index': i, | |
| 'cropped': cropped, | |
| 'detection_info': detection_info, | |
| 'bbox': box, | |
| }) | |
| # ---- 第三步：逐一 CLIP 分类 ---- | |
| classified_objects = [] | |
| for obj in objects_data: | |
| i = obj['index'] | |
| # 等待该物体的 OCR 结果（最多 8s） | |
| ocr_text = None | |
| if i in ocr_futures: | |
| try: | |
| ocr_text = ocr_futures[i].result(timeout=8) | |
| except Exception: | |
| pass | |
| # CLIP 分类（带异常保护，单个失败不影响其它物体） | |
| try: | |
| result = self._classify_single_crop( | |
| obj['cropped'], obj['detection_info'], ocr_text | |
| ) | |
| result['index'] = i | |
| result['bbox'] = obj['bbox'] | |
| classified_objects.append(result) | |
| except Exception as e: | |
| logger.warning(f"[多物体] 物体 {i} 分类失败，跳过: {e}") | |
| continue | |
| # 如果所有物体都失败了，降级为单物体模式 | |
| if not classified_objects: | |
| logger.info("[多物体] 所有物体分类失败，降级为单物体模式") | |
| single = self.classify(image_path) | |
| single["multi"] = False | |
| single["total_objects"] = 0 | |
| single["objects"] = [single] | |
| return single | |
| # ---- 第四步：汇总 ---- | |
| # 按置信度排序，取最高置信度作为 overall 类别 | |
| classified_objects.sort(key=lambda x: x['confidence'], reverse=True) | |
| # ...但展示时按检测到的空间顺序（左到右）排列 | |
| display_objects = sorted(classified_objects, key=lambda x: ( | |
| (x.get('bbox') or {}).get('xmin', 0) + (x.get('bbox') or {}).get('xmax', 0) | |
| ) / 2) | |
| # 确定整体分类（统计各类别，取置信度加权最高的类别） | |
| category_weights = {} | |
| for obj in classified_objects: | |
| cat = obj['category'] | |
| category_weights[cat] = category_weights.get(cat, 0) + obj['confidence'] | |
| overall_category = max(category_weights, key=category_weights.get) | |
| elapsed = time.time() - t_start | |
| logger.info(f"[多物体] 总耗时 {elapsed:.1f}s, 分类 {len(classified_objects)} 个物体") | |
| return { | |
| "success": True, | |
| "multi": True, | |
| "total_objects": len(classified_objects), | |
| "objects": display_objects, | |
| "overall": overall_category, | |
| "overall_info": self.category_info.get( | |
| overall_category, self.category_info["其他垃圾"] | |
| ), | |
| } | |
| def _detect_all_objects(self, image: Image.Image) -> List[Dict]: | |
| """ | |
| 检测图片中所有与垃圾相关的物体 | |
| Returns: | |
| [{label, score, box}, ...] 按置信度降序排列 | |
| """ | |
| try: | |
| self._load_detector() | |
| detections = self.detector(image) | |
| if not detections: | |
| return [] | |
| # 过滤掉无关类别 | |
| valid = [ | |
| d for d in detections | |
| if d['label'] not in self.SKIP_CATEGORIES | |
| ] | |
| # 如果没有任何有效物体，放宽限制（只排除 person） | |
| if not valid: | |
| valid = [d for d in detections if d['label'] not in {'person'}] | |
| if not valid: | |
| valid = detections | |
| # 按置信度降序排列 | |
| valid.sort(key=lambda x: x['score'], reverse=True) | |
| # 过滤太低置信度的 | |
| valid = [d for d in valid if d['score'] >= 0.25] | |
| logger.info(f"[DETR] 检测到 {len(detections)} 个物体，" | |
| f"过滤后 {len(valid)} 个有效物体") | |
| for d in valid: | |
| box = d['box'] | |
| bw = box['xmax'] - box['xmin'] | |
| bh = box['ymax'] - box['ymin'] | |
| logger.info(f" - {d['label']}: {d['score']:.3f} [{bw}x{bh}]") | |
| return valid | |
| except Exception as e: | |
| logger.warning(f"[DETR] 目标检测异常: {e}") | |
| return [] | |
| def generate_daily_challenge(self, challenge_date: Optional[date_type] = None, count: int = 10) -> List[Dict]: | |
| """ | |
| 生成每日挑战题目，基于日期确定性随机选取物品 | |
| 策略： | |
| 1. 用日期字符串做种子，确保同日同题 | |
| 2. 从 label_to_category 的垃圾标签中选取 | |
| 3. 尽量覆盖 4 个类别（保证每个类别至少 1 题） | |
| Args: | |
| challenge_date: 挑战日期，默认为当天 | |
| count: 题目数量，默认 10 | |
| Returns: | |
| [{name_en, name_zh, category}, ...] | |
| """ | |
| if challenge_date is None: | |
| challenge_date = date_type.today() | |
| # 用日期做确定性种子 | |
| seed_str = challenge_date.isoformat() # "2026-06-15" | |
| seed = hash(seed_str) & 0xFFFFFFFF | |
| rng = random.Random(seed) | |
| # 收集所有有效标签（排除非垃圾类别） | |
| items_by_category: Dict[str, List[Dict]] = {} | |
| for label_en, cat in self.label_to_category.items(): | |
| label_zh = self._get_label_zh(label_en) | |
| if cat not in items_by_category: | |
| items_by_category[cat] = [] | |
| items_by_category[cat].append({ | |
| "name_en": label_en, | |
| "name_zh": label_zh, | |
| "category": cat, | |
| }) | |
| categories = list(items_by_category.keys()) | |
| challenge = [] | |
| # 第一步：从每个类别至少选 1 题（4 个类别） | |
| for cat in categories: | |
| items = items_by_category[cat] | |
| rng.shuffle(items) | |
| challenge.append(items[0]) | |
| # 第二步：从剩余池中补全到 count 题 | |
| remaining_pool = [] | |
| for cat in categories: | |
| items = items_by_category[cat] | |
| already_taken = {c["name_en"] for c in challenge} | |
| remaining_pool.extend([it for it in items if it["name_en"] not in already_taken]) | |
| rng.shuffle(remaining_pool) | |
| needed = count - len(challenge) | |
| challenge.extend(remaining_pool[:needed]) | |
| # 第三步：整体打乱顺序 | |
| rng.shuffle(challenge) | |
| return challenge | |
| # 命令行测试 | |
| if __name__ == "__main__": | |
| import sys | |
| logging.basicConfig(level=logging.INFO) | |
| if len(sys.argv) < 2: | |
| print("用法: python model.py <图片路径>") | |
| sys.exit(1) | |
| clf = GarbageClassifier() | |
| result = clf.classify(sys.argv[1]) | |
| print(f"\n{'='*50}") | |
| print(f"检测物品: {result['item_zh']}") | |
| print(f"置信度: {result['confidence']:.2%}") | |
| print(f"分类: {result['category']}") | |
| print(f"垃圾桶: {result['category_info']['bin_name']}") | |
| print(f"提示: {result['category_info']['tip']}") | |
| if result.get('detection') and result['detection'].get('cropped'): | |
| print(f"\n[DETR] 检测到主体: {result['detection']['object']} " | |
| f"(置信度 {result['detection']['confidence']:.2%})") | |
| print(f"\nTop-3 预测:") | |
| for p in result['predictions']: | |
| print(f" - {p['label_zh']}: {p['score']:.2%} ({p['category']})") | |