zhangpu
da90a803e4
- 新增碗尺寸识别主应用界面，支持图片拖拽、选择和摄像头实时检测 - 界面显示碗的尺寸类别、直径、置信度及算法方法 - 支持计算食物重量，依据秤示数和碗重差值计算 - 实现核心检测模块，采用霍夫圆变换和轮廓法两种检测方法融合 - 检测结果包含尺寸分类和置信度估计，支持绘制调试信息 - 提供交互式标定工具，支持批量加载样本图片并自动检测直径 - 标定工具计算阈值和碗重量，生成并保存配置文件config.json - 界面友好，提供当前状态提示和结果预览功能
247 lines
9.6 KiB
Python
247 lines
9.6 KiB
Python
"""
|
||
bowl_detector.py
|
||
碗尺寸检测核心模块
|
||
|
||
检测逻辑：
|
||
1. 霍夫圆变换（HoughCircles）- 最优先，精度高
|
||
2. 轮廓法（Contour + 最大类圆轮廓）- 备用
|
||
3. 两路结果取平均（如果都成功）- 提高鲁棒性
|
||
|
||
输入：BGR 图像（numpy array 或 文件路径）
|
||
输出：BowlDetectionResult，包含 size_label / pixel_diameter / confidence / debug_image
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import math
|
||
import cv2
|
||
import numpy as np
|
||
from dataclasses import dataclass
|
||
from typing import Optional, Tuple
|
||
from enum import Enum
|
||
|
||
|
||
class BowlSize(Enum):
|
||
SMALL = "小碗"
|
||
MEDIUM = "中碗"
|
||
LARGE = "大碗"
|
||
UNKNOWN = "未知"
|
||
|
||
|
||
@dataclass
|
||
class BowlDetectionResult:
|
||
size: BowlSize = BowlSize.UNKNOWN
|
||
pixel_diameter: float = 0.0
|
||
confidence: float = 0.0
|
||
method: str = "none"
|
||
center: Tuple[int, int] = (0, 0)
|
||
debug_image: Optional[object] = None
|
||
|
||
@property
|
||
def label(self) -> str:
|
||
return self.size.value
|
||
|
||
@property
|
||
def is_valid(self) -> bool:
|
||
"""直径检测成功即视为有效（未标定时 size=UNKNOWN 但直径仍可用于标定）"""
|
||
return self.pixel_diameter > 0
|
||
|
||
@property
|
||
def is_classified(self) -> bool:
|
||
"""不仅检测到直径，且完成了尺寸分类"""
|
||
return self.size != BowlSize.UNKNOWN and self.pixel_diameter > 0
|
||
|
||
|
||
class BowlDetector:
|
||
DEFAULT_CONFIG_PATH = os.path.join(os.path.dirname(__file__), "config.json")
|
||
|
||
def __init__(self, config_path: Optional[str] = None):
|
||
self.config_path = config_path or self.DEFAULT_CONFIG_PATH
|
||
self.config = self._load_config()
|
||
|
||
def detect(self, image_input, draw_debug: bool = True) -> BowlDetectionResult:
|
||
img = self._load_image(image_input)
|
||
if img is None:
|
||
return BowlDetectionResult()
|
||
|
||
h, w = img.shape[:2]
|
||
cfg = self.config.get("detection", {})
|
||
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
|
||
# CLAHE 对比度增强 —— 解决浅色碗+浅色背景对比度低的问题
|
||
clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
|
||
enhanced = clahe.apply(gray)
|
||
|
||
blur_k = cfg.get("blur_kernel_size", 11)
|
||
blur_k = blur_k if blur_k % 2 == 1 else blur_k + 1
|
||
blurred = cv2.GaussianBlur(enhanced, (blur_k, blur_k), 0)
|
||
|
||
hough_result = self._detect_by_hough(blurred, w, h, cfg)
|
||
contour_result = self._detect_by_contour(blurred, w, h, cfg)
|
||
|
||
# 如果两路都失败，用更宽松参数再试一次（备用策略）
|
||
if hough_result is None and contour_result is None:
|
||
relaxed_cfg = dict(cfg)
|
||
relaxed_cfg["hough_param2"] = max(10, cfg.get("hough_param2", 25) // 2)
|
||
relaxed_cfg["canny_threshold1"] = max(5, cfg.get("canny_threshold1", 20) // 2)
|
||
relaxed_cfg["canny_threshold2"] = max(20, cfg.get("canny_threshold2", 60) // 2)
|
||
relaxed_cfg["min_circularity"] = 0.4
|
||
hough_result = self._detect_by_hough(blurred, w, h, relaxed_cfg)
|
||
contour_result = self._detect_by_contour(blurred, w, h, relaxed_cfg)
|
||
merged = self._merge_results(hough_result, contour_result)
|
||
|
||
if merged["diameter"] > 0:
|
||
size, confidence = self._classify_size(merged["diameter"])
|
||
cx, cy = merged["center"]
|
||
radius = int(merged["diameter"] / 2)
|
||
detection = BowlDetectionResult(
|
||
size=size,
|
||
pixel_diameter=merged["diameter"],
|
||
confidence=confidence,
|
||
method=merged["method"],
|
||
center=(cx, cy),
|
||
)
|
||
if draw_debug:
|
||
detection.debug_image = self._draw_debug(img.copy(), cx, cy, radius, detection)
|
||
else:
|
||
detection = BowlDetectionResult()
|
||
if draw_debug:
|
||
detection.debug_image = img.copy()
|
||
|
||
return detection
|
||
|
||
def reload_config(self):
|
||
self.config = self._load_config()
|
||
|
||
def is_calibrated(self) -> bool:
|
||
return self.config.get("calibrated", False)
|
||
|
||
def get_bowl_weight(self, size: BowlSize) -> float:
|
||
key = {BowlSize.SMALL: "small", BowlSize.MEDIUM: "medium", BowlSize.LARGE: "large"}.get(size, "")
|
||
return self.config.get("bowl_weights_grams", {}).get(key, 0.0)
|
||
|
||
def _load_config(self) -> dict:
|
||
if os.path.exists(self.config_path):
|
||
with open(self.config_path, "r", encoding="utf-8") as f:
|
||
return json.load(f)
|
||
return {}
|
||
|
||
def _load_image(self, image_input) -> Optional[object]:
|
||
if isinstance(image_input, str):
|
||
return cv2.imdecode(np.fromfile(image_input, dtype=np.uint8), cv2.IMREAD_COLOR)
|
||
elif isinstance(image_input, np.ndarray):
|
||
return image_input.copy()
|
||
return None
|
||
|
||
def _detect_by_hough(self, blurred_gray, w, h, cfg) -> Optional[dict]:
|
||
min_r = int(min(w, h) * cfg.get("hough_min_radius_ratio", 0.05))
|
||
max_r = int(min(w, h) * cfg.get("hough_max_radius_ratio", 0.50))
|
||
min_dist = int(min(w, h) * cfg.get("hough_min_dist_ratio", 0.30))
|
||
circles = cv2.HoughCircles(
|
||
blurred_gray, cv2.HOUGH_GRADIENT,
|
||
dp=cfg.get("hough_dp", 1.2), minDist=min_dist,
|
||
param1=cfg.get("hough_param1", 80), param2=cfg.get("hough_param2", 35),
|
||
minRadius=min_r, maxRadius=max_r,
|
||
)
|
||
if circles is not None:
|
||
circles = np.round(circles[0, :]).astype(int)
|
||
best = max(circles, key=lambda c: c[2])
|
||
cx, cy, r = best
|
||
return {"center": (cx, cy), "diameter": float(r * 2), "method": "hough"}
|
||
return None
|
||
|
||
def _detect_by_contour(self, blurred_gray, w, h, cfg) -> Optional[dict]:
|
||
edges = cv2.Canny(blurred_gray, cfg.get("canny_threshold1", 30), cfg.get("canny_threshold2", 100))
|
||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
|
||
edges = cv2.dilate(edges, kernel, iterations=1)
|
||
contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
if not contours:
|
||
return None
|
||
|
||
img_area = w * h
|
||
min_area = img_area * cfg.get("min_area_ratio", 0.05)
|
||
max_area = img_area * cfg.get("max_area_ratio", 0.95)
|
||
min_circ = cfg.get("min_circularity", 0.60)
|
||
best_score, best_info = -1, None
|
||
|
||
for cnt in contours:
|
||
area = cv2.contourArea(cnt)
|
||
if area < min_area or area > max_area:
|
||
continue
|
||
perimeter = cv2.arcLength(cnt, True)
|
||
if perimeter == 0:
|
||
continue
|
||
circularity = 4 * math.pi * area / (perimeter ** 2)
|
||
if circularity < min_circ:
|
||
continue
|
||
(cx, cy), radius = cv2.minEnclosingCircle(cnt)
|
||
score = circularity * area
|
||
if score > best_score:
|
||
best_score = score
|
||
best_info = {"center": (int(cx), int(cy)), "diameter": float(radius * 2), "method": "contour"}
|
||
|
||
return best_info
|
||
|
||
def _merge_results(self, hough, contour) -> dict:
|
||
if hough and contour:
|
||
d = (hough["diameter"] + contour["diameter"]) / 2
|
||
cx = (hough["center"][0] + contour["center"][0]) // 2
|
||
cy = (hough["center"][1] + contour["center"][1]) // 2
|
||
return {"center": (cx, cy), "diameter": d, "method": "hough+contour"}
|
||
elif hough:
|
||
return hough
|
||
elif contour:
|
||
return contour
|
||
return {"center": (0, 0), "diameter": 0.0, "method": "none"}
|
||
|
||
def _classify_size(self, diameter: float) -> Tuple[BowlSize, float]:
|
||
thresholds = self.config.get("thresholds", {})
|
||
small_max = thresholds.get("small_max_diameter", 0)
|
||
medium_max = thresholds.get("medium_max_diameter", 0)
|
||
if not self.is_calibrated() or small_max == 0 or medium_max == 0:
|
||
return BowlSize.UNKNOWN, 0.0
|
||
if diameter <= small_max:
|
||
conf = min(1.0, 0.7 + (small_max - diameter) / small_max * 0.3)
|
||
return BowlSize.SMALL, round(conf, 3)
|
||
elif diameter <= medium_max:
|
||
margin = medium_max - small_max
|
||
conf = min(1.0, 0.7 + min(diameter - small_max, medium_max - diameter) / margin * 0.3)
|
||
return BowlSize.MEDIUM, round(conf, 3)
|
||
else:
|
||
conf = min(1.0, 0.7 + (diameter - medium_max) / medium_max * 0.3)
|
||
return BowlSize.LARGE, round(conf, 3)
|
||
|
||
def _draw_debug(self, img, cx, cy, radius, result) -> object:
|
||
color_map = {
|
||
BowlSize.SMALL: (0, 200, 100),
|
||
BowlSize.MEDIUM: (0, 165, 255),
|
||
BowlSize.LARGE: (0, 80, 255),
|
||
BowlSize.UNKNOWN: (128, 128, 128),
|
||
}
|
||
color = color_map[result.size]
|
||
cv2.circle(img, (cx, cy), radius, color, 3)
|
||
cv2.circle(img, (cx, cy), 5, color, -1)
|
||
cv2.line(img, (cx - radius, cy), (cx + radius, cy), color, 1)
|
||
lines = [
|
||
f"{result.label}",
|
||
f"Diameter: {result.pixel_diameter:.0f}px",
|
||
f"Confidence: {result.confidence:.0%}",
|
||
f"Method: {result.method}",
|
||
]
|
||
for i, line in enumerate(lines):
|
||
y = 35 + i * 30
|
||
cv2.putText(img, line, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.85, (0, 0, 0), 4, cv2.LINE_AA)
|
||
cv2.putText(img, line, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.85, color, 2, cv2.LINE_AA)
|
||
return img
|
||
|
||
|
||
_default_detector: Optional[BowlDetector] = None
|
||
|
||
|
||
def detect_bowl_size(image_input, config_path: Optional[str] = None) -> BowlDetectionResult:
|
||
global _default_detector
|
||
if _default_detector is None or config_path:
|
||
_default_detector = BowlDetector(config_path)
|
||
return _default_detector.detect(image_input)
|