minya_hotomoe/modules/MazeGenerator.py

import cv2
import numpy
import random
import os

from module_interface import ModuleInterface


class MazeGenerator(ModuleInterface):
    def __init__(self):
        super().__init__()

        self.name = "MazeGenerator"
        self.funcs["generate_maze"] = self.generateMaze

        print("[MazeGenerator] MazeGenerator Module loaded.")

    # generate maze image
    async def generateMaze(self):
        width = 3280
        height = 3280
        img = numpy.zeros((width, height, 3), dtype = "uint8")

        LEFT = 1
        RIGHT = 2
        UP = 3
        DOWN = 4

        maze_x = 50
        maze_y = 50
        maze_width = width - maze_x * 2
        maze_height = height - maze_y * 2
        maze_col = 30
        maze_row = 30

        wall_size = 8

        cell_x = 0
        cell_y = 0
        
        cell_stack = [[0, 0]] # [[x, y]]
        cells_visited = [[0, 0]] # [[x, y]]

        # fill image
        cv2.rectangle(img, (0, 0), (width, height), (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)), -1)

        # draw border
        border_color = (0, 0, 0)
        cv2.rectangle(img, (int(maze_x - wall_size / 2), int(maze_y - wall_size / 2)), (int(width - maze_x + wall_size / 2), int(height - maze_y + wall_size / 2)), border_color, -1)

        ## draw maze ##
        size = int(maze_width / maze_col)

        cv2.rectangle(
            img,
            (cell_x * size + wall_size + maze_x, cell_y * size + wall_size + maze_y),
            (cell_x * size + size - wall_size + maze_x, cell_y * size + size - wall_size + maze_y),
            (255, 255, 255),
            -1
        )

        while len(cell_stack) > 0:
            # log pos
            if not [cell_x, cell_y] in cells_visited:
                cells_visited.append([cell_x, cell_y])

            # get around_cells
            around_cells = []
            around_cells.append([cell_x + 1, cell_y, LEFT])
            around_cells.append([cell_x - 1, cell_y, RIGHT])
            around_cells.append([cell_x, cell_y - 1, UP])
            around_cells.append([cell_x, cell_y + 1, DOWN])

            temp = around_cells.copy()

            for around_cell in temp:
                if [around_cell[0], around_cell[1]] in cells_visited or around_cell[0] < 0 or around_cell[1] < 0 or around_cell[0] >= maze_col or around_cell[1] >= maze_row:
                    around_cells.pop(around_cells.index(around_cell))

            if len(around_cells) > 0:
                # draw cell
                next_pos = random.choice(around_cells)
                cell_stack.append([next_pos[0], next_pos[1]])
                
                cv2.rectangle(
                    img,
                    (next_pos[0] * size + wall_size + maze_x, next_pos[1] * size + wall_size + maze_y),
                    (next_pos[0] * size + size - wall_size + maze_x, next_pos[1] * size + size - wall_size + maze_y),
                    (255, 255, 255),
                    -1
                )

                wall_x = next_pos[0] * size
                wall_y = next_pos[1] * size

                if next_pos[2] == LEFT:
                    wall_x -= size / 2
                elif next_pos[2] == RIGHT:
                    wall_x += size / 2
                elif next_pos[2] == UP:
                    wall_y += size / 2
                elif next_pos[2] == DOWN:
                    wall_y -= size / 2
                
                cv2.rectangle(
                    img,
                    (int(wall_x + wall_size + maze_x), int(wall_y + wall_size + maze_y)),
                    (int(wall_x + size - wall_size + maze_x), int(wall_y + size - wall_size + maze_y)),
                    (255, 255, 255),
                    -1
                )

                cell_x = next_pos[0]
                cell_y = next_pos[1]
                cell_stack.append([cell_x, cell_y])
            else:
                if cell_x == 0 and cell_y == 0:
                    break

                # go back
                cell_x = cell_stack[-1][0]
                cell_y = cell_stack[-1][1]
                cell_stack.pop()
        ##

        # draw start point
        cv2.circle(img, (int(size - size / 2 + maze_x), int(size - size / 2 + maze_y)), 15, (0, 0, 0), -1)

        # draw goal point
        cv2.circle(img, (int(maze_width - size / 2 + maze_x), int(maze_height - size / 2 + maze_y)), 15, (22, 183, 65), -1)

        # download image
        cv2.imwrite(os.path.dirname(os.path.abspath(__file__)) + "/maze/maze.png", img)

        print("[MazeGenerator] Maze was generated")