-
Notifications
You must be signed in to change notification settings - Fork 0
/
render_tree.py
142 lines (123 loc) · 6.29 KB
/
render_tree.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
# render_tree.py
from math2d_vector import Vector
from math2d_aa_rect import AxisAlignedRectangle
from math2d_affine_transform import AffineTransform
class RenderNode(object):
def __init__(self, person):
super().__init__()
self.person = person
self.sub_node_map = {}
self.label_string = person.name + '\n' + person.family_search_id
self.label_box = AxisAlignedRectangle()
self.bounding_box = None
def calculate_size(self):
size = 0
for node in self.all_nodes():
size += 1
return size
def render_graph(self, draw, image, font, person_subset):
image_rect = AxisAlignedRectangle(Vector(0.0, 0.0), Vector(float(image.width), float(image.height)))
world_rect = self.bounding_box.Copy()
world_rect.ExpandToMatchAspectRatioOf(image_rect)
for node in self.all_nodes():
node.render_edges(draw, font, image_rect, world_rect)
for node in self.all_nodes():
node.render_label_box(draw, font, image_rect, world_rect, person_subset)
def render_edges(self, draw, font, image_rect, world_rect):
point_a = world_rect.Map(self.label_box.Center(), image_rect)
for key in self.sub_node_map:
sub_node = self.sub_node_map[key]
point_b = world_rect.Map(sub_node.label_box.Center(), image_rect)
draw.line((point_a.x, point_a.y, point_b.x, point_b.y), fill=(50, 100, 100), width=1)
point_c = (point_a + point_b) * 0.5
text_size = draw.textsize(text=key, font=font)
draw.text((point_c.x - text_size[0] / 2, point_c.y - text_size[1] / 2), text=key, font=font, fill=(0, 0, 0))
def render_label_box(self, draw, font, image_rect, world_rect, person_subset):
point_a = world_rect.Map(self.label_box.min_point, image_rect)
point_b = world_rect.Map(self.label_box.max_point, image_rect)
color = (80, 32, 32) if self.person in person_subset else (32, 32, 32)
if self.person.any_proxy_work_available:
color = (32, 80, 23)
draw.rectangle((point_a.x, point_a.y, point_b.x, point_b.y), fill=color)
polygon = self.label_box.GeneratePolygon()
for world_line in polygon.GenerateLineSegments():
image_line = world_rect.Map(world_line, image_rect)
draw.line((image_line.point_a.x, image_line.point_a.y, image_line.point_b.x, image_line.point_b.y), fill=(255, 0, 0), width=1)
point = world_rect.Map(self.label_box.Center(), image_rect)
text_size = draw.textsize(text=self.label_string, font=font)
draw.text((point.x - text_size[0] / 2, point.y - text_size[1] / 2), text=self.label_string, font=font, fill=(200, 200, 200))
def transform_graph_layout(self, transform):
for node in self.all_nodes():
node.label_box = transform.Transform(node.label_box)
node.bounding_box = transform.Transform(node.bounding_box)
def calculate_graph_layout(self, draw, font):
text_size = draw.textsize(text=self.label_string, font=font)
text_width = text_size[0]
text_height = text_size[1]
total_width = 0.0
margin = 2.0
for key in self.sub_node_map:
sub_node = self.sub_node_map[key]
sub_node.calculate_graph_layout(draw, font)
sub_node.calculate_bounding_box()
total_width += sub_node.bounding_box.Width() + 2.0 * margin
location = Vector(-total_width / 2.0 + margin, -2.0 * text_height)
for key in self.sub_node_map:
sub_node = self.sub_node_map[key]
transform = AffineTransform()
upper_left_corner = Vector(sub_node.bounding_box.min_point.x, sub_node.bounding_box.max_point.y)
transform.Translation(location - upper_left_corner)
sub_node.transform_graph_layout(transform)
location.x += sub_node.bounding_box.Width() + 2.0 * margin
self.label_box.min_point = Vector(-text_width / 2.0 - 3.0, -text_height / 2.0 - 3.0)
self.label_box.max_point = Vector(text_width / 2.0 + 3.0, text_height / 2.0 + 3.0)
def calculate_bounding_box(self):
self.bounding_box = self.label_box.Copy()
for node in self.all_nodes():
self.bounding_box.GrowFor(node.label_box)
def all_nodes(self):
queue = [self]
while len(queue) > 0:
node = queue.pop()
yield node
for key in node.sub_node_map:
sub_node = node.sub_node_map[key]
queue.append(sub_node)
def prune_tree(self, person_set):
del_key_list = []
for key in self.sub_node_map:
sub_node = self.sub_node_map[key]
if not sub_node.any_person_found_in(person_set):
del_key_list.append(key)
for key in del_key_list:
del self.sub_node_map[key]
for key in self.sub_node_map:
sub_node = self.sub_node_map[key]
sub_node.prune_tree(person_set)
def any_person_found_in(self, person_set):
for node in self.all_nodes():
if node.person in person_set:
return True
# TODO: Write an optimizer for the tree. E.g., your father's spouse's son is just your brother.
# Note that some people can't be removed from the tree if they fall within a given set.
# I.e., the set of people we want to see in the tree.
def construct_using_path(self, path, i=0):
if i < len(path):
component = path[i]
if component[0] == 'mother':
next_person = self.person.mother
key = 'Mother'
elif component[0] == 'father':
next_person = self.person.father
key = 'Father'
elif component[0] == 'child':
next_person = self.person.child_list[component[1]]
key = 'Child %d' % (component[1] + 1)
elif component[0] == 'spouse':
next_person = self.person.spouse_list[component[1]]
key = 'Spouse %d' % (component[1] + 1)
else:
raise Exception('Unknown component: %s' % component[0])
if key not in self.sub_node_map:
self.sub_node_map[key] = RenderNode(person=next_person)
self.sub_node_map[key].construct_using_path(path, i + 1)