image_preprocessing.py - master - Arthur159/chat_bot_BLS

image_preprocessing.py
104 строки5 KB
Arthur159
upload files
02 июн 2025, 16:35
02 июн 2025, 16:35b4c2ea2
import cv2
import os
import shutil

def preprocess_image():

    # Путь к папке с изображениями
    input_folder = 'input_images'
    # Папка для сохранения отдельных изображений
    output_folder = 'output_images'

    # Настройки для размеров контуров (площадь в пикселях)
    min_area = 4  # Минимальная площадь
    max_area = 800  # Максимальная площадь
    min_height = 25  # Минимальная высота элемента
    max_width = 22  # Максимальная ширина элемента

    # Создаем папку для выходных изображений, если она не существует
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # Получаем список всех изображений в папке
    image_files = [f for f in os.listdir(input_folder) if f.endswith(('.png', '.jpg', '.jpeg'))]

    for image_file in image_files:
        # Полный путь к изображению
        image_path = os.path.join(input_folder, image_file)

        # Создаем папку для текущего изображения
        current_output_folder = os.path.join(output_folder, os.path.splitext(image_file)[0])
        if not os.path.exists(current_output_folder):
            os.makedirs(current_output_folder)

        # Загружаем изображение
        image = cv2.imread(image_path)

        # Преобразуем изображение в оттенки серого
        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # Инверсия цвета
        inverted_image = cv2.bitwise_not(gray_image)

        # Удаление шумов с помощью фильтра Гаусса
        gaussian_blurred = cv2.GaussianBlur(inverted_image, (5, 5), 0)

        # Изменение контраста
        alpha = 1.5  # Контраст
        beta = 0  # Яркость
        contrasted_image = cv2.convertScaleAbs(gaussian_blurred, alpha=alpha, beta=beta)

        # Начальный порог бинаризации
        threshold_value = 10

        while True:
            # Применяем пороговую бинаризацию
            _, binary_image = cv2.threshold(contrasted_image, threshold_value, 255, cv2.THRESH_BINARY)

            # Находим контуры на бинарном изображении
            contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            # Фильтруем контуры по площади, высоте и ширине
            filtered_contours = []
            for contour in contours:
                area = cv2.contourArea(contour)
                x, y, w, h = cv2.boundingRect(contour)
                if min_area <= area <= max_area and h >= min_height and w <= max_width:
                    filtered_contours.append(contour)

            # Проверяем, найдено ли достаточно контуров
            if len(filtered_contours) >= 3:
                break

            # Увеличиваем порог бинаризации
            threshold_value += 1
            if threshold_value > 255:  # Ограничиваем максимальное значение порога
                print(f"Не удалось найти 3 контура в изображении {image_file} с заданными параметрами.")
                break

        # Сортируем контуры по координате x
        filtered_contours.sort(key=lambda c: cv2.boundingRect(c)[0])

        # Сохраняем отдельные бинаризованные изображения (максимум 3)
        saved_count = 0  # Счетчик сохраненных изображений
        for contour in filtered_contours[:3]:  # Сохраняем только первые 3 контура
            x, y, w, h = cv2.boundingRect(contour)
            # Извлекаем бинаризованную область
            roi = binary_image[y:y + h, x:x + w]
            cv2.imwrite(os.path.join(current_output_folder, f'contour_{saved_count}.png'), roi)
            saved_count += 1

        print(f'Сохранено {saved_count} бинаризованных изображений в папку "{current_output_folder}".')

def delete_directory_contents(directory):
    # Проверяем, существует ли директория
    if os.path.exists(directory):
         # Удаляем все содержимое директории
        for item in os.listdir(directory):
            item_path = os.path.join(directory, item)
            if os.path.isfile(item_path):
                os.remove(item_path)  # Удаляем файл
            elif os.path.isdir(item_path):
                shutil.rmtree(item_path)  # Удаляем директорию
    else:
        print(f"Директория {directory} не найдена.")