import os import cv2 import json import numpy as np from pathlib import Path from dataclasses import dataclass, field from typing import Dict, List, Tuple, Optional, Any from enum import Enum import matplotlib.pyplot as plt from sklearn.linear_model import RANSACRegressor # matplotlib解决中文乱码 plt.rcParams["font.sans-serif"] = ["SimHei"] plt.rcParams["font.family"] = "sans-serif" plt.rcParams['axes.unicode_minus'] = False def fry_algo_print(level_str: str, info_str: str): print(f"[{level_str}] : {info_str}") def fry_cv2_imread(filename, flags=cv2.IMREAD_COLOR): """支持中文路径的imread""" try: with open(filename, 'rb') as f: chunk = f.read() chunk_arr = np.frombuffer(chunk, dtype=np.uint8) img = cv2.imdecode(chunk_arr, flags) if img is None: fry_algo_print("警告", f"Warning: Unable to decode image: {filename}") return img except IOError as e: fry_algo_print("错误", f"IOError: Unable to read file: {filename}") fry_algo_print("错误", f"Error details: {str(e)}") return None def fry_cv2_imwrite(filename, img, params=None): """支持中文路径的imwrite""" try: ext = os.path.splitext(filename)[1].lower() result, encoded_img = cv2.imencode(ext, img, params) if result: with open(filename, 'wb') as f: encoded_img.tofile(f) return True else: fry_algo_print("警告", f"Warning: Unable to encode image: {filename}") return False except Exception as e: fry_algo_print("错误", f"Error: Unable to write file: {filename}") fry_algo_print("错误", f"Error details: {str(e)}") return False def fry_opencv_Chinese_path_init(): """覆盖OpenCV的原始函数""" cv2.imread = fry_cv2_imread cv2.imwrite = fry_cv2_imwrite OPENCV_IO_ALREADY_INIT = False if not OPENCV_IO_ALREADY_INIT: fry_opencv_Chinese_path_init() OPENCV_IO_ALREADY_INIT = True def to_json_serializable(obj): """将包含自定义对象的数据结构转换为JSON可序列化的格式""" if obj is None or isinstance(obj, (bool, int, float, str)): return obj elif isinstance(obj, dict): return {key: to_json_serializable(value) for key, value in obj.items()} elif isinstance(obj, (list, tuple)): return [to_json_serializable(item) for item in obj] elif isinstance(obj, set): return [to_json_serializable(item) for item in obj] elif isinstance(obj, bytes): return obj.decode('utf-8', errors='ignore') else: if hasattr(obj, '__dict__'): return to_json_serializable(obj.__dict__) elif hasattr(obj, 'to_dict'): return to_json_serializable(obj.to_dict()) elif hasattr(obj, 'to_json'): return to_json_serializable(obj.to_json()) else: return str(obj) class EdgeCornerType(Enum): """边角类型枚举""" TOP_LEFT = "左上角" TOP_RIGHT = "右上角" BOTTOM_LEFT = "左下角" BOTTOM_RIGHT = "右下角" TOP = "上边" BOTTOM = "下边" LEFT = "左边" RIGHT = "右边" @dataclass class EdgeCornerDefect: """边角缺陷数据结构""" type: EdgeCornerType protrusion_area: float = 0.0 # 突出面积（像素） depression_area: float = 0.0 # 凹陷面积（像素） protrusion_pixels: int = 0 # 突出像素数 depression_pixels: int = 0 # 凹陷像素数 contour_points: List[List[float]] = field(default_factory=list) # 该边角的轮廓点 fitted_points: List[List[float]] = field(default_factory=list) # 拟合的边角点 region_points: List[List[float]] = field(default_factory=list) # 区域边界点 def to_dict(self) -> Dict[str, Any]: """转换为字典""" return { "type": self.type.value, "protrusion_area": float(self.protrusion_area), "depression_area": float(self.depression_area), "protrusion_pixels": int(self.protrusion_pixels), "depression_pixels": int(self.depression_pixels), "contour_points": self.contour_points, "fitted_points": self.fitted_points, "region_points": self.region_points } @classmethod def load_from_dict(cls, data: Dict[str, Any]) -> 'EdgeCornerDefect': """从字典加载""" return cls( type=EdgeCornerType(data["type"]), protrusion_area=data.get("protrusion_area", 0.0), depression_area=data.get("depression_area", 0.0), protrusion_pixels=data.get("protrusion_pixels", 0), depression_pixels=data.get("depression_pixels", 0), contour_points=data.get("contour_points", []), fitted_points=data.get("fitted_points", []), region_points=data.get("region_points", []) ) @dataclass class FryAlgoParamsBase: """算法参数基类""" algo_name: Optional[str] = None debug_level: str = "no" # no, normal, detail save_json_params: bool = False is_first_algo: bool = False is_last_algo: bool = False def to_dict(self) -> Dict[str, Any]: """转换为字典""" result = {} for key, value in self.__dict__.items(): if hasattr(value, 'to_dict'): result[key] = value.to_dict() elif isinstance(value, Enum): result[key] = value.value elif isinstance(value, Path): result[key] = str(value) else: result[key] = value return result @dataclass class CardDefectDetectionParams(FryAlgoParamsBase): """卡片缺陷检测算法参数""" algo_name: str = "卡片边角缺陷检测" label_name: str = "outer_box" # JSON中的框标签名称 long_edge_corner_length: float = 200 # 矩形长边的角的长度 short_edge_corner_length: float = 200 # 矩形短边的角的长度 edge_width: float = 50 # 边区域向内延伸的宽度 iteration_rounds: int = 5 # 步骤d和e迭代轮数 ransac_residual_threshold: float = 5.0 # RANSAC残差阈值 save_intermediate_results: bool = True # 是否保存中间结果 save_overlay_images: bool = True # 是否保存叠加图像 class MaskProcessor: """Mask处理工具类""" @staticmethod def polygon_to_mask(polygon_points: np.ndarray, img_shape: Tuple[int, int]) -> np.ndarray: """将多边形转换为mask Args: polygon_points: 多边形顶点坐标 img_shape: (height, width) Returns: 二值mask图像 """ mask = np.zeros(img_shape[:2], dtype=np.uint8) if len(polygon_points) > 0: cv2.fillPoly(mask, [polygon_points.astype(np.int32)], 255) return mask @staticmethod def calculate_inner_points( fitted_quad: np.ndarray, long_edge_indices: List[int], short_edge_indices: List[int], long_corner_length: float, short_corner_length: float ) -> np.ndarray: """计算内部四个点B1-B4 Args: fitted_quad: 拟合四边形的四个顶点 long_edge_indices: 长边索引 short_edge_indices: 短边索引 long_corner_length: 长边角长度 short_corner_length: 短边角长度 Returns: 内部四个点坐标 """ inner_points = np.zeros((4, 2)) for i in range(4): # 获取当前顶点 current_point = fitted_quad[i] # 获取相邻两条边的向量 prev_idx = (i - 1) % 4 next_idx = (i + 1) % 4 edge_to_prev = fitted_quad[prev_idx] - current_point edge_to_next = fitted_quad[next_idx] - current_point # 判断边的类型并设置对应的长度 # 边索引: 0(0-1), 1(1-2), 2(2-3), 3(3-0) edge_prev_idx = prev_idx if prev_idx < i else 3 edge_next_idx = i # 根据边类型设置移动距离 if edge_prev_idx in long_edge_indices: dist_prev = long_corner_length else: dist_prev = short_corner_length if edge_next_idx in long_edge_indices: dist_next = long_corner_length else: dist_next = short_corner_length # 单位化向量 edge_to_prev_unit = edge_to_prev / np.linalg.norm(edge_to_prev) edge_to_next_unit = edge_to_next / np.linalg.norm(edge_to_next) # 计算内部点 inner_points[i] = current_point + dist_prev * edge_to_prev_unit + dist_next * edge_to_next_unit return inner_points @staticmethod def get_corner_region_mask( corner_idx: int, fitted_quad: np.ndarray, inner_points: np.ndarray, img_shape: Tuple[int, int] ) -> np.ndarray: """获取角区域mask Args: corner_idx: 角索引(0-3) fitted_quad: 拟合四边形顶点 inner_points: 内部四个点 img_shape: 图像尺寸 Returns: 角区域mask """ mask = np.zeros(img_shape[:2], dtype=np.uint8) # 获取角区域的四个顶点 prev_idx = (corner_idx - 1) % 4 next_idx = (corner_idx + 1) % 4 # 计算角区域的边界点 corner_point = fitted_quad[corner_idx] inner_point = inner_points[corner_idx] # 从内部点向两条边做垂线得到边界点 edge_to_prev = fitted_quad[prev_idx] - corner_point edge_to_next = fitted_quad[next_idx] - corner_point edge_to_prev_unit = edge_to_prev / np.linalg.norm(edge_to_prev) edge_to_next_unit = edge_to_next / np.linalg.norm(edge_to_next) # 计算垂足 proj_prev = np.dot(inner_point - corner_point, edge_to_prev_unit) proj_next = np.dot(inner_point - corner_point, edge_to_next_unit) point_on_prev_edge = corner_point + proj_prev * edge_to_prev_unit point_on_next_edge = corner_point + proj_next * edge_to_next_unit # 创建角区域多边形 corner_polygon = np.array([ corner_point, point_on_prev_edge, inner_point, point_on_next_edge ], dtype=np.int32) cv2.fillPoly(mask, [corner_polygon], 255) return mask @staticmethod def get_edge_region_mask( edge_idx: int, fitted_quad: np.ndarray, inner_points: np.ndarray, edge_width: float, img_shape: Tuple[int, int] ) -> np.ndarray: """获取边区域mask Args: edge_idx: 边索引(0-3) fitted_quad: 拟合四边形顶点 inner_points: 内部四个点 edge_width: 边区域宽度 img_shape: 图像尺寸 Returns: 边区域mask """ mask = np.zeros(img_shape[:2], dtype=np.uint8) # 获取边的两个端点 p1_idx = edge_idx p2_idx = (edge_idx + 1) % 4 p1 = fitted_quad[p1_idx] p2 = fitted_quad[p2_idx] # 获取对应的内部点 inner_p1 = inner_points[p1_idx] inner_p2 = inner_points[p2_idx] # 计算边向量和法向量 edge_vec = p2 - p1 edge_unit = edge_vec / np.linalg.norm(edge_vec) # 法向量（向内） normal = np.array([-edge_unit[1], edge_unit[0]]) # 确保法向量指向内部 center = np.mean(fitted_quad, axis=0) if np.dot(normal, center - (p1 + p2) / 2) < 0: normal = -normal # 从内部点向边做垂线 proj1 = np.dot(inner_p1 - p1, edge_unit) proj2 = np.dot(inner_p2 - p1, edge_unit) point_on_edge_1 = p1 + proj1 * edge_unit point_on_edge_2 = p1 + proj2 * edge_unit # 创建边区域梯形 inner_edge_p1 = point_on_edge_1 + edge_width * normal inner_edge_p2 = point_on_edge_2 + edge_width * normal edge_polygon = np.array([ point_on_edge_1, point_on_edge_2, inner_edge_p2, inner_edge_p1 ], dtype=np.int32) cv2.fillPoly(mask, [edge_polygon], 255) return mask @staticmethod def calculate_defect_pixels( contour_mask: np.ndarray, fitted_mask: np.ndarray, region_mask: np.ndarray ) -> Tuple[int, int]: """计算缺陷像素数 Args: contour_mask: 轮廓mask fitted_mask: 拟合mask region_mask: 区域mask Returns: (突出像素数, 凹陷像素数) """ # 限制在区域内 contour_in_region = cv2.bitwise_and(contour_mask, region_mask) fitted_in_region = cv2.bitwise_and(fitted_mask, region_mask) # 计算并集 union_mask = cv2.bitwise_or(contour_in_region, fitted_in_region) # 突出：并集减去拟合区域 protrusion_mask = cv2.bitwise_and(union_mask, cv2.bitwise_not(fitted_in_region)) protrusion_pixels = np.sum(protrusion_mask > 0) # 凹陷：并集减去轮廓区域 depression_mask = cv2.bitwise_and(union_mask, cv2.bitwise_not(contour_in_region)) depression_pixels = np.sum(depression_mask > 0) return protrusion_pixels, depression_pixels class CardDefectDetectionAlgo: """卡片边角缺陷检测算法类""" def __init__(self): """初始化""" self.intermediate_results = {} self.mask_processor = MaskProcessor() @staticmethod def load_json_data(json_path) -> Dict[str, Any]: """加载JSON数据""" with open(json_path, 'r', encoding='utf-8') as f: return json.load(f) @staticmethod def _create_blank_image(width: int, height: int) -> np.ndarray: """创建空白图像""" return np.ones((height, width, 3), dtype=np.uint8) * 255 @staticmethod def load_or_create_image(image_path: str) -> np.ndarray: """加载图像或创建空白图像""" if os.path.exists(image_path): img = cv2.imread(str(image_path)) if img is not None: return img raise ValueError("警告", f"图像不存在或无法读取: {image_path}，创建空白图像") @staticmethod def _get_unique_shape(shapes: List[Dict], label_name: str) -> Optional[Dict]: """获取唯一的形状""" target_shapes = [s for s in shapes if s.get('label') == label_name] if not target_shapes: return None if len(target_shapes) == 1: return target_shapes[0] def get_shape_score(shape): confidence = shape.get('probability', 0) points = np.array(shape['points']) area = cv2.contourArea(points.astype(np.float32)) return (confidence, area) target_shapes.sort(key=get_shape_score, reverse=True) return target_shapes[0] @staticmethod def _get_min_area_rect(points: np.ndarray) -> Tuple[np.ndarray, float, float, List[int], List[int]]: """获取最小外接矩形并识别长短边 Returns: (矩形顶点, 长边长度, 短边长度, 长边索引列表, 短边索引列表) """ rect = cv2.minAreaRect(points.astype(np.float32)) box = cv2.boxPoints(rect) box = np.intp(box) # 计算每条边的长度 edge_lengths = [] for i in range(4): p1 = box[i] p2 = box[(i + 1) % 4] length = np.linalg.norm(p2 - p1) edge_lengths.append(length) # 识别长边和短边 sorted_indices = np.argsort(edge_lengths) short_edge_indices = sorted_indices[:2].tolist() long_edge_indices = sorted_indices[2:].tolist() long_edge = np.mean([edge_lengths[i] for i in long_edge_indices]) short_edge = np.mean([edge_lengths[i] for i in short_edge_indices]) return box, long_edge, short_edge, long_edge_indices, short_edge_indices @staticmethod def _classify_points_to_edges_corners( contour_points: np.ndarray, rect_corners: np.ndarray, long_edge: float, short_edge: float, long_corner_length: float, short_corner_length: float ) -> Dict[EdgeCornerType, np.ndarray]: """将轮廓点分类到边和角""" classified_points = {edge_type: [] for edge_type in EdgeCornerType} edges = [] for i in range(4): p1 = rect_corners[i] p2 = rect_corners[(i + 1) % 4] edges.append((p1, p2)) edge_lengths = [np.linalg.norm(e[1] - e[0]) for e in edges] long_edge_indices = np.argsort(edge_lengths)[-2:] short_edge_indices = np.argsort(edge_lengths)[:2] for point in contour_points: min_dist = float('inf') min_type = None for i, corner in enumerate(rect_corners): dist = np.linalg.norm(point - corner) if dist < min_dist: min_dist = dist if i == 0: min_type = EdgeCornerType.TOP_LEFT elif i == 1: min_type = EdgeCornerType.TOP_RIGHT elif i == 2: min_type = EdgeCornerType.BOTTOM_RIGHT else: min_type = EdgeCornerType.BOTTOM_LEFT for i, (p1, p2) in enumerate(edges): dist = CardDefectDetectionAlgo._point_to_segment_distance(point, p1, p2) edge_vector = p2 - p1 edge_length = np.linalg.norm(edge_vector) projection = np.dot(point - p1, edge_vector) / edge_length if i in long_edge_indices: corner_ratio = long_corner_length / edge_length else: corner_ratio = short_corner_length / edge_length if corner_ratio < projection / edge_length < (1 - corner_ratio): if dist < min_dist: min_dist = dist if i == 0: min_type = EdgeCornerType.TOP elif i == 1: min_type = EdgeCornerType.RIGHT elif i == 2: min_type = EdgeCornerType.BOTTOM else: min_type = EdgeCornerType.LEFT if min_type: classified_points[min_type].append(point) for key in classified_points: if classified_points[key]: classified_points[key] = np.array(classified_points[key]) else: classified_points[key] = np.array([]).reshape(0, 2) return classified_points @staticmethod def _point_to_segment_distance(point: np.ndarray, seg_p1: np.ndarray, seg_p2: np.ndarray) -> float: """计算点到线段的距离""" line_vec = seg_p2 - seg_p1 point_vec = point - seg_p1 line_len = np.linalg.norm(line_vec) if line_len == 0: return np.linalg.norm(point - seg_p1) line_unitvec = line_vec / line_len proj_length = np.dot(point_vec, line_unitvec) if proj_length < 0: return np.linalg.norm(point - seg_p1) elif proj_length > line_len: return np.linalg.norm(point - seg_p2) else: proj_point = seg_p1 + proj_length * line_unitvec return np.linalg.norm(point - proj_point) @staticmethod def _fit_line_with_ransac(points: np.ndarray, max_iterations: int = 1000, threshold: float = 5.0) -> Tuple[float, float, float]: """使用RANSAC拟合直线 Args: points: 点集 (N, 2) max_iterations: 最大迭代次数 threshold: 内点判定阈值 Returns: (a, b, c) 直线方程系数 ax + by + c = 0 """ if len(points) < 2: return None best_line = None best_inliers = 0 for _ in range(max_iterations): # 随机选择两个点 idx = np.random.choice(len(points), 2, replace=False) p1, p2 = points[idx] # 计算直线参数 if abs(p2[0] - p1[0]) < 1e-6: # 垂直线 a, b, c = 1, 0, -p1[0] else: # 一般直线 y = kx + m => kx - y + m = 0 k = (p2[1] - p1[1]) / (p2[0] - p1[0]) m = p1[1] - k * p1[0] a, b, c = k, -1, m # 归一化 norm = np.sqrt(a * a + b * b) a, b, c = a / norm, b / norm, c / norm # 计算内点数 distances = np.abs(a * points[:, 0] + b * points[:, 1] + c) inliers = np.sum(distances < threshold) if inliers > best_inliers: best_inliers = inliers best_line = (a, b, c) # 使用所有内点重新拟合（精细化） if best_line: a, b, c = best_line distances = np.abs(a * points[:, 0] + b * points[:, 1] + c) inlier_points = points[distances < threshold] if len(inlier_points) >= 2: # 使用OpenCV的fitLine进行最终拟合 vx, vy, x0, y0 = cv2.fitLine(inlier_points, cv2.DIST_L2, 0, 0.01, 0.01) # 转换为一般式 ax + by + c = 0 # 直线方向向量(vx, vy)，点(x0, y0)在直线上 # 法向量为(-vy, vx) a = -vy[0] b = vx[0] c = -(a * x0[0] + b * y0[0]) # 归一化 norm = np.sqrt(a * a + b * b) return (a / norm, b / norm, c / norm) return best_line @staticmethod def _fit_lines_opencv( classified_points: Dict[EdgeCornerType, np.ndarray], threshold: float = 5.0 ) -> Dict[EdgeCornerType, Tuple[float, float, float]]: """使用OpenCV拟合四条边的直线（只使用边点，不使用角点） Args: classified_points: 分类后的点 threshold: RANSAC阈值 Returns: 拟合的直线字典 {边类型: (a, b, c)} """ fitted_lines = {} # 边类型映射 edge_only_types = { EdgeCornerType.TOP: EdgeCornerType.TOP, EdgeCornerType.BOTTOM: EdgeCornerType.BOTTOM, EdgeCornerType.LEFT: EdgeCornerType.LEFT, EdgeCornerType.RIGHT: EdgeCornerType.RIGHT } for edge_type in edge_only_types: # 只使用纯边点，不包含角点 edge_points = classified_points.get(edge_type, np.array([])) if len(edge_points) < 2: fry_algo_print("警告", f"{edge_type.value}边点数不足，跳过拟合") continue # 使用RANSAC拟合直线 line_params = CardDefectDetectionAlgo._fit_line_with_ransac( edge_points, threshold=threshold ) if line_params: fitted_lines[edge_type] = line_params fry_algo_print("信息", f"{edge_type.value}拟合成功，点数: {len(edge_points)}") else: fry_algo_print("警告", f"{edge_type.value}拟合失败") return fitted_lines @staticmethod def _get_line_intersection(line1: Tuple[float, float, float], line2: Tuple[float, float, float]) -> Optional[np.ndarray]: """计算两条直线的交点 Args: line1: 直线1参数 (a1, b1, c1) line2: 直线2参数 (a2, b2, c2) Returns: 交点坐标 [x, y] 或 None """ a1, b1, c1 = line1 a2, b2, c2 = line2 det = a1 * b2 - a2 * b1 if abs(det) < 1e-10: return None x = (b1 * c2 - b2 * c1) / det y = (a2 * c1 - a1 * c2) / det return np.array([x, y]) @staticmethod def _get_quadrilateral_from_lines(fitted_lines: Dict[EdgeCornerType, Tuple[float, float, float]]) -> np.ndarray: """从拟合的直线获取四边形的四个角点 Args: fitted_lines: 拟合的直线字典 Returns: 四个角点坐标，按左上、右上、右下、左下顺序 """ # 检查是否有四条边 required_edges = [EdgeCornerType.TOP, EdgeCornerType.RIGHT, EdgeCornerType.BOTTOM, EdgeCornerType.LEFT] for edge in required_edges: if edge not in fitted_lines: fry_algo_print("警告", f"缺少{edge.value}的拟合直线") return None # 计算四个角点 corners = {} # 左上角 = 左边 ∩ 上边 top_left = CardDefectDetectionAlgo._get_line_intersection( fitted_lines[EdgeCornerType.LEFT], fitted_lines[EdgeCornerType.TOP] ) if top_left is not None: corners['top_left'] = top_left # 右上角 = 右边 ∩ 上边 top_right = CardDefectDetectionAlgo._get_line_intersection( fitted_lines[EdgeCornerType.RIGHT], fitted_lines[EdgeCornerType.TOP] ) if top_right is not None: corners['top_right'] = top_right # 右下角 = 右边 ∩ 下边 bottom_right = CardDefectDetectionAlgo._get_line_intersection( fitted_lines[EdgeCornerType.RIGHT], fitted_lines[EdgeCornerType.BOTTOM] ) if bottom_right is not None: corners['bottom_right'] = bottom_right # 左下角 = 左边 ∩ 下边 bottom_left = CardDefectDetectionAlgo._get_line_intersection( fitted_lines[EdgeCornerType.LEFT], fitted_lines[EdgeCornerType.BOTTOM] ) if bottom_left is not None: corners['bottom_left'] = bottom_left # 检查是否获得了所有角点 if len(corners) != 4: fry_algo_print("警告", f"只计算出{len(corners)}个角点") return None # 按顺序返回 return np.array([ corners['top_left'], corners['top_right'], corners['bottom_right'], corners['bottom_left'] ]) def _calculate_defects_with_mask( self, contour_points: np.ndarray, fitted_quad: np.ndarray, classified_points: Dict[EdgeCornerType, np.ndarray], img_shape: Tuple[int, int], long_corner_length: float, short_corner_length: float, edge_width: float ) -> Tuple[Dict[EdgeCornerType, EdgeCornerDefect], np.ndarray]: """使用mask方法计算缺陷 Returns: (缺陷字典, 内部四个点) """ defects = {} # 获取拟合四边形的最小外接矩形，识别长短边 _, _, _, long_edge_indices, short_edge_indices = self._get_min_area_rect(fitted_quad) # 计算内部四个点 inner_points = self.mask_processor.calculate_inner_points( fitted_quad, long_edge_indices, short_edge_indices, long_corner_length, short_corner_length ) # 创建轮廓和拟合四边形的mask contour_mask = self.mask_processor.polygon_to_mask(contour_points, img_shape) fitted_mask = self.mask_processor.polygon_to_mask(fitted_quad, img_shape) # 处理四个角 corner_types = [ EdgeCornerType.TOP_LEFT, EdgeCornerType.TOP_RIGHT, EdgeCornerType.BOTTOM_RIGHT, EdgeCornerType.BOTTOM_LEFT ] for i, corner_type in enumerate(corner_types): defect = EdgeCornerDefect(type=corner_type) # 获取角区域mask region_mask = self.mask_processor.get_corner_region_mask( i, fitted_quad, inner_points, img_shape ) # 计算缺陷 protrusion_pixels, depression_pixels = self.mask_processor.calculate_defect_pixels( contour_mask, fitted_mask, region_mask ) defect.protrusion_pixels = protrusion_pixels defect.depression_pixels = depression_pixels defect.protrusion_area = float(protrusion_pixels) defect.depression_area = float(depression_pixels) # 保存该区域的点 edge_corner_points = classified_points.get(corner_type, np.array([])) if len(edge_corner_points) > 0: defect.contour_points = edge_corner_points.tolist() defects[corner_type] = defect # 处理四条边 edge_types = [ EdgeCornerType.TOP, EdgeCornerType.RIGHT, EdgeCornerType.BOTTOM, EdgeCornerType.LEFT ] for i, edge_type in enumerate(edge_types): defect = EdgeCornerDefect(type=edge_type) # 获取边区域mask region_mask = self.mask_processor.get_edge_region_mask( i, fitted_quad, inner_points, edge_width, img_shape ) # 计算缺陷 protrusion_pixels, depression_pixels = self.mask_processor.calculate_defect_pixels( contour_mask, fitted_mask, region_mask ) defect.protrusion_pixels = protrusion_pixels defect.depression_pixels = depression_pixels defect.protrusion_area = float(protrusion_pixels) defect.depression_area = float(depression_pixels) # 保存该区域的点 edge_points = classified_points.get(edge_type, np.array([])) if len(edge_points) > 0: defect.contour_points = edge_points.tolist() defects[edge_type] = defect return defects, inner_points def _draw_text_with_background( self, img: np.ndarray, text: str, position: Tuple[int, int], font_scale: float = 0.6, font_thickness: int = 1, text_color: Tuple[int, int, int] = (255, 255, 255), bg_color: Tuple[int, int, int] = (0, 0, 0), bg_opacity: float = 0.7 ) -> np.ndarray: """在图像上绘制带半透明背景的文字""" font = cv2.FONT_HERSHEY_SIMPLEX (text_width, text_height), baseline = cv2.getTextSize( text, font, font_scale, font_thickness ) x, y = position padding = 5 bg_pt1 = (x - padding, y - text_height - padding) bg_pt2 = (x + text_width + padding, y + baseline + padding) overlay = img.copy() cv2.rectangle(overlay, bg_pt1, bg_pt2, bg_color, -1) cv2.addWeighted(overlay, bg_opacity, img, 1 - bg_opacity, 0, img) cv2.putText(img, text, position, font, font_scale, text_color, font_thickness) return img def _draw_overlay_on_image( self, image: np.ndarray, step_name: str, contour_points: Optional[np.ndarray] = None, min_rect: Optional[np.ndarray] = None, fitted_quad: Optional[np.ndarray] = None, inner_points: Optional[np.ndarray] = None, classified_points: Optional[Dict[EdgeCornerType, np.ndarray]] = None, defects: Optional[Dict[EdgeCornerType, EdgeCornerDefect]] = None ) -> np.ndarray: """在图像上绘制叠加可视化""" overlay = image.copy() # 绘制轮廓 if contour_points is not None and len(contour_points) > 0: cv2.polylines(overlay, [contour_points.astype(np.int32)], True, (255, 0, 0), 2) # 绘制最小外接矩形 if min_rect is not None and len(min_rect) == 4: cv2.polylines(overlay, [min_rect.astype(np.int32)], True, (0, 255, 0), 2) # 绘制拟合四边形 if fitted_quad is not None and len(fitted_quad) == 4: cv2.polylines(overlay, [fitted_quad.astype(np.int32)], True, (0, 0, 255), 3) for i, point in enumerate(fitted_quad): cv2.circle(overlay, tuple(point.astype(int)), 8, (0, 0, 255), -1) self._draw_text_with_background( overlay, f"C{i}", tuple(point.astype(int)), font_scale=0.8, text_color=(255, 255, 255) ) # 绘制内部四个点 if inner_points is not None and len(inner_points) == 4: for i, point in enumerate(inner_points): cv2.circle(overlay, tuple(point.astype(int)), 10, (255, 0, 255), -1) self._draw_text_with_background( overlay, f"B{i + 1}", tuple(point.astype(int) + np.array([15, 0])), font_scale=0.8, text_color=(255, 255, 0), bg_color=(128, 0, 128) ) # 绘制连接线（可选） if fitted_quad is not None: cv2.line(overlay, tuple(fitted_quad[i].astype(int)), tuple(point.astype(int)), (255, 0, 255), 1, cv2.LINE_AA) # 绘制分类点 if classified_points: colors = { EdgeCornerType.TOP_LEFT: (255, 0, 0), EdgeCornerType.TOP_RIGHT: (0, 255, 0), EdgeCornerType.BOTTOM_LEFT: (255, 255, 0), EdgeCornerType.BOTTOM_RIGHT: (0, 255, 255), EdgeCornerType.TOP: (255, 0, 255), EdgeCornerType.BOTTOM: (128, 0, 255), EdgeCornerType.LEFT: (0, 128, 255), EdgeCornerType.RIGHT: (255, 128, 0) } for edge_type, points in classified_points.items(): if len(points) > 0: color = colors.get(edge_type, (128, 128, 128)) for point in points: cv2.circle(overlay, tuple(point.astype(int)), 3, color, -1) # 绘制缺陷信息 if defects: y_offset = 30 for edge_type, defect in defects.items(): if defect.protrusion_pixels >= 0 or defect.depression_pixels >= 0: text = f"{edge_type.name}: protrusion={defect.protrusion_pixels}px, depression={defect.depression_pixels}px" self._draw_text_with_background( overlay, text, (10, y_offset), font_scale=0.5, text_color=(255, 255, 255), bg_color=(0, 0, 128), bg_opacity=0.7 ) y_offset += 25 # 添加步骤标题 self._draw_text_with_background( overlay, step_name, (10, overlay.shape[0] - 20), font_scale=1.0, text_color=(255, 255, 255), bg_color=(255, 0, 0), bg_opacity=0.7 ) return overlay def _save_intermediate_result( self, step_num: int, step_name: str, data: Any, output_dir: Path, image: Optional[np.ndarray] = None, save_overlay: bool = True, inner_points: Optional[np.ndarray] = None ): """保存中间结果""" step_dir = output_dir / "intermediate_results" step_dir.mkdir(parents=True, exist_ok=True) json_path = step_dir / f"step_{step_num:02d}_{step_name}.json" with open(json_path, 'w', encoding='utf-8') as f: json.dump(to_json_serializable(data), f, ensure_ascii=False, indent=2) if save_overlay and image is not None: overlay_image = self._draw_overlay_on_image( image=image, step_name=f"Step {step_num}: {step_name}", contour_points=np.array(data.get("contour", [])) if "contour" in data else None, min_rect=np.array(data.get("min_rect", [])) if "min_rect" in data else None, fitted_quad=np.array(data.get("fitted_quad", [])) if "fitted_quad" in data else None, inner_points=inner_points, classified_points={EdgeCornerType(k): np.array(v) for k, v in data.get("classified_points", {}).items()} if "classified_points" in data else None, defects={EdgeCornerType(k): EdgeCornerDefect.load_from_dict(v) for k, v in data.get("defects", {}).items()} if "defects" in data else None ) overlay_path = step_dir / f"step_{step_num:02d}_{step_name}_overlay.jpg" cv2.imwrite(str(overlay_path), overlay_image) def _interpolate_contour_points(self, contour_points: np.ndarray, interpolation_method: str = "linear") -> np.ndarray: """ 在轮廓点之间插入新点，增加点的密度 Args: contour_points: 原始轮廓点 interpolation_method: 插值方法，可选 "linear" 或 "midpoint" Returns: 插值后的轮廓点 """ if len(contour_points) < 2: return contour_points interpolated_points = [] for i in range(len(contour_points)): # 添加原始点 interpolated_points.append(contour_points[i]) # 获取下一个点（形成闭合轮廓） next_idx = (i + 1) % len(contour_points) next_point = contour_points[next_idx] if interpolation_method == "midpoint": # 简单的中点插值 mid_point = (contour_points[i] + next_point) / 2.0 interpolated_points.append(mid_point) elif interpolation_method == "linear": # 线性插值，可以插入多个点 distance = np.linalg.norm(next_point - contour_points[i]) # 根据距离决定插入点的数量（每50像素插入一个点） num_insert = max(1, int(distance / 50)) for j in range(1, num_insert + 1): t = j / (num_insert + 1) interp_point = (1 - t) * contour_points[i] + t * next_point interpolated_points.append(interp_point) return np.array(interpolated_points, dtype=np.int32) def process_single_group_entrance( self, data_dict: Dict, algo_params: CardDefectDetectionParams ) -> Dict: """处理单个数据组的核心逻辑""" if algo_params.debug_level == 'detail': fry_algo_print("信息", f"开始处理数据，参数: {algo_params.to_dict()}") image = data_dict.get("image") json_data = data_dict.get("label", {}) output_dir = Path(data_dict.get("output_dir", "./output")) output_dir.mkdir(parents=True, exist_ok=True) # 获取图像尺寸 if image is not None: img_shape = image.shape else: img_shape = (json_data.get("imageHeight", 2048), json_data.get("imageWidth", 2048), 3) # 步骤1：读取并筛选唯一形状 shapes = json_data.get("shapes", []) unique_shape = self._get_unique_shape(shapes, algo_params.label_name) if unique_shape is None: fry_algo_print("错误", f"未找到标签为{algo_params.label_name}的形状") return data_dict contour_points = np.array(unique_shape["points"]) if algo_params.save_intermediate_results: self._save_intermediate_result( 1, "筛选唯一形状", {"contour": contour_points.tolist(), "shape_info": unique_shape}, output_dir, image, algo_params.save_overlay_images ) # ========== 步骤1.5: 轮廓点插值（新增） ========== interpolated_contour = self._interpolate_contour_points(contour_points, "linear") if algo_params.debug_level == 'detail': fry_algo_print("信息", f"轮廓点插值: {len(contour_points)} -> {len(interpolated_contour)} 点") # 使用插值后的轮廓点进行后续处理 contour_points = interpolated_contour # 步骤2：获取最小外接矩形 min_rect, long_edge, short_edge, _, _ = self._get_min_area_rect(contour_points) if algo_params.save_intermediate_results: self._save_intermediate_result( 2, "最小外接矩形", { "contour": contour_points.tolist(), "min_rect": min_rect.tolist(), "long_edge": float(long_edge), "short_edge": float(short_edge) }, output_dir, image, algo_params.save_overlay_images ) # 步骤3：初始分类 classified_points = self._classify_points_to_edges_corners( contour_points, min_rect, long_edge, short_edge, algo_params.long_edge_corner_length, algo_params.short_edge_corner_length ) if algo_params.save_intermediate_results: classified_dict = {k.value: v.tolist() if len(v) > 0 else [] for k, v in classified_points.items()} self._save_intermediate_result( 3, "初始点分类", {"classified_points": classified_dict, "min_rect": min_rect.tolist()}, output_dir, image, algo_params.save_overlay_images ) # 步骤4-5：迭代拟合 fitted_quad = min_rect for iteration in range(algo_params.iteration_rounds): # 使用新的OpenCV拟合方法（只用边点） fitted_lines = self._fit_lines_opencv( classified_points, algo_params.ransac_residual_threshold ) new_quad = self._get_quadrilateral_from_lines(fitted_lines) if new_quad is not None: fitted_quad = new_quad # 重新分类点 classified_points = self._classify_points_to_edges_corners( contour_points, fitted_quad, long_edge, short_edge, algo_params.long_edge_corner_length, algo_params.short_edge_corner_length ) if algo_params.save_intermediate_results and iteration == algo_params.iteration_rounds - 1: classified_dict = {k.value: v.tolist() if len(v) > 0 else [] for k, v in classified_points.items()} # 转换fitted_lines格式用于保存 fitted_lines_dict = {k.value: v for k, v in fitted_lines.items()} self._save_intermediate_result( 4 + iteration, f"迭代{iteration + 1}_拟合结果", { "classified_points": classified_dict, "fitted_quad": fitted_quad.tolist(), "fitted_lines": fitted_lines_dict }, output_dir, image, algo_params.save_overlay_images ) # 步骤6：计算缺陷 defects, inner_points = self._calculate_defects_with_mask( contour_points, fitted_quad, classified_points, img_shape, algo_params.long_edge_corner_length, algo_params.short_edge_corner_length, algo_params.edge_width ) defects_dict = {k.value: v.to_dict() for k, v in defects.items()} if algo_params.save_intermediate_results: self._save_intermediate_result( 5 + algo_params.iteration_rounds, "最终缺陷结果", { "defects": defects_dict, "fitted_quad": fitted_quad.tolist(), "inner_points": inner_points.tolist(), "contour": contour_points.tolist() }, output_dir, image, algo_params.save_overlay_images, inner_points ) # 保存最终结果 result_path = output_dir / "defect_detection_result.json" final_result = { "defects": defects_dict, "fitted_quad": fitted_quad.tolist(), "inner_points": inner_points.tolist() } with open(result_path, 'w', encoding='utf-8') as f: json.dump(final_result, f, ensure_ascii=False, indent=2) if algo_params.debug_level in ['normal', 'detail']: fry_algo_print("成功", f"缺陷检测完成，结果已保存到: {result_path}") data_dict["defects"] = defects_dict data_dict["fitted_quad"] = fitted_quad.tolist() data_dict["inner_points"] = inner_points.tolist() return data_dict def process_batch_group_entrance( self, input_dir_str: str, output_dir_str: str, algo_params: CardDefectDetectionParams ) -> None: """批量处理多个JSON文件""" input_dir = Path(input_dir_str) output_dir = Path(output_dir_str) output_dir.mkdir(parents=True, exist_ok=True) json_files = list(input_dir.rglob("*.json")) if not json_files: fry_algo_print("警告", f"未在{input_dir}中找到JSON文件") return fry_algo_print("信息", f"找到{len(json_files)}个JSON文件") for json_file in json_files: fry_algo_print("组开始", f"处理文件: {json_file.name}") try: json_data = self.load_json_data(json_file) image_stem = json_file.stem image_extensions = ['.jpg', '.jpeg', '.png', '.bmp'] image_path = None for ext in image_extensions: potential_path = json_file.parent / f"{image_stem}{ext}" if potential_path.exists(): image_path = potential_path break image_width = json_data.get("imageWidth", 2048) image_height = json_data.get("imageHeight", 2048) if image_path: image = self.load_or_create_image(image_path, image_width, image_height) else: raise ValueError("警告", f"未找到对应的图像文件") relative_path = json_file.relative_to(input_dir) file_output_dir = output_dir / relative_path.parent / json_file.stem file_output_dir.mkdir(parents=True, exist_ok=True) data_dict = { "image": image, "label": json_data, "output_dir": str(file_output_dir), "answer_json": {} } result = self.process_single_group_entrance(data_dict, algo_params) fry_algo_print("成功", f"文件{json_file.name}处理完成") except Exception as e: fry_algo_print("错误", f"处理文件{json_file.name}时出错: {str(e)}") import traceback traceback.print_exc() fry_algo_print("组结束", f"文件{json_file.name}处理结束") fry_algo_print("成功", "批量处理完成") def _test_card_defect_detection_real(): """测试真实数据""" temp_dir = Path("_test_real").resolve() input_dir = temp_dir / "input" output_dir = temp_dir / "output" try: algo = CardDefectDetectionAlgo() params = CardDefectDetectionParams( # debug_level="detail", debug_level="no", save_intermediate_results=True, save_overlay_images=True, label_name="outer_box", long_edge_corner_length=200, short_edge_corner_length=200, edge_width=50, iteration_rounds=3, ransac_residual_threshold=5.0 ) algo.process_batch_group_entrance( str(input_dir), str(output_dir), params ) fry_algo_print("成功", f"测试完成，结果保存在: {output_dir}") result_files = list(output_dir.rglob("*")) for result_file in result_files: if result_file.is_file(): fry_algo_print("信息", f"生成文件: {result_file}") finally: pass def _test_one_img(): algo = CardDefectDetectionAlgo() image_path = r"C:\Code\ML\Image\Card\_250915_many_capture_img\_250915_1743_reflect_nature_defrct\1_front_coaxial_1_0.jpg" json_path = r"C:\Code\ML\Project\卡片缺陷检测项目组\_250903_1644_边角不直脚本优化\_test_real\input\1_front_coaxial_1_0.json" image = algo.load_or_create_image(image_path) json_data = algo.load_json_data(json_path) temp_out_dir = r"temp_out_dir" data_dict = { "image": image, "label": json_data, "output_dir": str(temp_out_dir), "answer_json": {} } algo_params = CardDefectDetectionParams( # debug_level="detail", debug_level="no", save_intermediate_results=True, save_overlay_images=True, label_name="outer_box", long_edge_corner_length=200, short_edge_corner_length=200, edge_width=50, iteration_rounds=3, ransac_residual_threshold=5.0 ) result = algo.process_single_group_entrance(data_dict, algo_params) print(result) if __name__ == "__main__": # _test_card_defect_detection_real() _test_one_img()