Utils/VocMaskFixer.py - master - rodion_1521/EntropyMapSegmentation

Utils/VocMaskFixer.py
61 строка2 KB
rodion_1521
benchmark
02 июн 2025, 11:32
02 июн 2025, 11:329a12897
import os
import cv2
import numpy as np
import tkinter as tk
from tkinter import filedialog, messagebox, ttk

def fill_mask_contours(input_folder, output_folder):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    mask_files = [f for f in os.listdir(input_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    if not mask_files:
        messagebox.showwarning("Нет масок", "В папке не найдено изображений масок.")
        return

    for filename in mask_files:
        input_path = os.path.join(input_folder, filename)
        output_path = os.path.join(output_folder, filename)

        mask = cv2.imread(input_path, cv2.IMREAD_GRAYSCALE)
        if mask is None:
            print(f"Пропущено (не удалось загрузить): {filename}")
            continue

        _, binary = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        filled_mask = np.zeros_like(binary)
        cv2.drawContours(filled_mask, contours, -1, 255, thickness=cv2.FILLED)

        cv2.imwrite(output_path, filled_mask)

    messagebox.showinfo("Готово", f"Все маски обработаны и сохранены в:\n{output_folder}")

def select_folder_and_process():
    folder_path = filedialog.askdirectory(title="Выберите папку с масками")
    if not folder_path:
        return
    output_folder = os.path.join(folder_path, "filled_masks")
    fill_mask_contours(folder_path, output_folder)

def run_gui():
    root = tk.Tk()
    root.title("Заливка контуров масок")
    root.geometry("400x150")
    root.resizable(False, False)

    style = ttk.Style()
    style.configure("TButton", font=('Helvetica', 10), padding=10)

    label = ttk.Label(root, text="Выберите папку с масками, содержащими только контуры:",
                      font=('Helvetica', 10), wraplength=380, justify="center")
    label.pack(pady=20)

    select_btn = ttk.Button(root, text="Выбрать папку и обработать", command=select_folder_and_process)
    select_btn.pack(pady=10)

    root.mainloop()

if __name__ == "__main__":
    run_gui()