Coverage for onnxcustom/plotting/plotting_onnx.py: 92%

1# flake8: noqa: F401 

2""" 

3@file 

4@brief Shortcuts to plotting. 

5""" 

6from mlprodict.plotting.plotting_onnx import plot_onnx 

7 

8 

9def plot_onnxs(*onx, ax=None, dpi=300, temp_dot=None, temp_img=None, 

10 show=False, title=None): 

11 """ 

12 Plots one or several ONNX graph into a :epkg:`matplotlib` graph. 

13 

14 :param onx: ONNX objects 

15 :param ax: existing axes 

16 :param dpi: resolution 

17 :param temp_dot: temporary file, 

18 if None, a file is created and removed 

19 :param temp_img: temporary image, 

20 if None, a file is created and removed 

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

22 :param title: graph title 

23 :return: axes 

24 """ 

25 if len(onx) == 1: 

26 if ax is None: 

27 import matplotlib.pyplot as plt # pylint: disable=C0415 

28 ax = plt.gca() 

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

30 import matplotlib.pyplot as plt # pylint: disable=C0415 

31 _, ax = plt.subplots(1, 1) 

32 ax = plot_onnx(onx[0], ax=ax, dpi=dpi, temp_dot=temp_dot, 

33 temp_img=temp_img) 

34 if title is not None: 

35 ax.set_title(title) 

36 return ax 

37 elif len(onx) > 1 and isinstance(ax, str) and ax == 'new': 

38 ax = None 

39 

40 if len(onx) == 0: 

41 raise ValueError( 

42 "Empty list of graph to plot.") 

43 

44 if ax is None: 

45 import matplotlib.pyplot as plt # pylint: disable=C0415 

46 fig, ax = plt.subplots(1, len(onx)) 

47 else: 

48 fig = None 

49 if ax.shape[0] != len(onx): 

50 raise ValueError( 

51 f"ax must be an array of shape ({len(onx)}, ).") 

52 for i, ox in enumerate(onx): 

53 plot_onnx(ox, ax=ax[i], dpi=dpi, temp_dot=temp_dot, 

54 temp_img=temp_img) 

55 if title is None or isinstance(title, str): 

56 continue 

57 if i < len(title): 

58 ax[i].set_title(title[i]) 

59 if len(onx) > 1 and isinstance(title, str): 

60 if fig is None: 

61 raise ValueError( # pragma: no cover 

62 "Main title cannot be set if fig is undefined (title=%r, " 

63 "len(onx)=%d)" % (title, len(onx))) 

64 fig.suptitle(title) 

65 elif len(onx) == 1: 

66 if isinstance(title, list): 

67 title = title[0] 

68 ax.set_title(title) 

69 return ax