Maze of Brownian tree

brownian_tree.py

#!/usr/bin/python

from __future__ import division
from __future__ import print_function

import sys
import random
import numpy as np
import cv2

move = [
    np.array([ 1,  0], dtype=int), 
    np.array([ 0,  1], dtype=int), 
    np.array([-1,  0], dtype=int), 
    np.array([ 0, -1], dtype=int), 
]

neighbor_grid = move

class BTree(object):
    def __init__(self, H, W):
        self.N = H*W
        self.H = H
        self.W = W
        self.tree = np.zeros((H, W), dtype=int)
        self.root = (random.randint(0, W - 1), random.randint(0, H - 1))
        self.tree[self.root[1], self.root[0]] = 1
        self.edges = []
    
    def grow(self):
        free_place = np.array(np.where(self.tree == 0)).T
        p = random.choice(free_place)[-1::-1]
        while not self._check_touched_at(p):
            p = p + random.choice(move)
            p[p < 0] = 0
            p[0] = min(p[0], self.W - 1)
            p[1] = min(p[1], self.H - 1)
        x, y = p
        self.tree[y, x] = 1

    def _check_touched_at(self, p):
        qs = []
        for dx in neighbor_grid:
            q = p + dx
            if q[0] < 0 or q[1] < 0 or q[1] >= self.H or q[0] >= self.W:
                continue
            if self.tree[q[1], q[0]] == 1:
                qs.append(q)
        if len(qs) > 0:
            self.edges.append((p, random.choice(qs)))
            return True

        return False
    
    def reset(self):
        self.tree = np.zeros((self.H, self.W), dtype=int)
        self.root = (random.randint(0, self.W - 1), random.randint(0, self.H - 1))
        self.tree[self.root[1], self.root[0]] = 1
        del self.edges
        self.edges = []
    
    def is_grown(self):
        return np.all(self.tree == 1)
    
    def print_tree(self, filename):
        np.savetxt(filename, self.tree, fmt='%d')
    
    def make_image(self, grid_px=11, edge_width=5):
        w = self.W*grid_px
        h = self.H*grid_px
        img = np.zeros((h, w, 3), dtype=np.uint8)
        center = int(grid_px*0.5 + 0.5)

        cx = self.root[0]*grid_px + center
        cy = self.root[1]*grid_px + center
        cv2.circle(img, (cx, cy), grid_px, (0, 120, 120), -1)

        for e in self.edges:
            p1, p2 = e
            c1 = p1*grid_px + center
            c2 = p2*grid_px + center
            cv2.line(img, tuple(c1), tuple(c2), (0, 200, 0), edge_width)
        return img
    
    def add_edge_to_image(self, img, last_n_edge=1, grid_px=11, edge_width=5):
        n = min(last_n_edge, len(self.edges))
        if n == 0:
            return img
        w = self.W*grid_px
        h = self.H*grid_px
        center = int(grid_px*0.5 + 0.5)
        for e in self.edges[-n:]:
            p1, p2 = e
            c1 = p1*grid_px + center
            c2 = p2*grid_px + center
            cv2.line(img, tuple(c1), tuple(c2), (0, 200, 0), edge_width)
        return img

    def fill_grid(self):
        while not self.is_grown():
            self.grow()

def main():
    T = BTree(64, 64)
    while True:
        img = T.make_image()
        while not T.is_grown():
            img = T.add_edge_to_image(img)
            cv2.imshow("test", img)
            cv2.waitKey(5)
            T.grow()
        print("grown!")
        img = T.add_edge_to_image(img)
        cv2.imshow("test", img)
        cv2.waitKey(1000)
        T.reset()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()

maze.py

#!/usr/bin/python

import networkx as nx
import numpy as np
import cv2
from brownian_tree import BTree

def main():
    H = 64 
    W = 128
    T = BTree(H, W)

    img = T.make_image()
    while not T.is_grown():
        img = T.add_edge_to_image(img)
        cv2.imshow("test", img)
        cv2.waitKey(1)
        T.grow()

    G = nx.Graph()
    for e in T.edges:
        u1, u2 = e
        index1 = u1[1]*W + u1[0]
        index2 = u2[1]*W + u2[0]
        G.add_edge(index1, index2)

    route = nx.dijkstra_path(G, 0, W*H - 1)

    grid_px = 11
    edge_width = 5
    center = int(grid_px*0.5 + 0.5)


    route = [-W] + route + [W*H + W - 1]
    T.edges.append((np.array([0, -1], dtype=int), np.array([0, 0], dtype=int)))
    T.edges.append((np.array([W - 1, H - 1], dtype=int), np.array([W - 1, H], dtype=int)))

    img = T.add_edge_to_image(img, last_n_edge=3)

    for u1, u2 in zip(route[:-1], route[1:]):
        p1 = ((u1%W)*grid_px + center, int(u1/W)*grid_px + center)
        p2 = ((u2%W)*grid_px + center, int(u2/W)*grid_px + center)
        cv2.line(img, p1, p2, (200, 0, 0), edge_width - 2)
        cv2.imshow("test", img)
        cv2.waitKey(20)
        
    cv2.imshow("test", img)
    cv2.waitKey(10000)
    cv2.destroyAllWindows()
        

if __name__ == '__main__':
    while True:
        main()