4. img.shape[:]、获取程序的执行的时间、图像的初始化和cv2.bitwise_not()

import cv2 import numpy as np from matplotlib import pyplot as plt from sklearn.cluster import KMeans import math # 纸张尺寸标准 (单位: mm) PAPER_SIZES = { "A4": (210, 297), "B5": (176, 250), "A5": (148, 210) } #使用Harris角点检测算法检测图像中的角点""" def detect_corners_with_harris(image, block_size=2, ksize=3, k=0.04, threshold=0.01): if len(image.shape) == 3: gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) else: gray = image gray = np.float32(gray) dst = cv2.cornerHarris(gray, block_size, ksize, k) dst = cv2.dilate(dst, None) corner_img = image.copy() corners = [] threshold_val = threshold * dst.max() for i in range(dst.shape[0]): for j in range(dst.shape[1]): if dst[i, j] > threshold_val: corners.append((j, i)) cv2.circle(corner_img, (j, i), 5, (0, 0, 255), 2) return corner_img, corners #从Harris角点中找出最可能的纸张四个角点 def find_paper_corners_using_harris(corners, image_shape, top_k=4): height, width = image_shape[:2] top_left = (0, 0) top_right = (width, 0) bottom_right = (width, height) bottom_left = (0, height) candidates = { 'top_left': [], 'top_right': [], 'bottom_right': [], 'bottom_left': [] } for corner in corners: x, y = corner distances = { 'top_left': np.sqrt((x - top_left[0]) ** 2 + (y - top_left[1]) ** 2), 'top_right': np.sqrt((x - top_right[0]) ** 2 + (y - top_right[1]) ** 2), 'bottom_right': np.sqrt((x - bottom_right[0]) ** 2 + (y - bottom_right[1]) ** 2), 'bottom_left': np.sqrt((x - bottom_left[0]) ** 2 + (y - bottom_left[1]) ** 2) } closest_category = min(distances, key=distances.get) candidates[closest_category].append((corner, distances[closest_category])) selected_corners = [] for category in candidates: if candidates[category]: candidates[category].sort(key=lambda x: x[1]) selected_corners.append(candidates[category][0][0]) if len(selected_corners) < 4: return None return order_points(np.array(selected_corners)) #对点进行排序：左上，右上，右下，左下 def order_points(pts): rect = np.zeros((4, 2), dtype="float32") s = pts.sum(axis=1) rect[0] = pts[np.argmin(s)] rect[2] = pts[np.argmax(s)] diff = np.diff(pts, axis=1) rect[1] = pts[np.argmin(diff)] rect[3] = pts[np.argmax(diff)] return rect #检测图像中的纸张轮廓 def detect_paper_contour(image): gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) blurred = cv2.GaussianBlur(gray, (5, 5), 0) edged = cv2.Canny(blurred, 50, 200) contours, _ = cv2.findContours(edged.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) contours = sorted(contours, key=cv2.contourArea, reverse=True)[:5] for contour in contours: peri = cv2.arcLength(contour, True) approx = cv2.approxPolyDP(contour, 0.02 * peri, True) if len(approx) == 4: return order_points(approx.reshape(4, 2)), gray return None, gray #透视变换校正图像 def perspective_transform(image, points): def max_width_height(pts): width1 = np.sqrt(((pts[0][0] - pts[1][0]) ** 2) + ((pts[0][1] - pts[1][1]) ** 2)) width2 = np.sqrt(((pts[2][0] - pts[3][0]) ** 2) + ((pts[2][1] - pts[3][1]) ** 2)) max_width = max(int(width1), int(width2)) height1 = np.sqrt(((pts[0][0] - pts[3][0]) ** 2) + ((pts[0][1] - pts[3][1]) ** 2)) height2 = np.sqrt(((pts[1][0] - pts[2][0]) ** 2) + ((pts[1][1] - pts[2][1]) ** 2)) max_height = max(int(height1), int(height2)) return max_width, max_height points = order_points(points) max_width, max_height = max_width_height(points) dst = np.array([ [0, 0], [max_width - 1, 0], [max_width - 1, max_height - 1], [0, max_height - 1] ], dtype="float32") M = cv2.getPerspectiveTransform(points.astype("float32"), dst) warped = cv2.warpPerspective(image, M, (max_width, max_height)) return warped, M, max_width, max_height, points #创建纸张区域掩码 def create_paper_mask(width, height): mask = np.ones((height, width), dtype=np.uint8) * 255 print(f"创建纸张掩码，尺寸: {mask.shape}, 类型: {mask.dtype}") return mask #预处理图像以进行形状检测 def preprocess_image_for_shape_detection(image, paper_mask=None): # 转换为HSV颜色空间 hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) # 可选：如果有纸张掩码，仅处理纸张内部区域 if paper_mask is not None: # 确保掩码尺寸与图像匹配 if paper_mask.shape[:2] != hsv.shape[:2]: print(f"调整掩码尺寸: {paper_mask.shape} -> {hsv.shape[:2]}") paper_mask = cv2.resize(paper_mask, (hsv.shape[1], hsv.shape[0]), interpolation=cv2.INTER_NEAREST) # 确保掩码是uint8类型 if paper_mask.dtype != np.uint8: print("转换掩码数据类型为uint8") paper_mask = paper_mask.astype(np.uint8) # 应用掩码 try: hsv = cv2.bitwise_and(hsv, hsv, mask=paper_mask) except cv2.error as e: print(f"应用掩码时出错: {e}") print(f"HSV形状: {hsv.shape}, 掩码形状: {paper_mask.shape}") # 高斯模糊去噪 blurred = cv2.GaussianBlur(hsv, (5, 5), 0) return hsv, blurred #使用颜色聚类检测图像中的颜色区域 def detect_color_regions(image, paper_mask=None, n_clusters=6, min_area_ratio=0.002): h, w = image.shape[:2] # 计算最小面积阈值（基于纸张面积比例） min_area = min_area_ratio * w * h print(f"最小面积阈值: {min_area} 像素") # 调整图像大小以提高聚类效率 if h > 500 or w > 500: scale = 500 / max(h, w) resized = cv2.resize(image, None, fx=scale, fy=scale) h_resized, w_resized = resized.shape[:2] print(f"调整图像大小: {h}x{w} -> {h_resized}x{w_resized}, 缩放比例: {scale}") else: resized = image.copy() scale = 1.0 h_resized, w_resized = h, w print(f"图像大小未调整: {h}x{w}") # 预处理图像 hsv, blurred = preprocess_image_for_shape_detection(resized, paper_mask) # 转换为2D数组 pixel_values = hsv.reshape((-1, 3)) pixel_values = np.float32(pixel_values) print(f"像素值数组形状: {pixel_values.shape}") # 排除背景颜色（白色背景：低饱和度、高亮度） not_background = np.logical_and( hsv[:, :, 1] > 30, # 饱和度高（目标颜色） hsv[:, :, 2] < 230 # 亮度不极端（背景很亮） ) if paper_mask is not None: # 调整掩码尺寸以匹配调整后的图像 paper_mask_resized = cv2.resize(paper_mask, (w_resized, h_resized), interpolation=cv2.INTER_NEAREST) not_background = np.logical_and(not_background, paper_mask_resized > 0) pixel_mask = not_background.reshape(-1) valid_pixels = pixel_values[pixel_mask] print(f"有效像素数量: {len(valid_pixels)}") # 执行KMeans聚类 regions = [] if len(valid_pixels) > 0: # 显式设置n_init参数以避免FutureWarning kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10) kmeans.fit(valid_pixels) print(f"KMeans聚类完成，聚类数: {n_clusters}") # 获取聚类标签并转换回原始图像尺寸 labels = kmeans.predict(pixel_values) labels = labels.reshape((h_resized, w_resized)) print(f"标签数组形状: {labels.shape}") # 为每个聚类创建掩码 for cluster_id in range(n_clusters): mask = np.uint8(labels == cluster_id) * 255 # 形态学操作优化掩码 kernel = np.ones((5, 5), np.uint8) mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=2) # 开运算去噪 mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2) # 闭运算连接区域 # 如果调整了图像大小，将掩码调整回原始尺寸 if scale != 1.0: mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST) print(f"调整掩码尺寸: {mask.shape}") # 如果有纸张掩码，仅保留纸张内部区域 if paper_mask is not None: # 确保掩码尺寸匹配 if mask.shape[:2] != paper_mask.shape[:2]: paper_mask_resized = cv2.resize(paper_mask, (mask.shape[1], mask.shape[0]), interpolation=cv2.INTER_NEAREST) else: paper_mask_resized = paper_mask mask = cv2.bitwise_and(mask, paper_mask_resized) # 查找轮廓 contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) print(f"聚类 {cluster_id} 检测到 {len(contours)} 个轮廓") for contour in contours: area = cv2.contourArea(contour) if area > min_area: # 过滤小区域 regions.append(contour) print(f" 保留轮廓，面积: {area} 像素") print(f"总共检测到 {len(regions)} 个颜色区域") return regions #识别物体形状，优化近似参数和判断逻辑 def identify_shape(contour): peri = cv2.arcLength(contour, True) if peri == 0: return "不规则形状" # 更严格的多边形近似（epsilon=0.02*peri） approx = cv2.approxPolyDP(contour, 0.02 * peri, True) num_sides = len(approx) area = cv2.contourArea(contour) circularity = 4 * np.pi * area / (peri ** 2) # 椭圆拟合（增加轮廓点数判断，避免报错） if len(contour) >= 5: try: ellipse = cv2.fitEllipse(contour) (center, axes, orientation) = ellipse major_axis = max(axes) minor_axis = min(axes) eccentricity = np.sqrt(1 - (minor_axis / major_axis) ** 2) if major_axis != 0 else 0 ellipse_area = np.pi * (major_axis / 2) * (minor_axis / 2) ellipse_ratio = area / ellipse_area if ellipse_area > 0 else 0 # 更严格的椭圆判定条件 if 0.8 <= ellipse_ratio <= 1.2 and 0.5 < eccentricity < 0.95: return "椭圆" except: pass # 多边形形状判断（优化长宽比计算） if num_sides == 3: return "三角形" elif num_sides == 4: rect = cv2.minAreaRect(contour) w, h = rect[1] if min(w, h) == 0: # 避免除以0 aspect_ratio = 0 else: aspect_ratio = max(w, h) / min(w, h) if 0.85 <= aspect_ratio <= 1.15: return "正方形" else: return "矩形" elif circularity > 0.85: # 圆形判定（圆度更高） return "圆形" elif 5 <= num_sides <= 10: return f"{num_sides}边形" else: return "不规则形状" #计算物体的几何属性 def calculate_properties(contour, pixel_to_mm=1.0): rect = cv2.minAreaRect(contour) width_px, height_px = rect[1] dimensions_px = sorted([width_px, height_px], reverse=True) length_px, width_px = dimensions_px length_mm = length_px * pixel_to_mm width_mm = width_px * pixel_to_mm area_px = cv2.contourArea(contour) area_mm = area_px * (pixel_to_mm ** 2) M = cv2.moments(contour) if M["m00"] != 0: cx = int(M["m10"] / M["m00"]) cy = int(M["m01"] / M["m00"]) else: cx, cy = 0, 0 return length_mm, width_mm, area_mm, (cx, cy), length_px, width_px, area_px #提取物体主要颜色(HSV) def extract_dominant_color(image, contour): mask = np.zeros(image.shape[:2], np.uint8) cv2.drawContours(mask, [contour], -1, 255, -1) hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) # 提取掩码区域的颜色 masked_pixels = hsv[mask == 255] if len(masked_pixels) == 0: return (0, 0, 0) # 计算主导颜色 (使用中位数更鲁棒) dominant_hsv = np.median(masked_pixels, axis=0) return tuple(map(int, dominant_hsv)) #可视化检测过程的各个步骤 def visualize_detection_steps(image, corrected_image, paper_mask, color_regions, result_img): """可视化检测过程的各个步骤，每个步骤显示在单独的窗口中""" # 显示原始图像 plt.figure(figsize=(10, 8)) plt.title('original image') plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) plt.axis('off') plt.tight_layout() plt.show() # 显示校正后的图像 plt.figure(figsize=(10, 8)) plt.title('corrected image') plt.imshow(cv2.cvtColor(corrected_image, cv2.COLOR_BGR2RGB)) plt.axis('off') plt.tight_layout() plt.show() # 显示颜色聚类区域 color_clusters = corrected_image.copy() colors = [(0, 0, 255), (0, 255, 0), (255, 0, 0), (255, 255, 0), (0, 255, 255), (255, 0, 255), (128, 0, 0), (0, 128, 0)] for i, contour in enumerate(color_regions): color = colors[i % len(colors)] cv2.drawContours(color_clusters, [contour.astype(int)], -1, color, 3) plt.figure(figsize=(10, 8)) plt.title('Color Cluster Area') plt.imshow(cv2.cvtColor(color_clusters, cv2.COLOR_BGR2RGB)) plt.axis('off') plt.tight_layout() plt.show() # 显示二值化结果 gray = cv2.cvtColor(corrected_image, cv2.COLOR_BGR2GRAY) binary = cv2.adaptiveThreshold( gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2 ) if paper_mask is not None: binary = cv2.bitwise_and(binary, binary, mask=paper_mask) plt.figure(figsize=(10, 8)) plt.title('Binary image') plt.imshow(binary, cmap='gray') plt.axis('off') plt.tight_layout() plt.show() # 显示最终识别结果 plt.figure(figsize=(10, 8)) plt.title('Shape recognition results') plt.imshow(cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)) plt.axis('off') plt.tight_layout() plt.show() #主函数 def main(image_path): image = cv2.imread(image_path) if image is None: print(f"错误：无法读取图像 {image_path}") return # 1. 使用Harris角点检测纸张边框 corner_img, corners = detect_corners_with_harris(image) paper_points = find_paper_corners_using_harris(corners, image.shape) # 初始化变量 corrected_image = image.copy() paper_mask = None pixel_to_mm = 1.0 paper_detected = False # 2. 透视变换校正图像 if paper_points is not None: # 执行透视变换 corrected_image, M, max_width, max_height, original_points = perspective_transform(image, paper_points) paper_detected = True print("成功检测并校正纸张轮廓") # 创建纸张掩码（仅处理纸张内部区域） paper_mask = create_paper_mask(max_width, max_height) # 计算物理尺寸转换因子 paper_width_mm, paper_height_mm = PAPER_SIZES["A4"] pixel_to_mm_x = paper_width_mm / max_width pixel_to_mm_y = paper_height_mm / max_height pixel_to_mm = (pixel_to_mm_x + pixel_to_mm_y) / 2.0 else: print("未检测到纸张轮廓 - 使用原始图像进行分析") h, w = image.shape[:2] paper_mask = create_paper_mask(w, h) # 3. 基于颜色聚类和形态学操作检测几何形状 color_regions = detect_color_regions( corrected_image, paper_mask=paper_mask, n_clusters=6, # 调整为6个聚类，匹配目标颜色数量 min_area_ratio=0.002 # 提高最小面积阈值，过滤小区域 ) if not color_regions: print("未检测到颜色区域 - 尝试使用二值化方法检测") # 备用方案：使用二值化方法检测 gray = cv2.cvtColor(corrected_image, cv2.COLOR_BGR2GRAY) binary = cv2.adaptiveThreshold( gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2 ) if paper_mask is not None: binary = cv2.bitwise_and(binary, binary, mask=paper_mask) # 形态学操作 kernel = np.ones((5, 5), np.uint8) binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel, iterations=2) binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, iterations=2) # 查找轮廓 contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) min_area = 0.001 * binary.shape[0] * binary.shape[1] color_regions = [c for c in contours if cv2.contourArea(c) > min_area] if not color_regions: print("未检测到几何图形") return # 初始化结果列表 results = [] result_img = corrected_image.copy() # 处理每个颜色区域 for i, contour in enumerate(color_regions): # 提取主导颜色 hsv_color = extract_dominant_color(corrected_image, contour) rgb_color = cv2.cvtColor(np.uint8([[hsv_color]]), cv2.COLOR_HSV2BGR)[0][0] color_name = f"RGB({rgb_color[2]}, {rgb_color[1]}, {rgb_color[0]})" # 识别形状 shape = identify_shape(contour) # 计算几何属性 length_mm, width_mm, area_mm, center_px, length_px, width_px, area_px = calculate_properties( contour, pixel_to_mm ) # 存储结果 results.append({ "id": i + 1, "shape": shape, "length_mm": length_mm, "width_mm": width_mm, "area_mm": area_mm, "color": color_name, "hsv_color": hsv_color, "center_px": center_px, "length_px": length_px, "width_px": width_px, "area_px": area_px }) # 在图像上标注物体 color = (int(rgb_color[2]), int(rgb_color[1]), int(rgb_color[0])) cv2.drawContours(result_img, [contour], -1, color, 3) cv2.putText(result_img, f"{i + 1}({shape})", (int(center_px[0]), int(center_px[1])), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2) # 可视化检测步骤 visualize_detection_steps(image, corrected_image, paper_mask, color_regions, result_img) # 打印分析结果 print("\n===== 图像分析结果 =====") print(f"检测到 {len(results)} 个几何图形") print("-" * 70) for obj in results: print(f"图形 {obj['id']}:") print(f" 形状: {obj['shape']}") print(f" 长度: {obj['length_mm']:.2f} mm ({obj['length_px']:.1f} 像素)") print(f" 宽度: {obj['width_mm']:.2f} mm ({obj['width_px']:.1f} 像素)") print(f" 面积: {obj['area_mm']:.2f} mm&sup2; ({obj['area_px']:.1f} 像素&sup2;)") print(f" 颜色: {obj['color']} (HSV: {obj['hsv_color']})") print(f" 中心坐标: ({obj['center_px'][0]}, {obj['center_px'][1]})") print("-" * 70) if __name__ == "__main__": image_path = "ALL2.jpg" # 替换为实际图像路径 main(image_path) 上述代码创建纸张区域掩码后要对图像进行处理再进一步识别图像内部的不同几何形状，上述代码所采用的是提取颜色进行识别，但是对于图片中的颜色可能是分配不均匀的，所以会导致识别效果不是很好，我想换成边缘检测的方法识别。先对图片进行二值化操作再进行边缘检测。请根据我的描述修改上述代码。

import cv2 import numpy as np from numpy.linalg import norm import sys import os import json SZ = 20 #训练图片长宽 MAX_WIDTH = 1000 #原始图片最大宽度 Min_Area = 2000 #车牌区域允许最大面积 PROVINCE_START = 1000 #读取图片文件 def imreadex(filename): return cv2.imdecode(np.fromfile(filename, dtype=np.uint8), cv2.IMREAD_COLOR) def point_limit(point): if point[0] < 0: point[0] = 0 if point[1] < 0: point[1] = 0 #根据设定的阈值和图片直方图，找出波峰，用于分隔字符 def find_waves(threshold, histogram): up_point = -1#上升点 is_peak = False if histogram[0] > threshold: up_point = 0 is_peak = True wave_peaks = [] for i,x in enumerate(histogram): if is_peak and x < threshold: if i - up_point > 2: is_peak = False wave_peaks.append((up_point, i)) elif not is_peak and x >= threshold: is_peak = True up_point = i if is_peak and up_point != -1 and i - up_point > 4: wave_peaks.append((up_point, i)) return wave_peaks #根据找出的波峰，分隔图片，从而得到逐个字符图片 def seperate_card(img, waves): part_cards = [] for wave in waves: part_cards.append(img[:, wave[0]:wave[1]]) return part_cards #来自opencv的sample，用于svm训练 def deskew(img): m = cv2.moments(img) if abs(m['mu02']) < 1e-2: return img.copy() skew = m['mu11']/m['mu02'] M = np.float32([[1, skew, -0.5*SZ*skew], [0, 1, 0]]) img = cv2.warpAffine(img, M, (SZ, SZ), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LINEAR) return img #来自opencv的sample，用于svm训练 def preprocess_hog(digits): samples = [] for img in digits: gx = cv2.Sobel(img, cv2.CV_32F, 1, 0) gy = cv2.Sobel(img, cv2.CV_32F, 0, 1) mag, ang = cv2.cartToPolar(gx, gy) bin_n = 16 bin = np.int32(bin_n*ang/(2*np.pi)) bin_cells = bin[:10,:10], bin[10:,:10], bin[:10,10:], bin[10:,10:] mag_cells = mag[:10,:10], mag[10:,:10], mag[:10,10:], mag[10:,10:] hists = [np.bincount(b.ravel(), m.ravel(), bin_n) for b, m in zip(bin_cells, mag_cells)] hist = np.hstack(hists) # transform to Hellinger kernel eps = 1e-7 hist /= hist.sum() + eps hist = np.sqrt(hist) hist /= norm(hist) + eps samples.append(hist) return np.float32(samples) #不能保证包括所有省份 provinces = [ "zh_cuan", "川", "zh_e", "鄂", "zh_gan", "赣", "zh_gan1", "甘", "zh_gui", "贵", "zh_gui1", "桂", "zh_hei", "黑", "zh_hu", "沪", "zh_ji", "冀", "zh_jin", "津", "zh_jing", "京", "zh_jl", "吉", "zh_liao", "辽", "zh_lu", "鲁", "zh_meng", "蒙", "zh_min", "闽", "zh_ning", "宁", "zh_qing", "靑", "zh_qiong", "琼", "zh_shan", "陕", "zh_su", "苏", "zh_sx", "晋", "zh_wan", "皖", "zh_xiang", "湘", "zh_xin", "新", "zh_yu", "豫", "zh_yu1", "渝", "zh_yue", "粤", "zh_yun", "云", "zh_zang", "藏", "zh_zhe", "浙" ] class StatModel(object): def load(self, fn): self.model = self.model.load(fn) def save(self, fn): self.model.save(fn) class SVM(StatModel): def __init__(self, C = 1, gamma = 0.5): self.model = cv2.ml.SVM_create() self.model.setGamma(gamma) self.model.setC(C) self.model.setKernel(cv2.ml.SVM_RBF) self.model.setType(cv2.ml.SVM_C_SVC) #训练svm def train(self, samples, responses): self.model.train(samples, cv2.ml.ROW_SAMPLE, responses) #字符识别 def predict(self, samples): r = self.model.predict(samples) return r[1].ravel() class CardPredictor: def __init__(self): #车牌识别的部分参数保存在js中，便于根据图片分辨率做调整 f = open('config.js') j = json.load(f) for c in j["config"]: if c["open"]: self.cfg = c.copy() break else: raise RuntimeError('没有设置有效配置参数') def __del__(self): self.save_traindata() def train_svm(self): #识别英文字母和数字 self.model = SVM(C=1, gamma=0.5) #识别中文 self.modelchinese = SVM(C=1, gamma=0.5) if os.path.exists("svm.dat"): self.model.load("svm.dat") else: chars_train = [] chars_label = [] for root, dirs, files in os.walk("train\\chars2"): if len(os.path.basename(root)) > 1: continue root_int = ord(os.path.basename(root)) for filename in files: filepath = os.path.join(root,filename) digit_img = cv2.imread(filepath) digit_img = cv2.cvtColor(digit_img, cv2.COLOR_BGR2GRAY) chars_train.append(digit_img) #chars_label.append(1) chars_label.append(root_int) chars_train = list(map(deskew, chars_train)) chars_train = preprocess_hog(chars_train) #chars_train = chars_train.reshape(-1, 20, 20).astype(np.float32) chars_label = np.array(chars_label) self.model.train(chars_train, chars_label) if os.path.exists("svmchinese.dat"): self.modelchinese.load("svmchinese.dat") else: chars_train = [] chars_label = [] for root, dirs, files in os.walk("train\\charsChinese"): if not os.path.basename(root).startswith("zh_"): continue pinyin = os.path.basename(root) index = provinces.index(pinyin) + PROVINCE_START + 1 #1是拼音对应的汉字 for filename in files: filepath = os.path.join(root,filename) digit_img = cv2.imread(filepath) digit_img = cv2.cvtColor(digit_img, cv2.COLOR_BGR2GRAY) chars_train.append(digit_img) #chars_label.append(1) chars_label.append(index) chars_train = list(map(deskew, chars_train)) chars_train = preprocess_hog(chars_train) #chars_train = chars_train.reshape(-1, 20, 20).astype(np.float32) chars_label = np.array(chars_label) print(chars_train.shape) self.modelchinese.train(chars_train, chars_label) def save_traindata(self): if not os.path.exists("svm.dat"): self.model.save("svm.dat") if not os.path.exists("svmchinese.dat"): self.modelchinese.save("svmchinese.dat") def accurate_place(self, card_img_hsv, limit1, limit2, color): row_num, col_num = card_img_hsv.shape[:2] xl = col_num xr = 0 yh = 0 yl = row_num #col_num_limit = self.cfg["col_num_limit"] row_num_limit = self.cfg["row_num_limit"] col_num_limit = col_num * 0.8 if color != "green" else col_num * 0.5#绿色有渐变 for i in range(row_num): count = 0 for j in range(col_num): H = card_img_hsv.item(i, j, 0) S = card_img_hsv.item(i, j, 1) V = card_img_hsv.item(i, j, 2) if limit1 < H <= limit2 and 34 < S and 46 < V: count += 1 if count > col_num_limit: if yl > i: yl = i if yh < i: yh = i for j in range(col_num): count = 0 for i in range(row_num): H = card_img_hsv.item(i, j, 0) S = card_img_hsv.item(i, j, 1) V = card_img_hsv.item(i, j, 2) if limit1 < H <= limit2 and 34 < S and 46 < V: count += 1 if count > row_num - row_num_limit: if xl > j: xl = j if xr < j: xr = j return xl, xr, yh, yl def predict(self, car_pic, resize_rate=1): if type(car_pic) == type(""): img = imreadex(car_pic) else: img = car_pic pic_hight, pic_width = img.shape[:2] if pic_width > MAX_WIDTH: pic_rate = MAX_WIDTH / pic_width img = cv2.resize(img, (MAX_WIDTH, int(pic_hight*pic_rate)), interpolation=cv2.INTER_LANCZOS4) if resize_rate != 1: img = cv2.resize(img, (int(pic_width*resize_rate), int(pic_hight*resize_rate)), interpolation=cv2.INTER_LANCZOS4) pic_hight, pic_width = img.shape[:2] print("h,w:", pic_hight, pic_width) blur = self.cfg["blur"] #高斯去噪 if blur > 0: img = cv2.GaussianBlur(img, (blur, blur), 0)#图片分辨率调整 oldimg = img img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) #equ = cv2.equalizeHist(img) #img = np.hstack((img, equ)) #去掉图像中不会是车牌的区域 kernel = np.ones((20, 20), np.uint8) img_opening = cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel) img_opening = cv2.addWeighted(img, 1, img_opening, -1, 0); #找到图像边缘 ret, img_thresh = cv2.threshold(img_opening, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) img_edge = cv2.Canny(img_thresh, 100, 200) #使用开运算和闭运算让图像边缘成为一个整体 kernel = np.ones((self.cfg["morphologyr"], self.cfg["morphologyc"]), np.uint8) img_edge1 = cv2.morphologyEx(img_edge, cv2.MORPH_CLOSE, kernel) img_edge2 = cv2.morphologyEx(img_edge1, cv2.MORPH_OPEN, kernel) #查找图像边缘整体形成的矩形区域，可能有很多，车牌就在其中一个矩形区域中 try: contours, hierarchy = cv2.findContours(img_edge2, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) except ValueError: image, contours, hierarchy = cv2.findContours(img_edge2, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) contours = [cnt for cnt in contours if cv2.contourArea(cnt) > Min_Area] print('len(contours)', len(contours)) #一一排除不是车牌的矩形区域 car_contours = [] for cnt in contours: rect = cv2.minAreaRect(cnt) area_width, area_height = rect[1] if area_width < area_height: area_width, area_height = area_height, area_width wh_ratio = area_width / area_height #print(wh_ratio) #要求矩形区域长宽比在2到5.5之间，2到5.5是车牌的长宽比，其余的矩形排除 if wh_ratio > 2 and wh_ratio < 5.5: car_contours.append(rect) box = cv2.boxPoints(rect) box = np.intp(box) #oldimg = cv2.drawContours(oldimg, [box], 0, (0, 0, 255), 2) #cv2.imshow("edge4", oldimg) #cv2.waitKey(0) print(len(car_contours)) print("精确定位") card_imgs = [] #矩形区域可能是倾斜的矩形，需要矫正，以便使用颜色定位 for rect in car_contours: if rect[2] > -1 and rect[2] < 1:#创造角度，使得左、高、右、低拿到正确的值 angle = 1 else: angle = rect[2] rect = (rect[0], (rect[1][0]+5, rect[1][1]+5), angle)#扩大范围，避免车牌边缘被排除 box = cv2.boxPoints(rect) heigth_point = right_point = [0, 0] left_point = low_point = [pic_width, pic_hight] for point in box: if left_point[0] > point[0]: left_point = point if low_point[1] > point[1]: low_point = point if heigth_point[1] < point[1]: heigth_point = point if right_point[0] < point[0]: right_point = point if left_point[1] <= right_point[1]:#正角度 new_right_point = [right_point[0], heigth_point[1]] pts2 = np.float32([left_point, heigth_point, new_right_point])#字符只是高度需要改变 pts1 = np.float32([left_point, heigth_point, right_point]) M = cv2.getAffineTransform(pts1, pts2) dst = cv2.warpAffine(oldimg, M, (pic_width, pic_hight)) point_limit(new_right_point) point_limit(heigth_point) point_limit(left_point) card_img = dst[int(left_point[1]):int(heigth_point[1]), int(left_point[0]):int(new_right_point[0])] card_imgs.append(card_img) #cv2.imshow("card", card_img) #cv2.waitKey(0) elif left_point[1] > right_point[1]:#负角度 new_left_point = [left_point[0], heigth_point[1]] pts2 = np.float32([new_left_point, heigth_point, right_point])#字符只是高度需要改变 pts1 = np.float32([left_point, heigth_point, right_point]) M = cv2.getAffineTransform(pts1, pts2) dst = cv2.warpAffine(oldimg, M, (pic_width, pic_hight)) point_limit(right_point) point_limit(heigth_point) point_limit(new_left_point) card_img = dst[int(right_point[1]):int(heigth_point[1]), int(new_left_point[0]):int(right_point[0])] card_imgs.append(card_img) #cv2.imshow("card", card_img) #cv2.waitKey(0) #开始使用颜色定位，排除不是车牌的矩形，目前只识别蓝、绿、黄车牌 colors = [] for card_index,card_img in enumerate(card_imgs): green = yello = blue = black = white = 0 card_img_hsv = cv2.cvtColor(card_img, cv2.COLOR_BGR2HSV) #有转换失败的可能，原因来自于上面矫正矩形出错 if card_img_hsv is None: continue row_num, col_num= card_img_hsv.shape[:2] card_img_count = row_num * col_num for i in range(row_num): for j in range(col_num): H = card_img_hsv.item(i, j, 0) S = card_img_hsv.item(i, j, 1) V = card_img_hsv.item(i, j, 2) if 11 < H <= 34 and S > 34:#图片分辨率调整 yello += 1 elif 35 < H <= 99 and S > 34:#图片分辨率调整 green += 1 elif 99 < H <= 124 and S > 34:#图片分辨率调整 blue += 1 if 0 < H <180 and 0 < S < 255 and 0 < V < 46: black += 1 elif 0 < H <180 and 0 < S < 43 and 221 < V < 225: white += 1 color = "no" limit1 = limit2 = 0 if yello*2 >= card_img_count: color = "yello" limit1 = 11 limit2 = 34#有的图片有色偏偏绿 elif green*2 >= card_img_count: color = "green" limit1 = 35 limit2 = 99 elif blue*2 >= card_img_count: color = "blue" limit1 = 100 limit2 = 124#有的图片有色偏偏紫 elif black + white >= card_img_count*0.7:#TODO color = "bw" print(color) colors.append(color) print(blue, green, yello, black, white, card_img_count) #cv2.imshow("color", card_img) #cv2.waitKey(0) if limit1 == 0: continue #以上为确定车牌颜色 #以下为根据车牌颜色再定位，缩小边缘非车牌边界 xl, xr, yh, yl = self.accurate_place(card_img_hsv, limit1, limit2, color) if yl == yh and xl == xr: continue need_accurate = False if yl >= yh: yl = 0 yh = row_num need_accurate = True if xl >= xr: xl = 0 xr = col_num need_accurate = True card_imgs[card_index] = card_img[yl:yh, xl:xr] if color != "green" or yl < (yh-yl)//4 else card_img[yl-(yh-yl)//4:yh, xl:xr] if need_accurate:#可能x或y方向未缩小，需要再试一次 card_img = card_imgs[card_index] card_img_hsv = cv2.cvtColor(card_img, cv2.COLOR_BGR2HSV) xl, xr, yh, yl = self.accurate_place(card_img_hsv, limit1, limit2, color) if yl == yh and xl == xr: continue if yl >= yh: yl = 0 yh = row_num if xl >= xr: xl = 0 xr = col_num card_imgs[card_index] = card_img[yl:yh, xl:xr] if color != "green" or yl < (yh-yl)//4 else card_img[yl-(yh-yl)//4:yh, xl:xr] #以上为车牌定位 #以下为识别车牌中的字符 predict_result = [] roi = None card_color = None for i, color in enumerate(colors): if color in ("blue", "yello", "green"): card_img = card_imgs[i] gray_img = cv2.cvtColor(card_img, cv2.COLOR_BGR2GRAY) #黄、绿车牌字符比背景暗、与蓝车牌刚好相反，所以黄、绿车牌需要反向 if color == "green" or color == "yello": gray_img = cv2.bitwise_not(gray_img) ret, gray_img = cv2.threshold(gray_img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) #查找水平直方图波峰 x_histogram = np.sum(gray_img, axis=1) x_min = np.min(x_histogram) x_average = np.sum(x_histogram)/x_histogram.shape[0] x_threshold = (x_min + x_average)/2 wave_peaks = find_waves(x_threshold, x_histogram) if len(wave_peaks) == 0: print("peak less 0:") continue #认为水平方向，最大的波峰为车牌区域 wave = max(wave_peaks, key=lambda x:x[1]-x[0]) gray_img = gray_img[wave[0]:wave[1]] #查找垂直直方图波峰 row_num, col_num= gray_img.shape[:2] #去掉车牌上下边缘1个像素，避免白边影响阈值判断 gray_img = gray_img[1:row_num-1] y_histogram = np.sum(gray_img, axis=0) y_min = np.min(y_histogram) y_average = np.sum(y_histogram)/y_histogram.shape[0] y_threshold = (y_min + y_average)/5#U和0要求阈值偏小，否则U和0会被分成两半 wave_peaks = find_waves(y_threshold, y_histogram) #for wave in wave_peaks: # cv2.line(card_img, pt1=(wave[0], 5), pt2=(wave[1], 5), color=(0, 0, 255), thickness=2) #车牌字符数应大于6 if len(wave_peaks) <= 6: print("peak less 1:", len(wave_peaks)) continue wave = max(wave_peaks, key=lambda x:x[1]-x[0]) max_wave_dis = wave[1] - wave[0] #判断是否是左侧车牌边缘 if wave_peaks[0][1] - wave_peaks[0][0] < max_wave_dis/3 and wave_peaks[0][0] == 0: wave_peaks.pop(0) #组合分离汉字 cur_dis = 0 for i,wave in enumerate(wave_peaks): if wave[1] - wave[0] + cur_dis > max_wave_dis * 0.6: break else: cur_dis += wave[1] - wave[0] if i > 0: wave = (wave_peaks[0][0], wave_peaks[i][1]) wave_peaks = wave_peaks[i+1:] wave_peaks.insert(0, wave) #去除车牌上的分隔点 point = wave_peaks[2] if point[1] - point[0] < max_wave_dis/3: point_img = gray_img[:,point[0]:point[1]] if np.mean(point_img) < 255/5: wave_peaks.pop(2) if len(wave_peaks) <= 6: print("peak less 2:", len(wave_peaks)) continue part_cards = seperate_card(gray_img, wave_peaks) for i, part_card in enumerate(part_cards): #可能是固定车牌的铆钉 if np.mean(part_card) < 255/5: print("a point") continue part_card_old = part_card #w = abs(part_card.shape[1] - SZ)//2 w = part_card.shape[1] // 3 part_card = cv2.copyMakeBorder(part_card, 0, 0, w, w, cv2.BORDER_CONSTANT, value = [0,0,0]) part_card = cv2.resize(part_card, (SZ, SZ), interpolation=cv2.INTER_AREA) #cv2.imshow("part", part_card_old) #cv2.waitKey(0) #cv2.imwrite("u.jpg", part_card) #part_card = deskew(part_card) part_card = preprocess_hog([part_card]) if i == 0: resp = self.modelchinese.predict(part_card) charactor = provinces[int(resp[0]) - PROVINCE_START] else: resp = self.model.predict(part_card) # charactor = chr(resp[0]) charactor = chr(int(resp[0])) # 显式转换为int #判断最后一个数是否是车牌边缘，假设车牌边缘被认为是1 if charactor == "1" and i == len(part_cards)-1: if part_card_old.shape[0]/part_card_old.shape[1] >= 8:#1太细，认为是边缘 print(part_card_old.shape) continue predict_result.append(charactor) roi = card_img card_color = color break return predict_result, roi, card_color#识别到的字符、定位的车牌图像、车牌颜色 #实现单图多车牌识别 def separate_card(self, gray_img, wave_peaks): """字符分割函数""" part_cards = [] for wave in wave_peaks: part_card = gray_img[:, wave[0]:wave[1]] if part_card.shape[1] > 0: part_cards.append(part_card) return part_cards def predict_multi(self, car_pic, resize_rate=1): if type(car_pic) == type(""): img = imreadex(car_pic) else: img = car_pic pic_hight, pic_width = img.shape[:2] if pic_width > MAX_WIDTH: pic_rate = MAX_WIDTH / pic_width img = cv2.resize(img, (MAX_WIDTH, int(pic_hight * pic_rate)), interpolation=cv2.INTER_LANCZOS4) if resize_rate != 1: img = cv2.resize(img, (int(pic_width * resize_rate), int(pic_hight * resize_rate)), interpolation=cv2.INTER_LANCZOS4) # 初始化参数 blur = self.cfg["blur"] if blur > 0: img = cv2.GaussianBlur(img, (blur, blur), 0) img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) kernel = np.ones((20, 20), np.uint8) img_opening = cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel) img_opening = cv2.addWeighted(img, 1, img_opening, -1, 0) ret, img_thresh = cv2.threshold(img_opening, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) img_edge = cv2.Canny(img_thresh, 100, 200) kernel = np.ones((self.cfg["morphologyr"], self.cfg["morphologyc"]), np.uint8) img_edge1 = cv2.morphologyEx(img_edge, cv2.MORPH_CLOSE, kernel) img_edge2 = cv2.morphologyEx(img_edge1, cv2.MORPH_OPEN, kernel) contours, hierarchy = cv2.findContours(img_edge2, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) contours = [cnt for cnt in contours if cv2.contourArea(cnt) > Min_Area] car_contours = [] for cnt in contours: rect = cv2.minAreaRect(cnt) area_width, area_height = rect[1] if area_width < area_height: area_width, area_height = area_height, area_width wh_ratio = area_width / area_height if wh_ratio > 2 and wh_ratio < 5.5: car_contours.append(rect) results = [] for rect in car_contours: # 矫正车牌区域 if rect[2] > -1 and rect[2] < 1: angle = 1 else: angle = rect[2] rect = (rect[0], (rect[1][0] + 5, rect[1][1] + 5), angle) box = cv2.boxPoints(rect) box = np.intp(box) heigth_point = right_point = [0, 0] left_point = low_point = [pic_width, pic_hight] for point in box: if left_point[0] > point[0]: left_point = point if low_point[1] > point[1]: low_point = point if heigth_point[1] < point[1]: heigth_point = point if right_point[0] < point[0]: right_point = point if left_point[1] <= right_point[1]: # 正角度 new_right_point = [right_point[0], heigth_point[1]] pts2 = np.float32([left_point, heigth_point, new_right_point]) pts1 = np.float32([left_point, heigth_point, right_point]) M = cv2.getAffineTransform(pts1, pts2) dst = cv2.warpAffine(img, M, (pic_width, pic_hight)) point_limit(new_right_point) point_limit(heigth_point) point_limit(left_point) card_img = dst[int(left_point[1]):int(heigth_point[1]), int(left_point[0]):int(new_right_point[0])] elif left_point[1] > right_point[1]: # 负角度 new_left_point = [left_point[0], heigth_point[1]] pts2 = np.float32([new_left_point, heigth_point, right_point]) pts1 = np.float32([left_point, heigth_point, right_point]) M = cv2.getAffineTransform(pts1, pts2) dst = cv2.warpAffine(img, M, (pic_width, pic_hight)) point_limit(right_point) point_limit(heigth_point) point_limit(new_left_point) card_img = dst[int(right_point[1]):int(heigth_point[1]), int(new_left_point[0]):int(right_point[0])] # 颜色识别（蓝/绿/黄） if len(card_img.shape) == 2: # 灰度图检测 card_img = cv2.cvtColor(card_img, cv2.COLOR_GRAY2BGR) card_img_hsv = cv2.cvtColor(card_img, cv2.COLOR_BGR2HSV) # 必须的转换步骤 color = self.detect_color(card_img_hsv) # 字符分割与识别 gray_img = cv2.cvtColor(card_img, cv2.COLOR_BGR2GRAY) if color in ["green", "yello"]: gray_img = cv2.bitwise_not(gray_img) ret, gray_img = cv2.threshold(gray_img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) # 水平投影分割 horizontal_sum = np.sum(gray_img, axis=1) wave_peaks = self.find_peaks(horizontal_sum) if len(wave_peaks) > 0: gray_img = gray_img[wave_peaks[0][0]:wave_peaks[-1][1], :] # 垂直投影分割 vertical_sum = np.sum(gray_img, axis=0) char_peaks = self.find_peaks(vertical_sum) # 字符识别 plate_chars = [] for i, (start, end) in enumerate(char_peaks): char_img = gray_img[:, start:end] char_img = cv2.resize(char_img, (20, 20)) if i == 0: # 第一个字符为中文 char_img = preprocess_hog([char_img]) resp = self.modelchinese.predict(char_img) plate_chars.append(provinces[int(resp[0]) - PROVINCE_START]) else: # 其他字符为字母/数字 char_img = preprocess_hog([char_img]) resp = self.model.predict(char_img) plate_chars.append(chr(int(resp[0]))) results.append({ "plate": "".join(plate_chars), "color": color, "roi": card_img }) return results def detect_color(self, hsv_img): """识别车牌颜色""" # 颜色阈值定义 blue_lower = np.array([100, 50, 50]) blue_upper = np.array([124, 255, 255]) green_lower = np.array([35, 50, 50]) green_upper = np.array([99, 255, 255]) yellow_lower = np.array([11, 50, 50]) yellow_upper = np.array([34, 255, 255]) # 计算各颜色像素占比 blue_mask = cv2.inRange(hsv_img, blue_lower, blue_upper) green_mask = cv2.inRange(hsv_img, green_lower, green_upper) yellow_mask = cv2.inRange(hsv_img, yellow_lower, yellow_upper) total_pixels = hsv_img.shape[0] * hsv_img.shape[1] blue_ratio = np.count_nonzero(blue_mask) / total_pixels green_ratio = np.count_nonzero(green_mask) / total_pixels yellow_ratio = np.count_nonzero(yellow_mask) / total_pixels # 确定主要颜色 if blue_ratio > 0.3: return "blue" elif green_ratio > 0.3: return "green" elif yellow_ratio > 0.3: return "yello" else: return "unknown" def find_peaks(self, histogram, min_width=2): """寻找直方图波峰""" peaks = [] start = 0 rising = False for i, value in enumerate(histogram): if value > 0 and not rising: start = i rising = True elif value == 0 and rising: if i - start >= min_width: peaks.append((start, i)) rising = False if rising and len(histogram) - start >= min_width: peaks.append((start, len(histogram))) return peaks def process_video(self, source=0, output_path=None, show=True): cap = cv2.VideoCapture(source) if not cap.isOpened(): print("无法打开摄像头") return while True: ret, frame = cap.read() if not ret: print("无法接收帧 (stream end?). Exiting ...") break r, roi, color = self.predict(frame) if roi is not None and show: cv2.imshow("识别的车牌", roi) if cv2.waitKey(1) & 0xFF == 27: # 按ESC键退出 break cap.release() cv2.destroyAllWindows() # if __name__ == '__main__': # c = CardPredictor() # c.train_svm() # r, roi, color = c.predict("10.jpg") # print(r) if __name__ == '__main__': c = CardPredictor() c.train_svm() # 用户交互菜单 while True: print("\n===== 车牌识别系统 =====") print("1. 单图片单车牌识别") print("2. 单图片多车牌识别") print("3. 摄像头实时识别") print("4. 视频文件识别") print("5. 退出") choice = input("请选择功能 (1-5): ") if choice == "1": # 单图片单车牌识别（原功能） img_path = input("请输入图片路径: ") r, roi, color = c.predict(img_path) print(f"识别结果: {''.join(r)}, 车牌颜色: {color}") # 显示识别的车牌 if roi is not None: cv2.imshow("识别的车牌", roi) cv2.waitKey(0) cv2.destroyAllWindows() elif choice == "2": # 单图片多车牌识别 img_path = input("请输入图片路径: ") results = c.predict_multi(img_path) if results: print(f"共识别到 {len(results)} 个车牌:") for i, result in enumerate(results, 1): try: # 安全获取字段并设置默认值 plate_num = ''.join(result.get("plate_number", ["识别失败"])) plate_color = result.get("plate_color", "未知颜色") roi_img = result.get("roi", None) print(f"车牌 {i}: {plate_num}, 颜色: {plate_color}") # 仅当存在ROI图像时显示 if roi_img is not None: cv2.imshow(f"车牌 {i}", roi_img) cv2.waitKey(1) # 实时刷新显示 else: print(f"警告: 车牌 {i} 无ROI图像") except Exception as e: print(f"处理车牌 {i} 时出错: {str(e)}") continue cv2.waitKey(0) cv2.destroyAllWindows() else: print("未识别到车牌") elif choice == "3": # 摄像头实时识别 print("按ESC键退出...") c.process_video(0, show=True) elif choice == "4": # 视频文件识别 video_path = input("请输入视频路径: ") output_path = input("请输入输出路径 (留空则不保存): ") output_path = output_path if output_path.strip() else None print("按ESC键退出...") c.process_video(video_path, output_path) elif choice == "5": print("程序已退出") break else: print("无效选择，请重新输入") 基于这个代码修改，让它可以识别出车牌号和车牌颜色以及有几个车牌

修正这段代码，import cv2 import numpy as np from sklearn.cluster import KMeans import matplotlib.pyplot as plt def identify_background_and_generate_patch(image_path, sample_size=10, expansion_factor=5, num_colors=3, tolerance=10): """ 识别图像背景并生成扩展填充样本 参数: image_path: 输入图像路径 sample_size: 背景采样区域大小（像素） expansion_factor: 背景样本扩展倍数 num_colors: 背景颜色聚类数量 tolerance: 背景颜色容差范围 """ # 1. 读取图像 img = cv2.imread(image_path) if img is None: raise FileNotFoundError(f"无法加载图像: {image_path}") img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) height, width, _ = img_rgb.shape # 2. 边缘采样获取初始背景样本 edge_samples = [] # 上边缘 edge_samples.append(img_rgb[0:sample_size, 0:width]) # 下边缘 edge_samples.append(img_rgb[height - sample_size:height, 0:width]) # 左边缘 edge_samples.append(img_rgb[0:height, 0:sample_size]) # 右边缘 edge_samples.append(img_rgb[0:height, width - sample_size:width]) background_samples = np.vstack(edge_samples) # 3. 使用K-means聚类识别主要背景色 pixels = background_samples.reshape(-1, 3).astype(np.float32) kmeans = KMeans(n_clusters=num_colors, random_state=0).fit(pixels) dominant_colors = kmeans.cluster_centers_.astype(np.uint8) # 4. 创建背景掩膜 background_mask = np.zeros((height, width), dtype=np.uint8) for color in dominant_colors: lower_bound = np.clip(color - tolerance, 0, 255) upper_bound = np.clip(color + tolerance, 0, 255) color_mask = cv2.inRange(img_rgb, lower_bound, upper_bound) background_mask = cv2.bitwise_or(background_mask, color_mask) # 5. 形态学操作优化背景掩膜 kernel = np.ones((5, 5), np.uint8) background_mask = cv2.morphologyEx(background_mask, cv2.MORPH_CLOSE, kernel, iterations=2) background_mask = cv2.morphologyEx(background_mask, cv2.MORPH_OPEN, kernel, iterations=1) # 6. 提取背景区域轮廓 contours, _ = cv2.findContours( background_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE ) # 7. 轮廓扩张生成大样本 expanded_mask = np.zeros_like(background_mask) for contour in contours: # 计算轮廓边界框 x, y, w, h = cv2.boundingRect(contour) # 扩展边界框 expand_x = int(w * (expansion_factor - 1) / 2) expand_y = int(h * (expansion_factor - 1) / 2) new_x = max(0, x - expand_x) new_y = max(0, y - expand_y) new_w = min(width - new_x, w * expansion_factor) new_h = min(height - new_y, h * expansion_factor) # 提取背景样本 sample = img_rgb[new_y:new_y + new_h, new_x:new_x + new_w] # 添加到扩展掩膜 expanded_mask[new_y:new_y + new_h, new_x:new_x + new_w] = 255 # 8. 生成最终背景样本图像 background_patch = cv2.bitwise_and( img_rgb, img_rgb, mask=expanded_mask ) # 9. 可视化结果 plt.figure(figsize=(15, 10)) plt.subplot(231) plt.imshow(img_rgb) plt.title('原始图像') plt.axis('off') plt.subplot(232) plt.imshow(background_mask, cmap='gray') plt.title('背景掩膜') plt.axis('off') plt.subplot(233) plt.imshow(expanded_mask, cmap='gray') plt.title('扩展掩膜') plt.axis('off') plt.subplot(234) # 显示聚类颜色 color_swatches = np.zeros((50, 300, 3), dtype=np.uint8) for i, color in enumerate(dominant_colors): color_swatches[:, i * 100:(i + 1) * 100] = color plt.imshow(color_swatches) plt.title(f'主要背景色 (k={num_colors})') plt.axis('off') plt.subplot(235) plt.imshow(background_samples) plt.title('边缘背景样本') plt.axis('off') plt.subplot(236) plt.imshow(background_patch) plt.title('扩展背景样本') plt.axis('off') plt.tight_layout() plt.savefig('background_analysis.jpg') return background_patch, expanded_mask, dominant_colors # 使用示例 if __name__ == "__main__": # 生成背景样本 background_patch, mask, colors = identify_background_and_generate_patch( "zdk.jpg", sample_size=20, expansion_factor=3, num_colors=3, tolerance=15 ) # 保存结果 cv2.imwrite("background_patch.jpg", cv2.cvtColor(background_patch, cv2.COLOR_RGB2BGR)) cv2.imwrite("background_mask.jpg", mask) print(f"生成背景样本尺寸: {background_patch.shape[1]}x{background_patch.shape[0]}") print(f"识别到的主要背景色: {colors.tolist()}") C:\Users\86152\爬虫学习文件\fastApiProjectone\.venv\Scripts\python.exe H:\img_ocr_server\img_util\input\背景去色扩张.py Traceback (most recent call last): File "H:\img_ocr_server\img_util\input\背景去色扩张.py", line 138, in <module> background_patch, mask, colors = identify_background_and_generate_patch( File "H:\img_ocr_server\img_util\input\背景去色扩张.py", line 38, in identify_background_and_generate_patch background_samples = np.vstack(edge_samples) File "<__array_function__ internals>", line 200, in vstack File "C:\Users\86152\爬虫学习文件\fastApiProjectone\.venv\lib\site-packages\numpy\core\shape_base.py", line 296, in vstack return _nx.concatenate(arrs, 0, dtype=dtype, casting=casting) File "<__array_function__ internals>", line 200, in concatenate ValueError: all the input array dimensions except for the concatenation axis must match exactly, but along dimension 1, the array at index 0 has size 2346 and the array at index 2 has size 20 Process finished with exit code 1