281 lines
10 KiB
Python
281 lines
10 KiB
Python
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()
|