import os import time import cv2 import numpy as np from ultralytics import YOLO from pathlib import Path from PIL import Image, ImageDraw, ImageFont from core.accel import resolve_yolo_weights # ==================== 【配置区域】 ==================== _ROOT = Path(__file__).resolve().parents[1] class Config: # 你的Plate检测模型路径（相对项目根目录）；若存在 best_openvino_model 会自动优先加载 MODEL_PATH = str(_ROOT / "best.pt") CONF_THRESHOLD = 0.90 IOU_THRESHOLD = 0.45 # 显式 imgsz，便于 OpenVINO / CPU 控制耗时；可用环境变量 DISH_YOLO_IMGSZ 覆盖 IMGSZ = int(os.environ.get("DISH_YOLO_IMGSZ", "640")) BOX_THICKNESS = 2 FONT_SIZE = 22 # 输出目录 OUTPUT_DIR = r"detect" Cutting = r"cutting" GENERATE_VISUALIZATION = True # 是否生成可视化图 # 只有 Plate 一个类别 LABEL_NAMES = {0: "Plate"} COLORS = {'Plate': (0, 255, 0)} # ======================================================== class YOLOv11PlateExtractor: def __init__(self, config): self.config = config weights = resolve_yolo_weights(config.MODEL_PATH) self.model = YOLO(weights) self.weights_path = weights print(f"✅ Plate 模型加载成功: {weights}") def predict(self, source): """source 可为图片路径或 BGR ndarray。""" return self.model.predict( source=source, conf=self.config.CONF_THRESHOLD, iou=self.config.IOU_THRESHOLD, imgsz=getattr(self.config, "IMGSZ", 640), verbose=False, )[0] def crop_plate(self, image, results): cropped_regions = [] if results.boxes is not None: boxes = results.boxes.xyxy.cpu().numpy() confs = results.boxes.conf.cpu().numpy() h, w = image.shape[:2] for box, conf in zip(boxes, confs): x1, y1, x2, y2 = map(int, box) x1 = max(0, x1) y1 = max(0, y1) x2 = min(w, x2) y2 = min(h, y2) if x2 > x1 and y2 > y1: crop = image[y1:y2, x1:x2] if crop.size > 0: cropped_regions.append({ "image": crop.copy(), "conf": float(conf), "box": [x1, y1, x2, y2], }) return cropped_regions # 标注默认：大库红框 + 白字；今日菜单由 style 覆盖为绿框绿字 LABEL_BOX_COLOR = (229, 57, 53) LABEL_CHIP_BG = (229, 57, 53, 220) LABEL_CHIP_TEXT = (255, 255, 255, 255) def _load_font(self, size=None): size = size or self.config.FONT_SIZE for font_name in ("simhei.ttf", "simsun.ttc", "msyh.ttc", "msyh.ttf"): try: return ImageFont.truetype(font_name, size) except Exception: continue return ImageFont.load_default() def _draw_box_and_chip(self, base_img, box, text, font, style=None): """描边框 + 标签芯片；style 可覆盖颜色（今日菜单绿 / 大库红 / 未知橙）。""" style = style or {} outline = style.get("box") or (229, 57, 53, 180) chip_bg = style.get("chip_bg") or self.LABEL_CHIP_BG chip_fg = style.get("chip_fg") or self.LABEL_CHIP_TEXT thickness = int(style.get("box_width") or self.config.BOX_THICKNESS or 2) x1, y1, x2, y2 = map(int, box) rgba = base_img.convert("RGBA") overlay = Image.new("RGBA", rgba.size, (0, 0, 0, 0)) odraw = ImageDraw.Draw(overlay) for t in range(max(1, thickness)): odraw.rectangle( [(x1 - t, y1 - t), (x2 + t, y2 + t)], outline=outline, ) try: bbox = font.getbbox(str(text)) tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1] except Exception: tw, th = len(str(text)) * 10, 16 pad_x, pad_y = 8, 4 chip_h = th + pad_y * 2 chip_w = max(tw + pad_x * 2, chip_h) chip_y = y1 - chip_h - 4 if y1 - chip_h - 4 > 0 else y1 + 4 chip_x = x1 # 圆形编号更贴近参考图 if str(text).isdigit() and len(str(text)) <= 2: r = max(chip_h, 22) // 2 + 2 cx, cy = x1 + r + 2, max(y1 - r - 4, r + 2) odraw.ellipse([(cx - r, cy - r), (cx + r, cy + r)], fill=chip_bg) odraw.text((cx, cy), str(text), fill=chip_fg, font=font, anchor="mm") else: odraw.rectangle( [(chip_x, chip_y), (chip_x + chip_w, chip_y + chip_h)], fill=chip_bg, ) odraw.text((chip_x + pad_x, chip_y + pad_y), str(text), fill=chip_fg, font=font) return Image.alpha_composite(rgba, overlay).convert("RGB") def visualize_results(self, image, results, save_path, labels=None): """ 在检测框上绘制标签。 labels 为 None 时画 Plate+置信度；传入列表时画菜品名称（认菜不认盘展示）。 """ img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) font = self._load_font() if results.boxes is not None: boxes = results.boxes.xyxy.cpu().numpy() confs = results.boxes.conf.cpu().numpy() for i, (box, conf) in enumerate(zip(boxes, confs)): if labels is not None and i < len(labels): text = str(labels[i]) else: text = f"Plate {conf:.2f}" img = self._draw_box_and_chip(img, box, text, font) os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True) img.save(save_path) return save_path def visualize_with_labels(self, image_or_path, boxes, labels, save_path, styles=None): """ 在原图上按框绘制菜品名称（不依赖 YOLO results）。 :param image_or_path: BGR ndarray 或图片路径 :param boxes: [[x1,y1,x2,y2], ...] :param labels: 与 boxes 对齐的菜品名称列表 :param styles: 可选，与 boxes 对齐的颜色样式 dict 列表 """ if isinstance(image_or_path, str): image = cv2.imread(image_or_path) if image is None: raise ValueError(f"无法读取图片: {image_or_path}") else: image = image_or_path img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) font = self._load_font() styles = styles or [] for i, (box, label) in enumerate(zip(boxes, labels)): style = styles[i] if i < len(styles) else None img = self._draw_box_and_chip(img, box, str(label), font, style=style) os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True) img.save(save_path) return save_path def _load_bgr(self, image_or_path): if isinstance(image_or_path, np.ndarray): if image_or_path.size == 0: return None return image_or_path if isinstance(image_or_path, (str, Path)): image = cv2.imread(str(image_or_path)) return image return None def process_image( self, image_or_path, stem=None, save_crops=True, visualize=None, ): """ 统一检测入口：支持路径或 BGR ndarray。 :return: (visualize_save_path, crop_items, boxes) crop_items: [{"image": ndarray, "path": str|None, "conf": float}, ...] """ image = self._load_bgr(image_or_path) if image is None: print(f"❌ 读取失败: {image_or_path}") return None, [], [] if stem is None: if isinstance(image_or_path, (str, Path)): stem = Path(image_or_path).stem else: stem = f"frame_{int(time.time() * 1000)}" # 内存直传，避免 YOLO 再读盘 results = self.predict(image) cropped_regions = self.crop_plate(image, results) do_viz = self.config.GENERATE_VISUALIZATION if visualize is None else bool(visualize) visualize_save_path = None if do_viz: os.makedirs(self.config.OUTPUT_DIR, exist_ok=True) visualize_save_path = self.visualize_results( image, results, os.path.join(self.config.OUTPUT_DIR, f"{stem}.jpg"), ) crop_items = [] boxes = [] if save_crops: os.makedirs(self.config.Cutting, exist_ok=True) for i, region in enumerate(cropped_regions): img_crop = region["image"] path = None if save_crops: filename = f"{stem}_Plate_{i + 1}.jpeg" path = os.path.join(self.config.Cutting, filename) cv2.imwrite(path, img_crop) crop_items.append({ "image": img_crop, "path": path, "conf": region.get("conf"), }) boxes.append(region["box"]) print(f"✅ {stem} → 裁剪完成，共 {len(cropped_regions)} 张") return visualize_save_path, crop_items, boxes def process_single_image(self, image_path): """ 处理单张图片：检测餐盘、裁剪、生成初始可视化图。 :return: (visualize_save_path, cropped_paths, boxes) boxes 与 cropped_paths 一一对应，供后续标注菜品名 """ viz, crop_items, boxes = self.process_image( image_path, save_crops=True, visualize=True, ) cropped_paths = [c["path"] for c in crop_items if c.get("path")] return viz, cropped_paths, boxes def main(): cfg = Config() app = YOLOv11PlateExtractor(cfg) img_path = r"C:\Users\王含波\Desktop\文件\classfiy\20260422095806.jpg" visualize_save_path, cropped_file_paths, boxes = app.process_single_image(img_path) print("\n📦 裁剪后的图片路径：") print(cropped_file_paths) print("boxes:", boxes) print("viz:", visualize_save_path) if __name__ == "__main__": main()