Coverage for mlprodict/plotting/plotting_onnx.py: 100%

12 statements  

« prev     ^ index     » next       coverage.py v7.1.0, created at 2023-02-04 02:28 +0100

1""" 

2@file 

3@brief Useful plots. 

4""" 

5from ..onnxrt import OnnxInference 

6 

7 

8def plot_onnx(onx, ax=None, dpi=300, temp_dot=None, temp_img=None, 

9 show=False): 

10 """ 

11 Plots an ONNX graph into a :epkg:`matplotlib` graph. 

12 

13 :param onx: ONNX object, @see cl OnnxInference 

14 :param ax: existing axes 

15 :param dpi: resolution 

16 :param temp_dot: temporary file, 

17 if None, a file is created and removed 

18 :param temp_img: temporary image, 

19 if None, a file is created and removed 

20 :param show: calls `plt.show()` 

21 :return: axes 

22 """ 

23 # delayed import, because too long 

24 from pyquickhelper.helpgen.graphviz_helper import plot_graphviz 

25 import matplotlib.pyplot as plt 

26 

27 if ax is None: 

28 ax = plt.gca() # pragma: no cover 

29 elif isinstance(ax, str) and ax == 'new': 

30 _, ax = plt.subplots(1, 1) # pragma: no cover 

31 if not isinstance(onx, OnnxInference): 

32 oinf = OnnxInference(onx, skip_run=True) 

33 else: 

34 oinf = onx # pragma: no cover 

35 dot = oinf.to_dot() 

36 plot_graphviz(dot, dpi=dpi, ax=ax, temp_dot=temp_dot, temp_img=temp_img) 

37 if show: 

38 plt.show() # pragma: no cover 

39 return ax