image to pixelart

main.py

import sys
import os
import copy
from collections import Counter
import cv2
import numpy
from PySide6.QtWidgets import (
    QApplication,
    QMainWindow,
    QWidget,
    QVBoxLayout,
    QHBoxLayout,
    QPushButton,
    QLabel,
    QFileDialog,
    QSpinBox,
    QDoubleSpinBox,
    QCheckBox,
    QGroupBox,
    QMessageBox,
    QSplitter,
    QScrollArea,
    QToolBar,
    QProgressBar,
    QStatusBar,
)
from PySide6.QtGui import QPixmap, QImage, QKeySequence, QShortcut, QAction
from PySide6.QtCore import Qt, QThread, Signal


def get_lines(edges):
    """Detect all lines in the image."""
    lines = cv2.HoughLines(edges, 1 / 2, numpy.pi / (180 * 2**6), 100)
    if lines is None:
        raise ValueError('Линии не обнаружены. Попробуйте другое изображение.')
    lines = lines.reshape(-1, 2).tolist()
    return lines


def get_angle_offset(lines):
    """Get the average angle offset (rad) of the lines."""
    angle_sum = 0
    for line in lines:
        angle_sum += ((line[1] + (numpy.pi / 4)) % (numpy.pi / 2)) - (numpy.pi / 4)
    avg_angle = angle_sum / len(lines)
    return avg_angle


def get_line_distances(lines):
    """Get the distances between all the lines."""
    line_distances = []
    for line_1 in lines:
        for line_2 in lines:
            line_distances.append(abs(abs(line_1[0]) - abs(line_2[0])))
    return line_distances


def get_average_line_distance(lines, pixel_width):
    """Get an average distance between all the lines that are 1 'pixel' apart."""
    line_distances = get_line_distances(lines)
    sorted_line_distances = Counter(line_distances).most_common()

    def filter_lambda(ld):
        return ld[0] > (pixel_width * 0.8) and ld[0] < (pixel_width * 1.2) and ld[1] > len(lines) / 2

    valid_lengths = list(filter(filter_lambda, sorted_line_distances))
    if len(valid_lengths) == 0:
        raise ValueError(f'Не найдено пикселей на основе ширины {pixel_width}. Попробуйте указать другое значение.')
    length_sum = 0
    count = 0
    for length in valid_lengths:
        length_sum += length[0] * length[1]
        count += length[1]
    average_line_distance = length_sum / count
    return average_line_distance


def get_average_pixel_offset(lines, average_line_distance):
    """Get the average x and y pixel offset."""
    offset_sum_x = 0
    offset_sum_y = 0
    for line in lines:
        if line[1] < numpy.pi:
            offset_sum_y += line[0] % average_line_distance
        else:
            offset_sum_x += line[0] % average_line_distance
    avg_offset_x = offset_sum_x / len(lines)
    avg_offset_y = offset_sum_y / len(lines)
    return (avg_offset_x, avg_offset_y)


def get_shape(image):
    """Get the dimensions of an image."""
    shape = image.shape
    height = shape[0]
    width = shape[1]
    return height, width


def get_pixel_image_and_coordinates(image, average_angle_offset, average_pixel_offset, average_line_distance):
    """Get the new image and the coordinates of the pixels in relation to the original image"""
    height, width = get_shape(image)
    pixel_height = height
    pixel_width = width
    pixel_image = numpy.full((pixel_height, pixel_width, 3), [255, 255, 255], dtype=numpy.uint8)
    cos = numpy.cos(average_angle_offset)
    sin = numpy.sin(average_angle_offset)
    pixel_offset_x = (average_pixel_offset[0] / average_line_distance) - pixel_width / 2
    pixel_offset_y = (average_pixel_offset[1] / average_line_distance) - pixel_height / 2
    pixel_coordinates = []
    for pixel_y in range(pixel_height):
        for pixel_x in range(pixel_width):
            pixel_x_unit = pixel_x + 0.5 + pixel_offset_x
            pixel_y_unit = pixel_y + 0.5 + pixel_offset_y
            x_unit = (pixel_x_unit * cos) - (pixel_y_unit * sin)
            y_unit = (pixel_x_unit * sin) + (pixel_y_unit * cos)
            x_scaled = int(average_line_distance * x_unit)
            y_scaled = int(average_line_distance * y_unit)
            if x_scaled < width and x_scaled >= 0 and y_scaled < height and y_scaled >= 0:
                pixel_coordinates.append((x_scaled, y_scaled))
                pixel_image[pixel_y, pixel_x] = image[y_scaled, x_scaled]
    return pixel_image, pixel_coordinates


def draw_lines(lines, image):
    """Draw lines on an image."""
    for line in lines:
        rho = line[0]
        theta = line[1]
        cos = numpy.cos(theta)
        sin = numpy.sin(theta)
        x_0 = cos * rho
        y_0 = sin * rho
        x_1 = int(x_0 + 10000 * (-sin))
        y_1 = int(y_0 + 10000 * (cos))
        x_2 = int(x_0 - 10000 * (-sin))
        y_2 = int(y_0 - 10000 * (cos))
        cv2.line(image, (x_1, y_1), (x_2, y_2), (0, 0, 0), 1)


def draw_points_on_image(image, points):
    """Draw a set of points on an image."""
    for point in points:
        image[int(point[1]), int(point[0])] = [0, 0, 0]


def crop_image(image, pixel_width):
    """Crop an image using fast vectorized NumPy actions."""
    non_white = numpy.any(image != [255, 255, 255], axis=-1)
    coords = numpy.argwhere(non_white)
    if coords.size == 0:
        raise ValueError(f'Не найдено цветных пикселей на основе ширины {pixel_width}. Попробуйте изменить настройки.')
    y_min, x_min = coords.min(axis=0)
    y_max, x_max = coords.max(axis=0)
    height, width = image.shape[:2]
    top = max(0, y_min - 1)
    bottom = min(height, y_max + 2)
    left = max(0, x_min - 1)
    right = min(width, x_max + 2)
    return image[top:bottom, left:right]


def get_background_mask(image):
    """Flood an image to create background mask."""
    height, width = get_shape(image)
    mask = numpy.zeros((height + 2, width + 2), numpy.uint8)
    diff = 10
    diff_array = [diff, diff, diff]
    cv2.floodFill(
        numpy.ascontiguousarray(image, dtype=numpy.uint8), mask, (0, 0), [0, 0, 0], loDiff=diff_array, upDiff=diff_array
    )
    return mask


def make_background_transparent(image, mask):
    """Make the background of an image transparent using a mask"""
    height, width = get_shape(image)
    pixel_image_transparent = numpy.zeros((height, width, 4), dtype=numpy.uint8)
    for y_pos in range(height):
        for x_pos in range(width):
            if not mask[y_pos + 1, x_pos + 1]:
                pixel = image[y_pos, x_pos]
                pixel_image_transparent[y_pos, x_pos] = [pixel[0], pixel[1], pixel[2], 255]
    return pixel_image_transparent


def check_left(x_pos, image, y_pos, border_pixels, stack, checked_pixels):
    if x_pos > 0:
        if image[y_pos, x_pos - 1][3]:
            border_pixels.add((x_pos, y_pos))
        else:
            stack.append((image, x_pos - 1, y_pos, border_pixels, checked_pixels))
            check_left_up(y_pos, image, x_pos, border_pixels, stack, checked_pixels)


def check_left_up(y_pos, image, x_pos, border_pixels, stack, checked_pixels):
    if y_pos > 0:
        if image[y_pos - 1, x_pos - 1][3]:
            border_pixels.add((x_pos, y_pos))
        else:
            stack.append((image, x_pos - 1, y_pos - 1, border_pixels, checked_pixels))


def check_down(y_pos, height, image, x_pos, border_pixels, stack, checked_pixels):
    if y_pos < height - 1:
        if image[y_pos + 1, x_pos][3]:
            border_pixels.add((x_pos, y_pos))
        else:
            stack.append((image, x_pos, y_pos + 1, border_pixels, checked_pixels))
            check_down_left(x_pos, image, y_pos, border_pixels, stack, checked_pixels)


def check_down_left(x_pos, image, y_pos, border_pixels, stack, checked_pixels):
    if x_pos > 0:
        if image[y_pos + 1, x_pos - 1][3]:
            border_pixels.add((x_pos, y_pos))
        else:
            stack.append((image, x_pos - 1, y_pos + 1, border_pixels, checked_pixels))


def check_right(x_pos, width, image, y_pos, border_pixels, stack, checked_pixels, height):
    if x_pos < width - 1:
        if image[y_pos, x_pos + 1][3]:
            border_pixels.add((x_pos, y_pos))
        else:
            stack.append((image, x_pos + 1, y_pos, border_pixels, checked_pixels))
            check_right_down(y_pos, height, image, x_pos, border_pixels, stack, checked_pixels)


def check_right_down(y_pos, height, image, x_pos, border_pixels, stack, checked_pixels):
    if y_pos < height - 1:
        if image[y_pos + 1, x_pos + 1][3]:
            border_pixels.add((x_pos, y_pos))
        else:
            stack.append((image, x_pos + 1, y_pos + 1, border_pixels, checked_pixels))


def check_up(y_pos, image, x_pos, border_pixels, stack, checked_pixels, width):
    if y_pos > 0:
        if image[y_pos - 1, x_pos][3]:
            border_pixels.add((x_pos, y_pos))
        else:
            stack.append((image, x_pos, y_pos - 1, border_pixels, checked_pixels))
            check_up_right(x_pos, width, image, y_pos, border_pixels, stack, checked_pixels)


def check_up_right(x_pos, width, image, y_pos, border_pixels, stack, checked_pixels):
    if x_pos < width - 1:
        if image[y_pos - 1, x_pos + 1][3]:
            border_pixels.add((x_pos, y_pos))
        else:
            stack.append((image, x_pos + 1, y_pos - 1, border_pixels, checked_pixels))


def find_border(image, x_pos, y_pos, border_pixels, checked_pixels):
    stack = [(image, x_pos, y_pos, border_pixels, checked_pixels)]
    while len(stack) > 0:
        arguments = stack.pop()
        image = arguments[0]
        x_pos = arguments[1]
        y_pos = arguments[2]
        border_pixels = arguments[3]
        checked_pixels = arguments[4]
        if (x_pos, y_pos) in checked_pixels:
            continue
        else:
            checked_pixels.add((x_pos, y_pos))
        height, width = get_shape(image)
        check_up(y_pos, image, x_pos, border_pixels, stack, checked_pixels, width)
        check_right(x_pos, width, image, y_pos, border_pixels, stack, checked_pixels, height)
        check_down(y_pos, height, image, x_pos, border_pixels, stack, checked_pixels)
        check_left(x_pos, image, y_pos, border_pixels, stack, checked_pixels)


def create_border(image):
    """Create a border around the item in the image."""
    border_pixels = set()
    checked_pixels = set()
    find_border(image, 0, 0, border_pixels, checked_pixels)
    white_pixel = [255, 255, 255, 255]
    for border_pixel in border_pixels:
        image[border_pixel[1], border_pixel[0]] = white_pixel


def crop_down(image):
    """Crop away the one pixel gap using fast NumPy slices."""
    height, width = image.shape[:2]
    if height <= 2 or width <= 2:
        return image
    return image[1:-1, 1:-1]


def scale_up(image, scale):
    """Scale up an image using optimized OpenCV interpolation."""
    height, width = image.shape[:2]
    return cv2.resize(image, (width * scale, height * scale), interpolation=cv2.INTER_NEAREST)


class ExtractionWorker(QThread):
    finished = Signal(numpy.ndarray, numpy.ndarray)
    error = Signal(str)

    def __init__(self, source_image, pixel_width, scale, add_border):
        super().__init__()
        self.source_image = source_image
        self.pixel_width = pixel_width
        self.scale = scale
        self.add_border = add_border

    def run(self):
        try:
            image_with_markings = copy.deepcopy(self.source_image)
            edges = cv2.Canny(self.source_image, 20, 50, L2gradient=True)
            lines = get_lines(edges)
            draw_lines(lines, image_with_markings)
            average_angle_offset = get_angle_offset(lines)
            average_line_distance = get_average_line_distance(lines, self.pixel_width)
            average_pixel_offset = get_average_pixel_offset(lines, average_line_distance)
            pixel_image, pixel_coordinates = get_pixel_image_and_coordinates(
                self.source_image, average_angle_offset, average_pixel_offset, average_line_distance
            )
            draw_points_on_image(image_with_markings, pixel_coordinates)
            processed = crop_image(pixel_image, self.pixel_width)
            mask = get_background_mask(processed)
            processed_transparent = make_background_transparent(processed, mask)
            if self.add_border:
                create_border(processed_transparent)
            else:
                processed_transparent = crop_down(processed_transparent)
            if self.scale > 1:
                processed_transparent = scale_up(processed_transparent, self.scale)
            self.finished.emit(image_with_markings, processed_transparent)
        except Exception as e:
            self.error.emit(str(e))


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle('Pixel Art Extractor')
        self.resize(1100, 700)
        self.setAcceptDrops(True)
        self.source_path = None
        self.cv_source_image = None
        self.cv_marked_image = None
        self.cv_output_image = None
        self.worker = None
        self.init_ui()
        self.init_shortcuts()

    def init_ui(self):
        self.toolbar = QToolBar('Главная панель инструментов')
        self.addToolBar(self.toolbar)
        self.toolbar.setMovable(False)
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        main_layout = QHBoxLayout(main_widget)
        control_panel = QWidget()
        control_layout = QVBoxLayout(control_panel)
        control_panel.setFixedWidth(270)
        load_group = QGroupBox('Операции с файлами')
        load_layout = QVBoxLayout(load_group)
        self.btn_load = QPushButton('Открыть файл')
        self.btn_load.clicked.connect(self.open_file_dialog)
        self.lbl_file_name = QLabel('Файл не выбран')
        self.lbl_file_name.setWordWrap(True)
        load_layout.addWidget(self.btn_load)
        load_layout.addWidget(self.lbl_file_name)
        control_layout.addWidget(load_group)
        settings_group = QGroupBox('Параметры сетки')
        settings_layout = QVBoxLayout(settings_group)
        settings_layout.addWidget(QLabel('Размер пикселя в файле (px):'))
        self.spin_width = QDoubleSpinBox()
        self.spin_width.setRange(0.1, 1000.0)
        self.spin_width.setValue(10.0)
        self.spin_width.setSingleStep(0.5)
        settings_layout.addWidget(self.spin_width)
        settings_layout.addWidget(QLabel('Масштаб вывода:'))
        self.spin_scale = QSpinBox()
        self.spin_scale.setRange(1, 100)
        self.spin_scale.setValue(1)
        settings_layout.addWidget(self.spin_scale)
        self.chk_border = QCheckBox('Белая рамка контура')
        settings_layout.addWidget(self.chk_border)
        control_layout.addWidget(settings_group)
        self.btn_process = QPushButton('Обработать')
        self.btn_process.clicked.connect(self.process_image)
        self.btn_process.setEnabled(False)
        self.btn_process.setStyleSheet('background-color: #2a82da; color: white; font-weight: bold; padding: 6px;')
        control_layout.addWidget(self.btn_process)
        self.btn_save = QPushButton('Сохранить результат')
        self.btn_save.clicked.connect(self.save_image)
        self.btn_save.setEnabled(False)
        control_layout.addWidget(self.btn_save)
        control_layout.addStretch()
        main_layout.addWidget(control_panel)
        self.splitter = QSplitter(Qt.Orientation.Horizontal)
        main_layout.addWidget(self.splitter)
        left_view_container = QWidget()
        left_view_layout = QVBoxLayout(left_view_container)
        left_view_layout.addWidget(QLabel('Исходное изображение и сетка:'))
        self.scroll_source = QScrollArea()
        self.scroll_source.setWidgetResizable(True)
        self.lbl_source_preview = QLabel('Перетащите изображение сюда\nили воспользуйтесь меню «Открыть»')
        self.lbl_source_preview.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.lbl_source_preview.setStyleSheet('border: 2px dashed #888; color: #777;')
        self.scroll_source.setWidget(self.lbl_source_preview)
        left_view_layout.addWidget(self.scroll_source)
        self.splitter.addWidget(left_view_container)
        right_view_container = QWidget()
        right_view_layout = QVBoxLayout(right_view_container)
        right_view_layout.addWidget(QLabel('Результат (Пиксель-арт):'))
        self.scroll_result = QScrollArea()
        self.scroll_result.setWidgetResizable(True)
        self.lbl_result_preview = QLabel('Ожидание обработки')
        self.lbl_result_preview.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.lbl_result_preview.setStyleSheet('background-color: #252525; color: #888; border: 1px solid #444;')
        self.scroll_result.setWidget(self.lbl_result_preview)
        right_view_layout.addWidget(self.scroll_result)
        self.splitter.addWidget(right_view_container)
        self.splitter.setSizes([550, 550])
        self.action_open = QAction('Открыть', self)
        self.action_open.triggered.connect(self.open_file_dialog)
        self.toolbar.addAction(self.action_open)
        self.action_run = QAction('Обработать', self)
        self.action_run.setEnabled(False)
        self.action_run.triggered.connect(self.process_image)
        self.toolbar.addAction(self.action_run)
        self.action_save = QAction('Сохранить', self)
        self.action_save.setEnabled(False)
        self.action_save.triggered.connect(self.save_image)
        self.toolbar.addAction(self.action_save)
        self.status_bar = QStatusBar()
        self.setStatusBar(self.status_bar)
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        self.progress_bar.setMaximumWidth(200)
        self.status_bar.addPermanentWidget(self.progress_bar)
        self.status_bar.showMessage('Готов к работе')

    def init_shortcuts(self):
        """Быстрые сочетания клавиш."""
        self.sc_open = QShortcut(QKeySequence('Ctrl+O'), self)
        self.sc_open.activated.connect(self.open_file_dialog)
        self.sc_process = QShortcut(QKeySequence('Ctrl+R'), self)
        self.sc_process.activated.connect(self.process_image)
        self.sc_save = QShortcut(QKeySequence('Ctrl+S'), self)
        self.sc_save.activated.connect(self.save_image)

    def dragEnterEvent(self, event):
        if event.mimeData().hasUrls():
            event.acceptProposedAction()

    def dropEvent(self, event):
        for url in event.mimeData().urls():
            file_path = url.toLocalFile()
            if os.path.isfile(file_path):
                ext = os.path.splitext(file_path)[1].lower()
                if ext in ['.png', '.jpg', '.jpeg', '.bmp', '.webp', '.tiff']:
                    self.load_image(file_path)
                    break

    def open_file_dialog(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, 'Открыть изображение', '', 'Изображения (*.png *.jpg *.jpeg *.bmp *.webp *.tiff);;Все файлы (*)'
        )
        if file_path:
            self.load_image(file_path)

    def load_image(self, file_path):
        try:
            img = cv2.imread(file_path)
            if img is None:
                raise ValueError('Не удалось открыть или декодировать графический файл.')
            self.source_path = file_path
            self.cv_source_image = img
            self.cv_marked_image = None
            self.cv_output_image = None
            self.lbl_file_name.setText(os.path.basename(file_path))
            self.display_image(self.cv_source_image, self.lbl_source_preview, self.scroll_source)
            self.lbl_result_preview.setPixmap(QPixmap())
            self.lbl_result_preview.setText('Изображение загружено. Запустите обработку.')
            self.lbl_result_preview.setStyleSheet('background-color: #252525; color: #888;')
            self.btn_process.setEnabled(True)
            self.action_run.setEnabled(True)
            self.btn_save.setEnabled(False)
            self.action_save.setEnabled(False)
            self.status_bar.showMessage('Файл загружен успешно', 4000)
        except Exception as e:
            QMessageBox.critical(self, 'Ошибка чтения', str(e))

    def process_image(self):
        if self.cv_source_image is None or (self.worker and self.worker.isRunning()):
            return
        self.set_ui_enabled(False)
        self.status_bar.showMessage('Анализ сетки и экстракция пикселей...')
        self.progress_bar.setRange(0, 0)
        self.progress_bar.setVisible(True)
        self.worker = ExtractionWorker(
            source_image=self.cv_source_image,
            pixel_width=self.spin_width.value(),
            scale=self.spin_scale.value(),
            add_border=self.chk_border.isChecked(),
        )
        self.worker.finished.connect(self.on_processing_finished)
        self.worker.error.connect(self.on_processing_error)
        self.worker.start()

    def on_processing_finished(self, marked_image, output_image):
        self.cv_marked_image = marked_image
        self.cv_output_image = output_image
        self.display_image(self.cv_marked_image, self.lbl_source_preview, self.scroll_source)
        self.display_image(self.cv_output_image, self.lbl_result_preview, self.scroll_result)
        self.set_ui_enabled(True)
        self.btn_save.setEnabled(True)
        self.action_save.setEnabled(True)
        self.progress_bar.setVisible(False)
        self.status_bar.showMessage('Обработка завершена успешно', 5000)

    def on_processing_error(self, err_msg):
        QMessageBox.warning(self, 'Ошибка расчета', err_msg)
        self.set_ui_enabled(True)
        self.progress_bar.setVisible(False)
        self.status_bar.showMessage('Ошибка в ходе выполнения операции')

    def set_ui_enabled(self, enabled):
        """Включение или отключение компонентов во время фоновой обработки."""
        self.btn_load.setEnabled(enabled)
        self.btn_process.setEnabled(enabled)
        self.btn_save.setEnabled(enabled and self.cv_output_image is not None)
        self.spin_width.setEnabled(enabled)
        self.spin_scale.setEnabled(enabled)
        self.chk_border.setEnabled(enabled)
        self.action_open.setEnabled(enabled)
        self.action_run.setEnabled(enabled)
        self.action_save.setEnabled(enabled and self.cv_output_image is not None)

    def save_image(self):
        if self.cv_output_image is None:
            return
        default_dir = os.path.dirname(self.source_path) if self.source_path else ''
        default_name = 'pixel_art_output.png'
        if self.source_path:
            base, _ = os.path.splitext(os.path.basename(self.source_path))
            default_name = f'{base}_extracted.png'
        default_path = os.path.join(default_dir, default_name)
        file_path, _ = QFileDialog.getSaveFileName(self, 'Сохранить результат', default_path, 'PNG Изображения (*.png)')
        if file_path:
            try:
                written = cv2.imwrite(file_path, self.cv_output_image)
                if written:
                    QMessageBox.information(self, 'Сохранение', f'Файл успешно записан:\n{file_path}')
                else:
                    raise IOError('Не удалось сохранить файл. Неизвестная ошибка файловой системы.')
            except Exception as e:
                QMessageBox.critical(self, 'Ошибка сохранения', str(e))

    def display_image(self, cv_img, target_label, scroll_area):
        """Преобразование цветового пространства OpenCV в QPixmap для отображения."""
        if cv_img is None:
            return
        h, w = cv_img.shape[:2]
        if len(cv_img.shape) == 3:
            if cv_img.shape[2] == 4:
                rgb_image = cv2.cvtColor(cv_img, cv2.COLOR_BGRA2RGBA)
                bytes_per_line = 4 * w
                q_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format.Format_RGBA8888)
            else:
                rgb_image = cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)
                bytes_per_line = 3 * w
                q_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format.Format_RGB888)
        else:
            bytes_per_line = w
            q_image = QImage(cv_img.data, w, h, bytes_per_line, QImage.Format.Format_Grayscale8)
        pixmap = QPixmap.fromImage(q_image.copy())
        view_w = scroll_area.viewport().width() - 4
        view_h = scroll_area.viewport().height() - 4
        scaled_pixmap = pixmap.scaled(
            view_w,
            view_h,
            Qt.AspectRatioMode.KeepAspectRatio,
            Qt.TransformationMode.SmoothTransformation,
        )
        target_label.setPixmap(scaled_pixmap)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())