208 lines
7.9 KiB
Python
208 lines
7.9 KiB
Python
from world_maker.District import District, Road
|
|
from world_maker.Position import Position
|
|
from PIL import Image
|
|
from random import randint
|
|
from world_maker.data_analysis import handle_import_image, detect_mountain
|
|
from typing import Union
|
|
import numpy as np
|
|
|
|
|
|
class City:
|
|
"""
|
|
Attributes:
|
|
districts (list): The list of districts in the city.
|
|
map_data (list): The 2D list representing the map of the city.
|
|
height_map (list): The 2D list representing the height map of the city.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""
|
|
The constructor for the City class.
|
|
"""
|
|
self.districts = []
|
|
self.map_data = []
|
|
self.height_map = []
|
|
self.init_maps()
|
|
|
|
def init_maps(self):
|
|
"""
|
|
Initialize the maps of the city. It reads the heightmap and watermap images and converts them into 2D lists.
|
|
"""
|
|
heightmap = Image.open('./world_maker/data/heightmap.png').convert('L')
|
|
watermap = Image.open('./world_maker/data/watermap.png').convert('L')
|
|
width, height = heightmap.size
|
|
self.map_data = [[-1 if watermap.getpixel((x, y)) > 0 else 0 for x in range(width)] for y in range(height)]
|
|
self.height_map = [[heightmap.getpixel((x, y)) for x in range(width)] for y in range(height)]
|
|
watermap.close()
|
|
heightmap.close()
|
|
|
|
def add_district(self, center: Position, district_type: str = ""):
|
|
"""
|
|
Add a new district to the city.
|
|
|
|
:param district_type:
|
|
:param center: The center position of the new district.
|
|
"""
|
|
self.districts.append(District(len(self.districts) + 1, center, district_type))
|
|
self.map_data[center.y][center.x] = len(self.districts)
|
|
|
|
def is_expend_finished(self):
|
|
"""
|
|
Check if the expansion of all districts in the city is finished.
|
|
|
|
:return: True if the expansion is finished, False otherwise.
|
|
"""
|
|
for district in self.districts:
|
|
if len(district.area_expend_from_point) > 0:
|
|
return False
|
|
return True
|
|
|
|
def choose_expend_point(self, point: Position, index_district: int):
|
|
"""
|
|
Choose a point to expand a district based on the distance between the center of the district and the point itself.
|
|
|
|
:param point: The point to be expanded.
|
|
:param index_district: The index of the district to be expanded.
|
|
"""
|
|
min_distance = point.distance_to(self.districts[index_district].center_expend)
|
|
index_district_chosen = index_district
|
|
for index in range(index_district + 1, len(self.districts)):
|
|
if point in self.districts[index].area_expend:
|
|
distance = point.distance_to(self.districts[index].center_expend)
|
|
if distance < min_distance:
|
|
min_distance = distance
|
|
self.districts[index_district_chosen].area_expend.remove(point)
|
|
index_district_chosen = index
|
|
else:
|
|
self.districts[index].area_expend.remove(point)
|
|
self.districts[index_district_chosen].area_expend_from_point.append(point)
|
|
self.districts[index_district_chosen].area_expend.remove(point)
|
|
self.map_data[point.y][point.x] = index_district_chosen + 1
|
|
|
|
def update_expend_district(self):
|
|
"""
|
|
Update the expansion points of all districts in the city.
|
|
"""
|
|
for district in self.districts:
|
|
if len(district.area_expend_from_point) > 0:
|
|
district.update_expend_points(district.area_expend_from_point[0], self.map_data, self.height_map)
|
|
for district in self.districts:
|
|
for point in district.area_expend:
|
|
self.choose_expend_point(point, district.tile_id - 1)
|
|
|
|
def loop_expend_district(self):
|
|
"""
|
|
Loop the expansion of all districts in the city until all districts are fully expanded.
|
|
"""
|
|
print("[City] Start expanding districts...")
|
|
while not self.is_expend_finished():
|
|
self.update_expend_district()
|
|
print("[City] Finished expanding districts.")
|
|
|
|
def district_draw_map(self):
|
|
"""
|
|
Draw the map of the city with different colors for each district.
|
|
"""
|
|
width, height = len(self.map_data[0]), len(self.map_data)
|
|
img = Image.new('RGB', (width, height))
|
|
colors = {id_district: (randint(0, 255), randint(0, 255), randint(0, 255))
|
|
for id_district in range(1, len(self.districts) + 1)}
|
|
|
|
for y in range(height):
|
|
for x in range(width):
|
|
if self.map_data[y][x] <= 0:
|
|
img.putpixel((x, y), (0, 0, 0))
|
|
else:
|
|
img.putpixel((x, y), colors[self.map_data[y][x]])
|
|
|
|
img.save('./world_maker/data/district.png')
|
|
print("[City] District map created.")
|
|
|
|
def draw_roads(self, size_road: int = 1) -> Image:
|
|
"""
|
|
Draw the roads of the city on the image.
|
|
|
|
:param size:
|
|
"""
|
|
image = Image.new('RGB', Image.open('./world_maker/data/heightmap.png').size)
|
|
for district in self.districts:
|
|
district.draw_roads(image, size_road)
|
|
return image
|
|
|
|
def district_generate_road(self) -> list[Road]:
|
|
"""
|
|
Generate the roads of the city for each district.
|
|
|
|
:return: The list of roads of the city.
|
|
"""
|
|
roads = []
|
|
for district in self.districts:
|
|
if district.type != "mountain":
|
|
district.generate_roads(self.map_data)
|
|
roads.extend(district.roads)
|
|
return roads
|
|
|
|
def point_in_which_district(self, point: Union[Position, tuple[int, int]]) -> int:
|
|
"""
|
|
Get the index of the district in which the point is located.
|
|
|
|
:param point: The point to check.
|
|
:return: The index of the district in which the point is located.
|
|
"""
|
|
if isinstance(point, Position):
|
|
point = (point.x, point.y)
|
|
return self.map_data[point[1]][point[0]]
|
|
|
|
def get_district_mountain_map(self) -> Image:
|
|
"""
|
|
Get the map of a district.
|
|
|
|
:param district_id: The id of the district.
|
|
:return: The map of the district.
|
|
"""
|
|
district_id = [district.tile_id for district in self.districts if district.type == "mountain"]
|
|
array = np.array([[True if self.map_data[y][x] in district_id else False for x in range(len(self.map_data[0]))]
|
|
for y in range(len(self.map_data))])
|
|
image = Image.fromarray(array)
|
|
image.save('./world_maker/data/mountain_map.png')
|
|
return image
|
|
|
|
def generate_district(self):
|
|
image = handle_import_image('./world_maker/data/smooth_sobel_watermap.png').convert('L')
|
|
array = np.array(image)
|
|
mountain_coo = detect_mountain()
|
|
self.add_district(Position(mountain_coo[0], mountain_coo[1]), "mountain")
|
|
print("[City] District added.")
|
|
remove_circle_data(array, mountain_coo)
|
|
area = get_area_array(array)
|
|
sizeX, sizeY = len(array[0]), len(array)
|
|
while area > sizeX * sizeY * 0.1:
|
|
x, y = randint(0, sizeX - 1), randint(0, sizeY - 1)
|
|
if array[y][x]:
|
|
self.add_district(Position(x, y))
|
|
remove_circle_data(array, (x, y))
|
|
area = get_area_array(array)
|
|
print("[City] District added.")
|
|
|
|
|
|
def remove_circle_data(array, center, radius=100):
|
|
y_indices, x_indices = np.indices(array.shape)
|
|
dist_sq = (y_indices - center[1]) ** 2 + (x_indices - center[0]) ** 2
|
|
mask = dist_sq <= radius ** 2
|
|
array[mask] = False
|
|
|
|
|
|
def get_area_array(array) -> int:
|
|
return np.sum(array)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
city = City()
|
|
for i in range(10):
|
|
city.add_district(Position(randint(0, 400), randint(0, 400)))
|
|
city.loop_expend_district()
|
|
city.district_draw_map()
|
|
city.district_generate_road()
|
|
image = city.draw_roads(Image.new('RGB', (401, 401)), 4)
|
|
image.save('./world_maker/data/roadmap.png')
|