main.py

# 导入必要的库
import cv2  # OpenCV库用于图像处理和计算机视觉任务
import numpy as np  # NumPy库用于高效的数值计算和数组操作
from skimage.measure import (
    label,
    regionprops,
)  # scikit-image库用于图像区域分析和特征提取
import os  # 操作系统接口，用于文件和目录操作
import time  # 时间处理模块，用于生成时间戳和计时


class 铜心织宙:
    """
    铜心织宙PCB分析系统

    本系统用于自动分析和处理PCB（印刷电路板）图像，识别电路板上的元件、
    引脚和铜层走线，并生成可视化的原理图。系统具有以下功能：

    1. 图像预处理和增强
    2. 元件和引脚检测
    3. 铜层提取和分析
    4. 原理图自动生成
    5. 结果保存和报告生成

    使用方法：
    1. 创建铜心织宙实例
    2. 调用run方法并提供PCB图像路径
    3. 查看output目录下的分析结果
    """

    def __init__(self) -> None:
        """
        初始化铜心织宙PCB分析系统

        创建系统实例并初始化所有必要的属性和数据结构。
        系统启动时会自动创建必要的输出目录结构。
        """
        super(铜心织宙, self).__init__()

        self.components = []  # 存储检测到的元件列表，每个元件为(x, y, w, h)格式
        self.mask = None  # 存储铜层掩码，用于分析铜层走线
        self.processed_image = None  # 存储处理后的图像，包含标记的元件和引脚
        self.output_dir = "output"  # 输出目录路径配置
        self.ensure_output_dirs()  # 确保输出目录结构存在

    def ensure_output_dirs(self) -> None:
        """
        确保输出目录结构存在

        检查并创建系统所需的所有输出目录，包括：
        - 主输出目录
        - 原理图输出目录
        - 处理后图像目录
        - 分析结果目录
        """
        directories = [
            f"{self.output_dir}",  # 主输出目录
            f"{self.output_dir}/schematics",  # 存储生成的原理图
            f"{self.output_dir}/processed",  # 存储处理后的PCB图像
            f"{self.output_dir}/analysis",  # 存储分析报告和数据
        ]
        for directory in directories:
            if not os.path.exists(directory):
                os.makedirs(directory)

    def save_results(self) -> None:
        """
        保存处理结果到指定目录

        将系统生成的所有结果保存到对应的输出目录中，包括：
        1. 原理图 - 保存到schematics目录
        2. 处理后的PCB图像 - 保存到processed目录
        3. 铜层轨迹图像 - 保存到processed目录
        4. 分析报告文本 - 保存到analysis目录

        所有文件名都包含时间戳以避免覆盖。

        输出文件结构：
        output/
        ├── schematics/    # 原理图输出
        │   └── schematic_{timestamp}.jpg
        ├── processed/     # 处理后的图像
        │   ├── pcb_{timestamp}.jpg
        │   └── copper_tracks_{timestamp}.jpg
        └── analysis/      # 分析结果
            └── results_{timestamp}.txt
        """
        if self.processed_image is None:
            return

        # 生成时间戳，格式为年月日_时分秒
        timestamp = time.strftime("%Y%m%d_%H%M%S")

        # 保存原理图（高质量JPEG格式）
        schematic_path = f"{self.output_dir}/schematics/schematic_{timestamp}.jpg"
        cv2.imwrite(schematic_path, self.schematic, [cv2.IMWRITE_JPEG_QUALITY, 100])

        # 保存处理后的PCB图像（包含标记的元件和引脚）
        processed_path = f"{self.output_dir}/processed/pcb_{timestamp}.jpg"
        cv2.imwrite(processed_path, self.processed_image)

        # 保存铜层轨迹图像（仅显示铜层部分）
        tracks_path = f"{self.output_dir}/processed/copper_tracks_{timestamp}.jpg"
        copper_tracks = cv2.bitwise_and(
            self.processed_image, self.processed_image, mask=self.mask
        )
        cv2.imwrite(tracks_path, copper_tracks)

        # 保存分析结果文本报告
        analysis_path = f"{self.output_dir}/analysis/results_{timestamp}.txt"
        with open(analysis_path, "w", encoding="utf-8") as f:
            f.write(f"PCB Analysis Report\n")
            f.write(f"Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}\n\n")
            f.write(f"Components Detected: {len(self.components)}\n")
            f.write(f"Component Positions:\n")
            for idx, (x, y, w, h) in enumerate(self.components, 1):
                f.write(f"Component {idx}: Position=({x}, {y}), Size={w}x{h}\n")

    def generate_schematic(self) -> None:
        """
        生成PCB原理图

        基于检测到的元件和铜层信息，生成美观专业的PCB原理图。
        原理图包含以下元素：
        1. 专业的网格背景和标题栏
        2. 元件符号和标签
        3. 引脚连接点
        4. 连接线表示
        5. 时间戳和边框

        生成的原理图会保存为类属性，并同时保存到输出目录。
        """
        if not self.components or self.mask is None:
            raise ValueError("需要先处理PCB图像")

        # 创建更大的画布以容纳标签和边距
        height, width = self.mask.shape[:2]
        canvas_height = height + 100  # 增加高度以适应标签
        canvas_width = width + 100  # 增加宽度以适应标签
        schematic = np.ones((canvas_height, canvas_width, 3), dtype=np.uint8) * 255

        # 添加基础网格背景
        grid_size = 20
        for x in range(0, canvas_width, grid_size):
            cv2.line(schematic, (x, 0), (x, canvas_height), (240, 240, 240), 1)
        for y in range(0, canvas_height, grid_size):
            cv2.line(schematic, (0, y), (canvas_width, y), (240, 240, 240), 1)

        # 创建专业背景色调（淡蓝色）
        for y in range(canvas_height):
            for x in range(canvas_width):
                schematic[y, x] = [
                    250,
                    252,
                    255,
                ]  # 淡蓝色调，提供专业外观

        # 增强网格，添加渐变效果
        grid_size = 40
        for x in range(0, canvas_width, grid_size):
            alpha = 0.7 - (x / canvas_width) * 0.3  # 创建水平方向的渐变
            color = tuple([int(240 + 15 * alpha)] * 3)
            cv2.line(schematic, (x, 0), (x, canvas_height), color, 1)

        for y in range(0, canvas_height, grid_size):
            alpha = 0.7 - (y / canvas_height) * 0.3  # 创建垂直方向的渐变
            color = tuple([int(240 + 15 * alpha)] * 3)
            cv2.line(schematic, (0, y), (canvas_width, y), color, 1)

        # 添加专业标题栏
        header_height = 60
        cv2.rectangle(
            schematic, (0, 0), (canvas_width, header_height), (240, 245, 250), -1
        )
        cv2.line(
            schematic,
            (0, header_height),
            (canvas_width, header_height),
            (200, 200, 200),
            2,
        )

        # 添加带阴影效果的标题
        title = "PCB Schematic Analysis"
        font = cv2.FONT_HERSHEY_SIMPLEX
        cv2.putText(schematic, title, (51, 41), font, 1.2, (150, 150, 150), 3)  # 阴影
        cv2.putText(schematic, title, (50, 40), font, 1.2, (50, 50, 50), 2)  # 主文本

        # 添加时间戳和信息
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        cv2.putText(
            schematic,
            f"Generated: {timestamp}",
            (canvas_width - 300, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (100, 100, 100),
            1,
        )

        # 绘制元件符号和标签
        for idx, (x, y, w, h) in enumerate(self.components, 1):
            # 计算元件在原理图上的中心位置（添加偏移）
            center_x = x + w // 2 + 50
            center_y = y + h // 2 + 50

            # 绘制美化的元件框（矩形表示）
            cv2.rectangle(
                schematic,
                (center_x - 30, center_y - 25),
                (center_x + 30, center_y + 25),
                (50, 50, 50),
                2,
            )

            # 添加内部细节（电阻符号横线）
            cv2.line(
                schematic,
                (center_x - 20, center_y),
                (center_x + 20, center_y),
                (100, 100, 100),
                1,
            )

            # 添加美化的引脚点（红色圆点带黑色边框）
            cv2.circle(
                schematic, (center_x - 30, center_y), 4, (0, 0, 255), -1
            )  # 左引脚
            cv2.circle(
                schematic, (center_x + 30, center_y), 4, (0, 0, 255), -1
            )  # 右引脚
            cv2.circle(
                schematic, (center_x - 30, center_y), 6, (0, 0, 0), 1
            )  # 左引脚边框
            cv2.circle(
                schematic, (center_x + 30, center_y), 6, (0, 0, 0), 1
            )  # 右引脚边框

            # 添加元件标签（编号）
            label_text = f"C{idx}"  # 元件编号，C代表组件
            font = cv2.FONT_HERSHEY_SIMPLEX
            cv2.putText(
                schematic,
                label_text,
                (center_x - 15, center_y - 35),
                font,
                0.6,
                (0, 0, 0),
                2,
            )

            # 添加引脚标签（更小字体）
            pin_font = cv2.FONT_HERSHEY_SIMPLEX
            pin_font_scale = 0.4

            # 左引脚标签
            cv2.putText(
                schematic,
                f"PIN{idx}_L",
                (center_x - 45, center_y - 5),
                pin_font,
                pin_font_scale,
                (0, 0, 0),
                1,
            )
            # 右引脚标签
            cv2.putText(
                schematic,
                f"PIN{idx}_R",
                (center_x + 15, center_y - 5),
                pin_font,
                pin_font_scale,
                (0, 0, 0),
                1,
            )

            # 元件类型标签
            cv2.putText(
                schematic,
                "RES",  # 电阻器缩写
                (center_x - 15, center_y + 5),
                pin_font,
                pin_font_scale,
                (100, 100, 100),
                1,
            )

        # 改进的连接线绘制（基于铜层掩码）
        _, copper_binary = cv2.threshold(self.mask, 127, 255, cv2.THRESH_BINARY)
        copper_lines = cv2.HoughLinesP(
            copper_binary, 1, np.pi / 180, 50, minLineLength=40, maxLineGap=20
        )

        if copper_lines is not None:
            for line in copper_lines:
                x1, y1, x2, y2 = line[0]
                # 添加偏移以匹配画布大小
                x1, x2 = x1 + 50, x2 + 50
                y1, y2 = y1 + 50, y2 + 50
                # 绘制美化的连接线
                cv2.line(schematic, (x1, y1), (x2, y2), (0, 0, 0), 2)
                # 添加连接点（小圆点）
                cv2.circle(schematic, (x1, y1), 3, (50, 50, 50), -1)
                cv2.circle(schematic, (x2, y2), 3, (50, 50, 50), -1)

        # 添加标题和边框
        title = ""  # 此处标题为空，因为已在上方添加了标题
        cv2.putText(
            schematic, title, (50, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2
        )
        cv2.rectangle(
            schematic, (10, 10), (canvas_width - 10, canvas_height - 10), (0, 0, 0), 2
        )

        # Store schematic as class attribute before saving
        self.schematic = schematic

        # 保存高质量原理图
        cv2.imwrite("output/schematic.jpg", schematic, [cv2.IMWRITE_JPEG_QUALITY, 100])

    def run(self, image_path: str = None) -> None:
        if image_path is None:
            raise ValueError("请提供有效的图像路径")

        # 验证文件是否存在
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"未找到图像文件：{image_path}")

        self.Image_Process(image_path)
        self.Neural_Network()
        self.generate_schematic()  # 添加生成原理图步骤
        self.save_results()

    def Image_Process(self, image_path: str) -> tuple:
        # Specific pin detection parameters
        pin_min_area = 50
        pin_max_area = 300
        min_area = 100  # Add this line for component detection
        pin_aspect_ratio_range = (0.8, 1.2)

        img = cv2.imread(image_path)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        # Enhance pin detection
        _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
        kernel = np.ones((3, 3), np.uint8)
        binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)

        # Find pin contours
        contours, _ = cv2.findContours(
            binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )

        pins = []
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if pin_min_area < area < pin_max_area:
                x, y, w, h = cv2.boundingRect(cnt)
                aspect_ratio = w / float(h)

                if pin_aspect_ratio_range[0] < aspect_ratio < pin_aspect_ratio_range[1]:
                    pins.append((x, y, w, h))
                    # Mark detected pins
                    cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
                    cv2.circle(img, (x + w // 2, y + h // 2), 2, (0, 0, 255), -1)

        self.pins = pins  # Store pin locations
        self.processed_image = img

        # Continue with existing component detection
        # 用于存储检测到的元件信息
        components = []

        # 读取PCB图像并进行预处理
        # 1. 读取原始图像
        # 2. 转换为灰度图
        # 3. 使用高斯模糊减少图像噪声
        # 4. 二值化处理，将图像转换为黑白图像
        img = cv2.imread(image_path)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        blur = cv2.GaussianBlur(gray, (5, 5), 0)
        _, binary = cv2.threshold(blur, 127, 255, cv2.THRESH_BINARY_INV)

        # 图像去噪和形态学处理
        # 使用开运算去除小的噪点
        kernel = np.ones((3, 3), np.uint8)
        denoised = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)

        # 检测电路板上的元件轮廓
        # 使用EXTERNAL方式仅检测外部轮廓
        contours, _ = cv2.findContours(
            denoised, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )

        # 遍历所有检测到的轮廓
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > min_area:  # Now min_area is defined
                x, y, w, h = cv2.boundingRect(cnt)

                component = (x, y, w, h)
                components.append(component)
                self.components.append(component)  # Add this line
                cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)

        # 提取铜层走线
        # 增强图像对比度
        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
        cl = clahe.apply(l)
        enhanced = cv2.merge((cl, a, b))
        enhanced = cv2.cvtColor(enhanced, cv2.COLOR_LAB2BGR)

        # 转换到HSV空间并提取铜层
        hsv = cv2.cvtColor(enhanced, cv2.COLOR_BGR2HSV)

        # 定义多个铜色范围
        copper_ranges = [
            (np.array([0, 50, 50]), np.array([20, 255, 255])),  # 红铜色
            (np.array([170, 50, 50]), np.array([180, 255, 255])),  # 深红铜色
            (np.array([20, 50, 50]), np.array([40, 255, 255])),  # 黄铜色
        ]

        # 合并多个颜色范围的掩码
        mask = np.zeros(hsv.shape[:2], dtype=np.uint8)
        for lower, upper in copper_ranges:
            copper_mask = cv2.inRange(hsv, lower, upper)
            mask = cv2.bitwise_or(mask, copper_mask)

        # 应用形态学操作改善掩码质量
        kernel = np.ones((3, 3), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)

        tracks = cv2.bitwise_and(img, img, mask=mask)

        # 将处理后的图像和掩码保存为类属性
        self.processed_image = img
        self.mask = mask

        # 返回检测到的元件信息和铜层轨迹
        return components, tracks

    def Neural_Network(self) -> list:
        # 确保已经执行过图像处理
        if self.mask is None:
            raise ValueError("请先执行图像处理步骤")

        # 对掩码进行标记分析
        labeled = label(self.mask)
        regions = regionprops(labeled)

        # 分析每个区域的特征
        results = []

        for region in regions:
            result = {
                "label": region.label,  # 区域标签
                "area": region.area,  # 区域面积
                "centroid": region.centroid,  # 区域中心点
            }
            results.append(result)

            print(
                f"网络区域 {result['label']}: 面积={result['area']}, 中心点={result['centroid']}"
            )

        return results

    def save_results(self) -> None:
        # Generate timestamp
        timestamp = time.strftime("%Y%m%d_%H%M%S")

        if self.processed_image is not None:
            # Save processed PCB image
            processed_path = f"{self.output_dir}/processed/pcb_{timestamp}.jpg"
            cv2.imwrite(processed_path, self.processed_image)

            # Save copper tracks image with timestamp
            copper_path = f"{self.output_dir}/processed/copper_tracks_{timestamp}.jpg"
            copper_tracks = cv2.bitwise_and(
                self.processed_image, self.processed_image, mask=self.mask
            )
            cv2.imwrite(copper_path, copper_tracks)

            # Save schematic with timestamp
            schematic_path = f"{self.output_dir}/schematics/schematic_{timestamp}.jpg"
            cv2.imwrite(schematic_path, self.schematic, [cv2.IMWRITE_JPEG_QUALITY, 100])

            # Save analysis results
            analysis_path = f"{self.output_dir}/analysis/results_{timestamp}.txt"
            with open(analysis_path, "w", encoding="utf-8") as f:
                f.write(f"PCB Analysis Report\n")
                f.write(f"Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}\n\n")
                f.write(f"Components Detected: {len(self.components)}\n")
                f.write(f"Component Positions:\n")
                for idx, (x, y, w, h) in enumerate(self.components, 1):
                    f.write(f"Component {idx}: Position=({x}, {y}), Size={w}x{h}\n")


if __name__ == "__main__":
    # 创建铜心织宙PCB分析系统实例
    app = 铜心织宙()
    # 提示用户输入PCB图像路径
    image_path = input("请将PCB图像拖入控制台（或输入图像路径）: ")
    # 去除路径中可能存在的引号
    image_path = image_path.strip("'\"")
    # 运行PCB分析
    app.run(image_path)