Skeleton highway

world_maker/data_analysis.py

@@ -1,12 +1,10 @@
 import World
-from PIL import Image
-from PIL import ImageFilter
+from PIL import Image, ImageFilter
 import numpy as np
-import networkx as nx
 from scipy import ndimage
-from scipy.ndimage import gaussian_gradient_magnitude
-from scipy.ndimage import label
 from Skeleton import Skeleton
+from typing import Union
+
 
 def get_data(world: World):
     heightmap, watermap, treemap = world.getData()
@@ -16,7 +14,13 @@ def get_data(world: World):
     return heightmap, watermap, treemap
 
 
-def filter_inverse(image: Image) -> Image:
+def handle_import_image(image: Union[str, Image]) -> Image:
+    if isinstance(image, str):
+        return Image.open(image)
+    return image
+
+
+def filter_negative(image: Image) -> Image:
     """
     Invert the colors of an image.
 
@@ -26,7 +30,7 @@ def filter_inverse(image: Image) -> Image:
     return Image.fromarray(np.invert(np.array(image)))
 
 
-def filter_sobel(image) -> Image:
+def filter_sobel(image: Union[str, Image]) -> Image:
     """
     Edge detection algorithms from an image.
 
@@ -35,8 +39,7 @@ def filter_sobel(image) -> Image:
     """
 
     # Open the image
-    if isinstance(image, str):
-        image = Image.open(image).convert('RGB')
+    image = handle_import_image(image).convert('RGB')
 
     img = np.array(image).astype(np.uint8)
 
@@ -97,7 +100,7 @@ def filter_sobel(image) -> Image:
     return image
 
 
-def filter_smooth(image, radius: int = 3):
+def filter_smooth(image: Union[str, Image], radius: int = 3):
     """
     :param image: white and black image representing the derivative of the terrain (sobel), where black is flat and white is very steep.
     :param radius: Radius of the Gaussian blur.
@@ -106,8 +109,7 @@ def filter_smooth(image, radius: int = 3):
         image: black or white image, with black as flat areas to be skeletonized
     """
 
-    if isinstance(image, str):
-        image = Image.open(image)
+    image = handle_import_image(image)
 
     # image = image.filter(ImageFilter.SMOOTH_MORE)
     # image = image.filter(ImageFilter.SMOOTH_MORE)
@@ -128,7 +130,8 @@ def filter_smooth(image, radius: int = 3):
     return Image.fromarray(bool_array)
 
 
-def remove_water_from_map(image: Image) -> Image:
+def remove_water_from_map(image: Union[str, Image]) -> Image:
+    image = handle_import_image(image)
     watermap = Image.open('./data/watermap.png').convert('L')
 
     array_heightmap = np.array(image)
@@ -141,42 +144,66 @@ def remove_water_from_map(image: Image) -> Image:
     return result_image
 
 
-def group_map(image1: Image, image2: Image) -> Image:
+def group_map(image1: Union[str, Image], image2: Union[str, Image]) -> Image:
+    image1 = handle_import_image(image1)
+    image2 = handle_import_image(image2)
+
     array1 = np.array(image1)
     array2 = np.array(image2)
 
     mask = array1 == 255
     array2[mask] = 255
 
-    result_image = Image.fromarray(array2)
-    return result_image
+    return Image.fromarray(array2)
+
+
+def filter_smooth_array(array: np.ndarray, radius: int = 3) -> np.ndarray:
+    image = Image.fromarray(array)
+    smooth_image = filter_smooth(image, radius)
+    array = np.array(smooth_image)
+    return array
 
 
 def highway_map() -> Image:
     smooth_sobel = filter_smooth("./data/sobelmap.png", 1)
-    inverse_sobel = filter_inverse(smooth_sobel)
+    inverse_sobel = filter_negative(smooth_sobel)
     sobel_no_water = remove_water_from_map(inverse_sobel)
-    sobel_no_water.save("./data/test.png")
+    sobel_no_water.save("./data/negative_sobel_water_map.png")
     array = np.array(sobel_no_water)
     array = ndimage.binary_erosion(array, iterations=10)
     array = ndimage.binary_dilation(array, iterations=5)
-    image = Image.fromarray(array)
-    smooth_image = filter_smooth(image, 5)
-    array = np.array(smooth_image)
+    array = filter_smooth_array(array, 5)
     array = ndimage.binary_erosion(array, iterations=17)
-    image = Image.fromarray(array)
-    smooth_image = filter_smooth(image, 6)
-    array = np.array(smooth_image)
+    array = filter_smooth_array(array, 6)
     array = ndimage.binary_dilation(array, iterations=3)
     image = Image.fromarray(array)
     image.save('./data/highwaymap.png')
     return image
 
-def skeletonnize_map(map: Image):
-    skeleton = Skeleton()
-    image_array = np.array(map)
-    skeleton.setSkeleton(image_array)
-    skeleton.parseGraph()
-    heightmap_skeleton, roadsArea = skeleton.map()
-    heightmap_skeleton.save('./data/skeleton.png')
-    roadsArea.save('./data/roads.png')
+
+def create_volume(surface: np.ndarray, heightmap: np.ndarray, make_it_flat: bool = False) -> np.ndarray:
+    volume = np.full((len(surface), 255, len(surface[0])), False)
+    for z in range(len(surface)):
+        for x in range(len(surface[0])):
+            if not make_it_flat:
+                volume[x][heightmap[z][x]][z] = surface[z][x]
+            else:
+                volume[x][0][z] = surface[z][x]
+    return volume
+
+
+def convert_2D_to_3D(image: Union[str, Image], make_it_flat: bool = False) -> np.ndarray:
+    image = handle_import_image(image)
+    heightmap = Image.open('./data/heightmap.png').convert('L')
+    heightmap = np.array(heightmap)
+    surface = np.array(image)
+    volume = create_volume(surface, heightmap, make_it_flat)
+    return volume
+
+
+def skeleton_highway_map(map: Union[str, Image]):
+    image_array = convert_2D_to_3D(map, True)
+    skeleton = Skeleton(image_array)
+    skeleton.parse_graph(True)
+    heightmap_skeleton = skeleton.map()
+    heightmap_skeleton.save('./data/skeleton_highway.png')