python matplolib/networkx image to byteIO

python 으로 matplolib/networkx 를 이용하여 Graph 를 그리는 작업을 하다보면 굳이 image 파일로 저장해야 하는가에 대한 고민이 생겼다.

 

아래와 같은 코드를 이용하면 Node(A, B, C) 3개를 가지는 Graph 가 보여지고 test.png 파일로 저장 할 수 있습니다.

 

import networkx as nx                                                                                                                                     
import matplotlib.pyplot as plt                                                                                                                           
                                                                                                                                                            
G = nx.Graph()                                                                                                                                            
G.add_nodes_from([('A', {'weight':5}), ('B', {'weight':3}), ('C', {'weight':3})])                                                                         
G.add_edges_from([('A', 'B', {'weight':20})])                                                                                                             
G.add_edges_from([('A', 'C', {'weight':20})])                                                                                                             
                                                                                                                                                            
pos = nx.shell_layout(G)  # positions for all nodes                                                                                                       
                                                                                                                                                            
nx.draw(G, pos=pos, node_size=1000, with_labels=False)                                                                                                    
nx.draw_networkx_labels(G,pos=pos,font_size=30)                                                                                                           
                                                                                                                                                            
plt.savefig('test.png')                                                                                                                                   
plt.show()

 

이렇게 파일로 저장이 되는 것도 좋지만, 보통 이런 Graph 는 실시간으로 화면에 보여주기 위한 용도로 많이 사용됩니다. 이런 경우 "그래프 생성 → 그래프 이미지 파일 저장 → 그래프 이미지 파일 열기 → 그래프 이미지 파일 보여주기" 의 4 단계를 거치는 것은 비효율적으로 보입니다.

 

이런 비효율적 부분 중 파일 저장/읽기 단계를 “그래프 생성 → 그래프 이미지 스트림 생성→ 그래프 이미지 파일 보여주기” 단계로 변경 하면 개선이 가능 합니다. 바로 아래 코드 처럼 말입니다.

import matplotlib.pyplot as plt
import networkx as nx
from io import BytesIO
import base64
G = nx.Graph()
G.add_nodes_from([('A', {'weight':5}), ('B', {'weight':3}), ('C', {'weight':3})])
G.add_edges_from([('A', 'B', {'weight':20})])
G.add_edges_from([('A', 'C', {'weight':20})])

pos = nx.shell_layout(G)  # positions for all nodes

nx.draw(G, pos=pos, node_size=1000, with_labels=False)                                                                                                    
nx.draw_networkx_labels(G,pos=pos,font_size=30)

figfile = BytesIO()
plt.savefig(figfile, format='png')
figfile.seek(0)
figdata_png = base64.b64encode(figfile.getvalue())
plt.clf()

print(figdata_png)

 

이 코드를 이용하면 아래와 같이 그래프 이미지가 base64 로 인코딩된 문자열로로 생성이 됩니다.

 

이렇된 문자열을 img 태그에 넣으면 브라우저에서 바로 보여집니다.