def gen_img(graph: Graph, node_num, rm_ratio: float = 0.2, root_folder='train', source_folder='', source_file='train.txt'): """ 在graph上,以rm_ratio的概率随机删除一些链路,然后用随机选取的一对节点作为源宿点,判断其是否联通。并将生成的图像放到root_folder目录中, 对应的文件名称和标签放到source_folder目录下的train.txt文件中。 """ # 生成原始数据 rds = rand_edges(graph, rm_ratio) graph.remove_edges_from(rds) src, dst = gen_src_dst(0, node_num) label = get_label(src, dst, graph) # 构造graphviz对象 name = str(time.time()) g = Graph(format='jpeg', engine='neato') g.attr('node', shape='circle') g.attr('edge') for node in graph.nodes(): g.node(name=str(node)) g.node(name=str(src), shape='triangle') g.node(name=str(dst), shape='triangle') for edge in graph.edges(): g.edge(str(edge[0]), str(edge[1])) for edge in rds: g.edge(str(edge[0]), str(edge[1]), color='white') g.render(filename=name, directory=root_folder, view=False, cleanup=True) # 存储到文件中去 file = open(os.path.join(source_folder, source_file), 'a') file.write(name + '.jpeg ' + str(label) + '\n') file.flush() file.close()