acdr/acdr-ui/tootls/imgtools.py

import os
import cv2
import numpy as np
import re
from tkinter import *
from tkinter import filedialog, messagebox, colorchooser
from PIL import Image
import tkinter as tk

# 主应用类
class ImageSplitterApp:
    def __init__(self, root):
        self.root = root
        self.root.title("图像分割工具")
        self.root.geometry("500x450")

        # 文件或文件夹路径
        self.file_path = None
        self.folder_path = None
        self.selected_color = (0, 0, 0)  # 默认黑色

        # 创建选择文件按钮
        Button(self.root, text="选择单个文件", command=self.select_file).pack(pady=10)
        Button(self.root, text="选择文件夹", command=self.select_folder).pack(pady=10)

        # 选择颜色按钮
        Label(self.root, text="选择分割线颜色:").pack(pady=10)
        Button(self.root, text="打开取色器", command=self.pick_color).pack(pady=5)

        # 当前选择的颜色显示
        self.color_label = Label(self.root, text=f"当前颜色: {self.selected_color}")
        self.color_label.pack(pady=5)

        # 命名规则
        Label(self.root, text="请输入保存图片的命名规则(支持正则):").pack(pady=10)
        self.naming_entry = Entry(self.root)
        self.naming_entry.insert(0, r"sub_image_{idx}")  # 默认命名规则
        self.naming_entry.pack(pady=5)

        # 分割按钮
        Button(self.root, text="开始分割", command=self.split_image).pack(pady=20)

        # 选择保存文件夹
        self.save_folder_path = None
        Button(self.root, text="选择保存文件夹", command=self.select_save_folder).pack(pady=10)

    def select_file(self):
        self.file_path = filedialog.askopenfilename(filetypes=[("Image Files", "*.png *.jpg *.jpeg")])
        if self.file_path:
            messagebox.showinfo("文件选择", f"已选择文件: {self.file_path}")

    def select_folder(self):
        self.folder_path = filedialog.askdirectory()
        if self.folder_path:
            messagebox.showinfo("文件夹选择", f"已选择文件夹: {self.folder_path}")

    def select_save_folder(self):
        self.save_folder_path = filedialog.askdirectory()
        if self.save_folder_path:
            messagebox.showinfo("保存文件夹选择", f"已选择保存文件夹: {self.save_folder_path}")

    # 颜色取色器
    def pick_color(self):
        color = colorchooser.askcolor(title="选择分割线颜色")
        if color:
            self.selected_color = tuple(map(int, color[0]))  # 颜色转换为整数的 (r, g, b)
            self.color_label.config(text=f"当前颜色: {self.selected_color}")

    # 查找分割线的位置 (通过指定颜色)
    def find_dividing_lines(self, image, line_color):
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # 将图像转换为灰度图
        _, binary = cv2.threshold(gray, 128, 255, cv2.THRESH_BINARY_INV)  # 进行二值化

        # 查找水平和垂直线条
        horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 1))
        vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 25))

        # 水平线
        horizontal_lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horizontal_kernel, iterations=2)

        # 垂直线
        vertical_lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, vertical_kernel, iterations=2)

        return horizontal_lines, vertical_lines

    # 分割图像
    def split_image(self):
        # 提取颜色
        line_color = self.selected_color

        # 读取命名规则
        naming_rule = self.naming_entry.get()

        # 处理单个文件
        if self.file_path:
            self.process_file(self.file_path, line_color, naming_rule)

        # 处理文件夹中的所有文件
        if self.folder_path:
            for file_name in os.listdir(self.folder_path):
                if file_name.endswith(('.png', '.jpg', '.jpeg')):
                    file_path = os.path.join(self.folder_path, file_name)
                    self.process_file(file_path, line_color, naming_rule)

    # 处理单个文件的分割
    def process_file(self, file_path, line_color, naming_rule):
        image = cv2.imread(file_path)
        if image is None:
            messagebox.showerror("错误", f"无法加载图像: {file_path}")
            return

        horizontal_lines, vertical_lines = self.find_dividing_lines(image, line_color)
        sub_images = self.split_image_by_lines(image, horizontal_lines, vertical_lines)

        # 保存分割后的图像
        self.save_sub_images(sub_images, naming_rule)

    # 根据分割线分割图像
    def split_image_by_lines(self, image, horizontal_lines, vertical_lines):
        horizontal_coords = np.where(np.any(horizontal_lines > 0, axis=1))[0]
        vertical_coords = np.where(np.any(vertical_lines > 0, axis=0))[0]

        sub_images = []
        for i in range(len(horizontal_coords) - 1):
            for j in range(len(vertical_coords) - 1):
                x1, x2 = vertical_coords[j], vertical_coords[j + 1]
                y1, y2 = horizontal_coords[i], horizontal_coords[i + 1]
                sub_image = image[y1:y2, x1:x2]
                sub_images.append(sub_image)

        return sub_images

    # 保存分割后的图像，支持正则命名规则
    def save_sub_images(self, sub_images, naming_rule):
        if not self.save_folder_path:
            messagebox.showerror("错误", "请先选择保存文件夹")
            return

        for idx, img in enumerate(sub_images):
            filename = re.sub(r'{idx}', str(idx), naming_rule)  # 使用正则表达式进行命名
            save_path = os.path.join(self.save_folder_path, f"{filename}.png")
            cv2.imwrite(save_path, img)

        messagebox.showinfo("完成", f"图像已分割并保存到 {self.save_folder_path}")

# 运行图形界面应用
if __name__ == "__main__":
    root = tk.Tk()
    app = ImageSplitterApp(root)
    root.mainloop()