def show_only_true_colour(screen, state: VisState, settings: VisSettings) -> Generator:
state.pygame_img = settings.image_loader_func(
state.points[0][0],
state.points[0][1],
state.points[2][0],
state.points[2][1],
)
state.pygame_img = pygame.transform.scale(state.pygame_img, settings.screen_size)
screen.blit(state.pygame_img, (0, 0))
yield
def show_only_heights(screen, state: VisState, settings: VisSettings) -> Generator:
height_array = state.selected_area_height_map - state.selected_area_height_map.min()
height_array = (height_array // (height_array.max() / 255)).astype("int32")
image = Image.fromarray(numpy.uint8(cm.gist_earth(height_array) * 255))
state.pygame_img = pygame.image.fromstring(image.tobytes(), image.size, image.mode)
state.pygame_img = pygame.transform.scale(
state.pygame_img, [int(state.pygame_img.get_rect().size[i] * state.float_pixel_size[i]) for i in (0, 1)]
)
screen.blit(state.pygame_img, (0, 0))
yield
def add_circles(screen, state: VisState, settings: VisSettings) -> Generator:
state.circles_surface = pygame.Surface(settings.screen_size, pygame.SRCALPHA, 32).convert_alpha()
_draw_circles(state.circles_surface, state, settings)
# Remove background image
num_steps = 40
state.resized_selected_surface = state.resized_selected_surface.convert()
for i in range(num_steps, -1, -1):
image_alpha = int(255 * i / num_steps)
state.resized_selected_surface.set_alpha(image_alpha)
screen.fill((0, 0, 0))
screen.blit(state.resized_selected_surface, (0, 0))
screen.blit(state.circles_surface, (0, 0))
yield
def scale_up_selection(screen, state: VisState, settings: VisSettings) -> Generator:
# Get the section of the image
# Scale up the section step by step
selected_surface = settings.full_size_height_image.subsurface(
pygame.Rect(
state.points[0][0],
state.points[0][1],
state.points[2][0] - state.points[0][0],
state.points[2][1] - state.points[0][1],
)
)
num_steps = 5
for i in range(num_steps + 1):
# Smoothly transition the scale of the selected surface to cover the whole screen
# TODO: change this section so that it looks like a flat surface moving closer to an observer
proportion_finished = (i / num_steps) ** 3
proportion_unfinished = 1 - proportion_finished
width = int(
state.selection_pixel_size[0] * settings.scale_ratio * proportion_unfinished
+ settings.screen_size[0] * proportion_finished
)
height = int(
state.selection_pixel_size[1] * settings.scale_ratio * proportion_unfinished
+ settings.screen_size[1] * proportion_finished
)
left = int(state.scaled_points[0][0] * proportion_unfinished)
top = int(state.scaled_points[0][1] * proportion_unfinished)
state.resized_selected_surface = pygame.transform.scale(selected_surface, (width, height))
screen.blit(state.resized_selected_surface, (left, top))
yield
def animate_watershed(event, screen, state: VisState,
settings: VisSettings) -> Generator:
if event.type == pygame.MOUSEBUTTONDOWN:
source = find_clicked_node(pygame.mouse.get_pos(), state)
elif event.type == pygame.KEYDOWN and event.key == pygame.K_RETURN:
source = None
else:
return
circle_radius = int(max(*state.float_pixel_size) * 0.35)
screen.blit(state.pygame_img, (0, 0))
state.node_flows, state.edge_flows = calculate_watershed(state,
source=source)
state.node_flow_items = list(state.node_flows.items())
node_flow_indexes = {
state.node_flow_items[i][0]: i
for i in range(len(state.node_flow_items))
}
j = 0
for i in range(len(state.node_flow_items)):
node, flow = state.node_flow_items[i]
if flow > 0.0001:
new_location = get_node_centerpoint(node)
for neighbour in state.graph[node]:
if neighbour in node_flow_indexes and node_flow_indexes[
neighbour] > i:
neighbour_location = get_node_centerpoint(neighbour)
_draw_line(screen, new_location, neighbour_location, state)
j += 1
circle_center = [
int(new_location[i] * state.float_pixel_size[i] +
state.center_offset[i]) for i in (0, 1)
]
colour = [i * 255 for i in cm.gist_heat(flow)[:3]]
pygame.draw.circle(screen, colour, circle_center, circle_radius)
if j % 20 == 0:
yield
yield
def highlight_low_nodes(screen, state: VisState, settings: VisSettings) -> Generator:
circle_radius = int(max(*state.float_pixel_size) * 0.37)
print(f"circle_radius = {circle_radius}")
state.low_nodes = find_low_nodes(state.graph, state)
print(f"Found {len(state.low_nodes)} low nodes")
for low_node in state.low_nodes:
new_location = get_node_centerpoint(low_node)
pygame.draw.circle(
screen,
(255, 0, 0),
(
int(new_location[0] * state.float_pixel_size[0] + state.center_offset[0]),
int(new_location[1] * state.float_pixel_size[1] + state.center_offset[1]),
),
circle_radius,
min(circle_radius, 3),
)
state.low_nodes = sorted(state.low_nodes, key=lambda key: get_height_by_key(key, state))
yield
def random_selection_polygon(screen, state: VisState, settings: VisSettings) -> Generator:
if state.click_location_1 and state.click_location_2:
return
shown_screen_dimensions = [int(i / settings.scale_ratio) for i in settings.screen_size]
state.selection_pixel_size = _compute_selection_pixel_size(settings.screen_size, settings.max_pixels)
num_tries = 0
highest_point_altitude = 0
while num_tries < 10 and highest_point_altitude <= 0:
left = randint(0, shown_screen_dimensions[0] - 1 - state.selection_pixel_size[0])
right = left + state.selection_pixel_size[0]
top = randint(0, shown_screen_dimensions[1] - 1 - state.selection_pixel_size[1])
bottom = top + state.selection_pixel_size[1]
state.points = ((left, top), (right, top), (right, bottom), (left, bottom))
state.selected_area_height_map = settings.height_map[
state.points[0][1] : state.points[0][1] + state.selection_pixel_size[1],
state.points[0][0] : state.points[0][0] + state.selection_pixel_size[0],
]
highest_point_altitude = state.selected_area_height_map.max()
state.scaled_points = [(int(x * settings.scale_ratio), int(y * settings.scale_ratio)) for x, y in state.points]
state.float_pixel_size = (
settings.screen_size[0] / state.selection_pixel_size[0],
settings.screen_size[1] / state.selection_pixel_size[1],
)
state.center_offset = (state.float_pixel_size[0] / 2, state.float_pixel_size[1] / 2)
pygame.draw.polygon(
settings.screen_size_height_image,
settings.selection_line_colour,
state.scaled_points,
settings.selection_line_width,
)
screen.blit(settings.screen_size_height_image, (0, 0))
yield
def flood_points(screen, state: VisState, settings: VisSettings) -> Generator:
"""
Highlight all the neighbours (and add them to a priority queue).
Loop through the priority queue;
Check if the node is lower than the lake
Merge the node with the lake
Add the other neighbours to the priority queue
"""
def does_node_touch_border(node):
if node[0] == 0:
return True
if node[1] == 0:
return True
if node[0] == state.selection_pixel_size[0] - 1:
return True
if node[1] == state.selection_pixel_size[1] - 1:
return True
return False
circle_radius = int(max(*state.float_pixel_size) * 0.35)
if not state.low_nodes:
yield
for low_node in state.low_nodes:
if low_node not in state.graph:
continue
lake_height = get_height_by_key(low_node, state)
for neighbour in state.graph[low_node]:
if get_height_by_key(neighbour, state) < lake_height:
continue
queue = [(lake_height, low_node)]
nodes_in_queue = {low_node}
merging_nodes = {low_node}
while True:
try:
node_height, node = heapq.heappop(queue)
except IndexError:
print("heap ran out of items but it shouldn't")
break
new_location = get_node_centerpoint(node)
if node_height < lake_height:
pygame.draw.circle(
screen,
(0, 255, 0),
(
int(new_location[0] * state.float_pixel_size[0] + state.center_offset[0]),
int(new_location[1] * state.float_pixel_size[1] + state.center_offset[1]),
),
circle_radius,
min(circle_radius, 3),
)
break
lake_height = node_height
merging_nodes.add(node)
# If node is a border then this means the flow can go off the edge. merging should stop after this node
if any(does_node_touch_border(i) for i in node):
pygame.draw.circle(
screen,
(0, 255, 0),
(
int(new_location[0] * state.float_pixel_size[0] + state.center_offset[0]),
int(new_location[1] * state.float_pixel_size[1] + state.center_offset[1]),
),
circle_radius,
min(circle_radius, 3),
)
break
pygame.draw.circle(
screen,
(255, 0, 0),
(
int(new_location[0] * state.float_pixel_size[0] + state.center_offset[0]),
int(new_location[1] * state.float_pixel_size[1] + state.center_offset[1]),
),
circle_radius,
min(circle_radius, 3),
)
for adjacent_node in state.graph[node]:
if adjacent_node not in nodes_in_queue:
new_location = get_node_centerpoint(adjacent_node)
pygame.draw.circle(
screen,
(255, 165, 0),
(
int(new_location[0] * state.float_pixel_size[0] + state.center_offset[0]),
int(new_location[1] * state.float_pixel_size[1] + state.center_offset[1]),
),
circle_radius,
min(circle_radius, 3),
)
nodes_in_queue.add(adjacent_node)
heapq.heappush(queue, (get_height_by_key(adjacent_node, state), adjacent_node))
yield
# Add all equal height nodes
while True:
try:
node_height, node = heapq.heappop(queue)
except IndexError:
break
# Don't merge the outflow points to the lake
if node_height < lake_height:
continue
elif node_height == lake_height:
# Merge this
merging_nodes.add(node)
new_location = get_node_centerpoint(node)
pygame.draw.circle(
screen,
(255, 0, 0),
(
int(new_location[0] * state.float_pixel_size[0] + state.center_offset[0]),
int(new_location[1] * state.float_pixel_size[1] + state.center_offset[1]),
),
circle_radius,
min(circle_radius, 3),
)
for adjacent_node in state.graph[node]:
if adjacent_node not in nodes_in_queue:
new_location = get_node_centerpoint(adjacent_node)
pygame.draw.circle(
screen,
(255, 165, 0),
(
int(new_location[0] * state.float_pixel_size[0] + state.center_offset[0]),
int(new_location[1] * state.float_pixel_size[1] + state.center_offset[1]),
),
circle_radius,
min(circle_radius, 3),
)
nodes_in_queue.add(adjacent_node)
heapq.heappush(queue, (get_height_by_key(adjacent_node, state), adjacent_node))
yield
else:
break
merged_node_key = tuple(sorted({node for node_key in merging_nodes for node in node_key}))
neighbours = {node for merging_node in merging_nodes for node in state.graph[merging_node]} - set(
merging_nodes
)
for neighbour in neighbours:
updated_neighbours = set(state.graph[neighbour]) - merging_nodes
updated_neighbours.add(merged_node_key)
state.graph[neighbour] = tuple(sorted(updated_neighbours))
state.graph[merged_node_key] = tuple(sorted(neighbours))
state.selected_area_height_map[merged_node_key[0][1], merged_node_key[0][0]] = lake_height
# Find all the nodes don't touch the merged area
untouched_nodes = state.graph.keys() - {merged_node_key} - set(merging_nodes)
untouched_surface = pygame.Surface(settings.screen_size, pygame.SRCALPHA, 32).convert_alpha()
height_array = state.selected_area_height_map - state.selected_area_height_map.min()
height_array = (height_array // (height_array.max() / 255)).astype("int32")
for node in untouched_nodes:
new_location = get_node_centerpoint(node)
for neighbour in state.graph[node]:
if neighbour < node:
neighbour_location = get_node_centerpoint(neighbour)
_draw_line(untouched_surface, new_location, neighbour_location, state)
for node in untouched_nodes:
new_location = get_node_centerpoint(node)
height = height_array[node[0][1], node[0][0]]
colour = get_colour_by_height(height)
pygame.draw.circle(
untouched_surface,
colour,
(
int(new_location[0] * state.float_pixel_size[0] + state.center_offset[0]),
int(new_location[1] * state.float_pixel_size[1] + state.center_offset[1]),
),
circle_radius,
)
if node in state.low_nodes:
pygame.draw.circle(
untouched_surface,
(255, 0, 0),
(
int(new_location[0] * state.float_pixel_size[0] + state.center_offset[0]),
int(new_location[1] * state.float_pixel_size[1] + state.center_offset[1]),
),
circle_radius,
min(circle_radius, 3),
)
num_steps = 60
new_location = get_node_centerpoint(merged_node_key)
for i in range(1, num_steps + 1):
screen.fill((0, 0, 0))
screen.blit(untouched_surface, (0, 0))
for original_node in merging_nodes:
original_node_position = (
sum(x for x, y in original_node) / len(original_node),
sum(y for x, y in original_node) / len(original_node),
)
actual_node_position = (
(new_location[0] - original_node_position[0]) * i / num_steps + original_node_position[0],
(new_location[1] - original_node_position[1]) * i / num_steps + original_node_position[1],
)
for neighbour in state.graph[original_node]:
if neighbour not in merging_nodes:
neighbour_position = get_node_centerpoint(neighbour)
# Move edge from to position
_draw_line(screen, neighbour_position, actual_node_position, state)
height = height_array[neighbour[0][1], neighbour[0][0]]
colour = get_colour_by_height(height)
pygame.draw.circle(
screen,
colour,
(
int(neighbour_position[0] * state.float_pixel_size[0] + state.center_offset[0]),
int(neighbour_position[1] * state.float_pixel_size[1] + state.center_offset[1]),
),
circle_radius,
)
height = height_array[original_node[0][1], original_node[0][0]]
colour = get_colour_by_height(height)
pygame.draw.circle(
screen,
colour,
(
int(actual_node_position[0] * state.float_pixel_size[0] + state.center_offset[0]),
int(actual_node_position[1] * state.float_pixel_size[1] + state.center_offset[1]),
),
circle_radius,
)
yield
for merging_node in merging_nodes:
del state.graph[merging_node]
def merge_equal_height_nodes(screen, state: VisState, settings: VisSettings) -> Generator:
# 1. Loop through all pixels
# 2. DFS to any adjancent nodes of equal height
# 3. Store co-ordinates of points to be merged
# create a surface with the circles which don't move
# for every point which moved calculate it's start and end position
# then render frames 1 by one by redrawing all the circles which move, and the connections to those points
node_movements, skip_nodes, node_merge_operations = equal_height_node_merge(state, settings)
state.circles_surface = pygame.Surface(settings.screen_size, pygame.SRCALPHA, 32).convert_alpha()
non_skip_nodes = [
(x, y)
for x in range(state.selection_pixel_size[0])
for y in range(state.selection_pixel_size[1])
if (x, y) not in skip_nodes
]
state.graph = create_graph(node_merge_operations, skip_nodes, non_skip_nodes, state)
non_skip_nodes_set = set(non_skip_nodes)
for from_node in non_skip_nodes_set:
if (from_node[0] + 1, from_node[1]) in non_skip_nodes_set:
_draw_line(state.circles_surface, from_node, (from_node[0] + 1, from_node[1]), state)
if (from_node[0], from_node[1] + 1) in non_skip_nodes_set:
_draw_line(state.circles_surface, from_node, (from_node[0], from_node[1] + 1), state)
_draw_circles(state.circles_surface, state, settings, skip_nodes, absolute_scale=False)
num_steps = 60
circle_radius = int(max(*state.float_pixel_size) * 0.35)
height_array = state.selected_area_height_map - state.selected_area_height_map.min()
height_array = (height_array // (height_array.max() / 255)).astype("int32")
def get_updated_node_position(node, new_position, progress):
x = node[0] + i / num_steps * (new_position[0] - node[0])
y = node[1] + i / num_steps * (new_position[1] - node[1])
return x, y
for i in range(1, num_steps + 1):
screen.fill((0, 0, 0))
moving_circles_surface = pygame.Surface(settings.screen_size, pygame.SRCALPHA, 32).convert_alpha()
for node, new_position in node_movements.items():
adjacent_nodes = _get_adjacent_nodes(node, state, lambda x, y: x < y or x in skip_nodes)
for adjacent_node in adjacent_nodes:
if adjacent_node in skip_nodes:
adjacent_node_position = get_updated_node_position(
adjacent_node, node_movements[adjacent_node], i / num_steps
)
else:
adjacent_node_position = adjacent_node
_draw_line(
moving_circles_surface,
get_updated_node_position(node, new_position, i / num_steps),
adjacent_node_position,
state,
)
for node, new_position in node_movements.items():
current_x, current_y = get_updated_node_position(node, new_position, i / num_steps)
colour = get_colour_by_height(height_array[node[1], node[0]])
pygame.draw.circle(
moving_circles_surface,
colour,
(
int(current_x * state.float_pixel_size[0] + state.center_offset[0]),
int(current_y * state.float_pixel_size[1] + state.center_offset[1]),
),
circle_radius,
)
screen.blit(moving_circles_surface, (0, 0))
screen.blit(state.circles_surface, (0, 0))
yield
def show_selection_polygon(event, screen, state: VisState,
settings: VisSettings) -> Generator:
if event.type == pygame.MOUSEBUTTONDOWN:
state.click_location_1 = pygame.mouse.get_pos()
elif event.type == pygame.MOUSEBUTTONUP:
state.click_location_2 = pygame.mouse.get_pos()
if state.click_location_1 == state.click_location_2:
state.click_location_1 = None
state.click_location_2 = None
else:
shown_screen_dimensions = [
int(i / settings.scale_ratio) for i in settings.screen_size
]
scaled_location_1 = (
int(state.click_location_1[0] / settings.scale_ratio),
int(state.click_location_1[1] / settings.scale_ratio),
)
scaled_location_2 = (
int(state.click_location_2[0] / settings.scale_ratio),
int(state.click_location_2[1] / settings.scale_ratio),
)
left = min(scaled_location_1[0], scaled_location_2[0])
right = max(scaled_location_1[0], scaled_location_2[0])
top = min(scaled_location_1[1], scaled_location_2[1])
bottom = max(scaled_location_1[1], scaled_location_2[1])
screen_fill_ratio_x = (right - left) / shown_screen_dimensions[0]
screen_fill_ratio_y = (bottom - top) / shown_screen_dimensions[1]
screen_fill_ratio = min(screen_fill_ratio_x, screen_fill_ratio_y)
if screen_fill_ratio_x < screen_fill_ratio_y:
bottom = top + screen_fill_ratio_x * shown_screen_dimensions[1]
else:
right = left + screen_fill_ratio_y * shown_screen_dimensions[0]
state.selection_pixel_size = _compute_selection_pixel_size(
settings.screen_size, max(right - left, bottom - top))
print(f"Selected area size: {state.selection_pixel_size}")
state.points = ((left, top), (right, top), (right, bottom),
(left, bottom))
state.selected_area_height_map = settings.height_map[
state.points[0][1]:state.points[0][1] +
state.selection_pixel_size[1],
state.points[0][0]:state.points[0][0] +
state.selection_pixel_size[0], ]
state.scaled_points = [(int(x * settings.scale_ratio),
int(y * settings.scale_ratio))
for x, y in state.points]
state.float_pixel_size = (
settings.screen_size[0] / state.selection_pixel_size[0],
settings.screen_size[1] / state.selection_pixel_size[1],
)
state.center_offset = (state.float_pixel_size[0] / 2,
state.float_pixel_size[1] / 2)
pygame.draw.polygon(
settings.screen_size_height_image,
settings.selection_line_colour,
state.scaled_points,
settings.selection_line_width,
)
screen.blit(settings.screen_size_height_image, (0, 0))
yield