framework/freecad_workbench/freecad/robossembler/assembly_graph_generation.py

42 lines
1.1 KiB
Python

import json
import networkx as nx
import matplotlib.pyplot as plt
from PIL import Image
import matplotlib.image as mpimg
from networkx.drawing.nx_agraph import graphviz_layout
def load_assembly_sequence(filepath):
with open(filepath, 'r') as file:
return json.load(file)
def create_assembly_graph(sequence):
G = nx.DiGraph()
for i, part in enumerate(sequence):
G.add_node(part, image=f'{part}.svg')
if i > 0:
G.add_edge(sequence[i-1], part)
return G
def draw_graph(G):
pos = graphviz_layout(G, prog='dot')
fig, ax = plt.subplots(figsize=(12, 8))
for node in G.nodes(data=True):
x, y = pos[node[0]]
img = mpimg.imread(node[1]['image'])
ax.imshow(img, aspect='auto', extent=(x-0.5, x+0.5, y-0.5, y+0.5), zorder=1)
ax.text(x, y-0.6, s=node[0], horizontalalignment='center', fontsize=12, zorder=2)
nx.draw_networkx_edges(G, pos, ax=ax, arrows=True)
plt.axis('off')
plt.show()
def main():
sequence = load_assembly_sequence('assembly_sequence.json')
assembly_graph = create_assembly_graph(sequence)
draw_graph(assembly_graph)