""" 垃圾图像分类模型 - 基于 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']})")