403 lines
15 KiB
Python
403 lines
15 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 # 是否生成可视化图
|
||
|
||
# 默认类别名（加载权重后会被 model.names 覆盖）
|
||
LABEL_NAMES = {0: "Plate"}
|
||
COLORS = {'Plate': (0, 255, 0)}
|
||
# 检测到后直接丢弃，不进入裁剪/识别
|
||
FILTER_CLASS_NAMES = ("手", "空盘")
|
||
|
||
|
||
# ========================================================
|
||
|
||
def _load_filter_names(config):
|
||
names = getattr(config, "FILTER_CLASS_NAMES", None)
|
||
if names is None:
|
||
try:
|
||
from dao.settings import DISH_DETECT_FILTER_NAMES
|
||
names = DISH_DETECT_FILTER_NAMES
|
||
except Exception:
|
||
names = ("手", "空盘")
|
||
return {str(n).strip() for n in (names or []) if str(n).strip()}
|
||
|
||
|
||
def _load_edge_filter_params(config):
|
||
"""贴边/过小框过滤参数；可用 Config 覆盖，否则读 settings。"""
|
||
def _get(name, default):
|
||
if hasattr(config, name):
|
||
return getattr(config, name)
|
||
try:
|
||
from dao import settings as S
|
||
return getattr(S, name, default)
|
||
except Exception:
|
||
return default
|
||
|
||
return {
|
||
"edge_margin": int(_get("DISH_DETECT_EDGE_MARGIN_PX", 8)),
|
||
"edge_min_area_ratio": float(_get("DISH_DETECT_EDGE_MIN_AREA_RATIO", 0.02)),
|
||
"edge_min_side": int(_get("DISH_DETECT_EDGE_MIN_SIDE_PX", 48)),
|
||
"min_area_ratio": float(_get("DISH_DETECT_MIN_AREA_RATIO", 0.008)),
|
||
}
|
||
|
||
|
||
class YOLOv11PlateExtractor:
|
||
def __init__(self, config):
|
||
self.config = config
|
||
weights = resolve_yolo_weights(config.MODEL_PATH)
|
||
self.model = YOLO(weights)
|
||
self.weights_path = weights
|
||
# 以权重内类别为准
|
||
try:
|
||
model_names = getattr(self.model, "names", None) or {}
|
||
if isinstance(model_names, dict) and model_names:
|
||
self.config.LABEL_NAMES = {int(k): str(v) for k, v in model_names.items()}
|
||
except Exception:
|
||
pass
|
||
self.filter_class_names = _load_filter_names(config)
|
||
self.edge_filter = _load_edge_filter_params(config)
|
||
print(f"✅ Plate 模型加载成功: {weights}")
|
||
print(f"[OK] detect classes: {self.config.LABEL_NAMES}")
|
||
if self.filter_class_names:
|
||
print(f"[OK] detect filter drop: {sorted(self.filter_class_names)}")
|
||
ef = self.edge_filter
|
||
print(
|
||
f"[OK] edge/small box filter: margin={ef['edge_margin']}px "
|
||
f"edge_area<{ef['edge_min_area_ratio']:.3f} "
|
||
f"edge_side<{ef['edge_min_side']} "
|
||
f"min_area<{ef['min_area_ratio']:.3f}"
|
||
)
|
||
|
||
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 _class_name(self, results, cls_id):
|
||
cid = int(cls_id)
|
||
names = getattr(results, "names", None) or self.config.LABEL_NAMES or {}
|
||
if isinstance(names, dict):
|
||
return str(names.get(cid, names.get(str(cid), cid)))
|
||
return str(cid)
|
||
|
||
def _should_drop_class(self, class_name):
|
||
name = (class_name or "").strip()
|
||
if not name:
|
||
return False
|
||
return name in self.filter_class_names
|
||
|
||
def _should_drop_incomplete_box(self, box, img_w, img_h):
|
||
"""贴边且过小 / 面积过小 → 视为入画不完整，丢弃。"""
|
||
try:
|
||
x1, y1, x2, y2 = [float(v) for v in box[:4]]
|
||
except Exception:
|
||
return True
|
||
bw = max(0.0, x2 - x1)
|
||
bh = max(0.0, y2 - y1)
|
||
if bw <= 1 or bh <= 1 or img_w <= 0 or img_h <= 0:
|
||
return True
|
||
area_ratio = (bw * bh) / float(img_w * img_h)
|
||
ef = self.edge_filter or {}
|
||
min_area = float(ef.get("min_area_ratio") or 0)
|
||
if min_area > 0 and area_ratio < min_area:
|
||
return True
|
||
|
||
margin = max(0, int(ef.get("edge_margin") or 0))
|
||
touches = (
|
||
x1 <= margin
|
||
or y1 <= margin
|
||
or x2 >= (img_w - 1 - margin)
|
||
or y2 >= (img_h - 1 - margin)
|
||
)
|
||
if not touches:
|
||
return False
|
||
edge_area = float(ef.get("edge_min_area_ratio") or 0)
|
||
edge_side = int(ef.get("edge_min_side") or 0)
|
||
# 贴边且「面积小」或「短边过短」→ 多半是只露一角
|
||
if edge_area > 0 and area_ratio < edge_area:
|
||
return True
|
||
if edge_side > 0 and min(bw, bh) < edge_side:
|
||
return True
|
||
return False
|
||
|
||
def iter_kept_detections(self, results, image_shape=None):
|
||
"""产出保留的检测框：(box, conf, class_name)。类别/贴边不完整已过滤。"""
|
||
if results is None or results.boxes is None:
|
||
return
|
||
boxes = results.boxes.xyxy.cpu().numpy()
|
||
confs = results.boxes.conf.cpu().numpy()
|
||
if getattr(results.boxes, "cls", None) is not None:
|
||
clss = results.boxes.cls.cpu().numpy()
|
||
else:
|
||
clss = [0] * len(boxes)
|
||
img_h = img_w = None
|
||
if image_shape is not None and len(image_shape) >= 2:
|
||
img_h, img_w = int(image_shape[0]), int(image_shape[1])
|
||
dropped_cls = 0
|
||
dropped_edge = 0
|
||
for box, conf, cls_id in zip(boxes, confs, clss):
|
||
cname = self._class_name(results, cls_id)
|
||
if self._should_drop_class(cname):
|
||
dropped_cls += 1
|
||
continue
|
||
if img_w and img_h and self._should_drop_incomplete_box(box, img_w, img_h):
|
||
dropped_edge += 1
|
||
continue
|
||
yield box, float(conf), cname
|
||
if dropped_cls:
|
||
print(f"[detect] filtered {dropped_cls} box(es) by class {sorted(self.filter_class_names)}")
|
||
if dropped_edge:
|
||
print(f"[detect] filtered {dropped_edge} incomplete/edge box(es)")
|
||
|
||
def crop_plate(self, image, results):
|
||
cropped_regions = []
|
||
h, w = image.shape[:2]
|
||
for box, conf, cname in self.iter_kept_detections(results, image_shape=image.shape):
|
||
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],
|
||
"class_name": cname,
|
||
})
|
||
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 时画 类别+置信度；传入列表时画菜品名称（认菜不认盘展示）。
|
||
"""
|
||
img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||
font = self._load_font()
|
||
kept = list(self.iter_kept_detections(results, image_shape=image.shape))
|
||
for i, (box, conf, cname) in enumerate(kept):
|
||
if labels is not None and i < len(labels):
|
||
text = str(labels[i])
|
||
else:
|
||
text = f"{cname} {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"),
|
||
"class_name": region.get("class_name"),
|
||
})
|
||
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()
|