def create_graph(s_objects, p_objects):
print 'Creating graph ...'
graph = Graph()
node_id = 0 # use this as node id as some segments might have similar id
# add the imp object nodes
for s_obj in s_objects:
node = Node(node_id, \
seg_id=s_obj.id, \
x_pos=s_obj.pos[1], \
y_pos=s_obj.pos[0], \
b_fix=True, \
size=s_obj.size, \
color_energy=s_obj.color_energy, \
color=s_obj.color, \
saliency=s_obj.saliency, \
radius=np.sqrt(s_obj.size)/4)
graph.add_node(node)
node_id += 1
max_snode_id = node_id
# add the people node
for p_obj in p_objects:
node = Node(node_id,\
seg_id=p_obj.id, \
x_pos=p_obj.pos[1], \
y_pos=p_obj.pos[0], \
b_fix=False, \
size=p_obj.size, \
color_energy=1.2*p_obj.color_energy, \
color=p_obj.color, \
saliency=1.2*p_obj.saliency, \
radius=np.sqrt(p_obj.size)/4)
graph.add_node(node)
# add this node to the list of people nodes
graph.add_pnode(node)
graph.num_people += 1
# create edges from the people node to rest of nodes
for dst_id in range(max_snode_id):
src_id = node_id
src_node = node
dst_node = graph.get_node(dst_id)
assert src_node is not None
assert dst_node is not None
graph.add_edge(src_node, dst_node, spring_length=0.0)
node_id += 1
return graph
