Coverage for supercat/viz.py: 100.00%

1from pathlib import Path 

2import numpy as np 

3import plotly.graph_objects as go 

4from PIL import Image 

5from plotly.subplots import make_subplots 

6import numpy as np 

7import typer 

8 

9app = typer.Typer() 

10 

11 

12def format_fig(fig): 

13 fig.update_layout( 

14 plot_bgcolor="white", 

15 title_font_color="black", 

16 font=dict( 

17 family="Linux Libertine Display O", 

18 size=18, 

19 color="black", 

20 ), 

21 ) 

22 

23 

24@app.command() 

25def comparison( 

26 hr: list[Path] = typer.Option( 

27 ..., help="Paths to the original high-resolution images" 

28 ), 

29 lr: list[Path] = typer.Option( 

30 ..., help="Paths to the downscaled low-resolution images" 

31 ), 

32 sr: list[Path] = typer.Option( 

33 ..., help="Paths to the upscaled super-resolution images" 

34 ), 

35 titles: list[str] = typer.Option( 

36 ..., help="Titles for each row, in the same order as the images" 

37 ), 

38 output: Path = typer.Option( 

39 None, help="Path to save the resulting figure (e.g., comparison.html)" 

40 ), 

41): 

42 assert ( 

43 len(hr) == len(lr) == len(sr) == len(titles) 

44 ), "The number of original, downscaled, upscaled images and titles must be the same" 

45 

46 fig = make_subplots( 

47 rows=len(hr), 

48 cols=4, 

49 subplot_titles=( 

50 "Original", 

51 "Downscaled", 

52 "Upscaled", 

53 "Difference", 

54 ), 

55 vertical_spacing=0.02, 

56 horizontal_spacing=0.02, 

57 ) 

58 

59 def read(x): 

60 if not isinstance(x, (Path, str)): 

61 return np.asarray(x) 

62 

63 x = str(x) 

64 if x.endswith(".mat"): 

65 from .data import read_mat 

66 

67 volume = read_mat(x) 

68 return np.asarray(volume[volume.shape[0] // 2]).astype(int) 

69 

70 return np.asarray(Image.open(x).convert("L")).astype(int) 

71 

72 for row, (original, downscaled, upscaled, title) in enumerate( 

73 zip(hr, lr, sr, titles) 

74 ): 

75 original_im = read(original) 

76 downscaled_im = read( 

77 downscaled 

78 ) # .resize( (original_im.size[0], original_im.size[1]), resample=PIL.Image.Resampling.NEAREST) 

79 upscaled = read(upscaled) 

80 

81 difference = upscaled - original_im 

82 # squared_error = np.power(difference.astype(float)/255, 2.0) 

83 

84 fig.add_trace( 

85 go.Heatmap( 

86 z=np.asarray(original_im).astype(int), 

87 colorscale="gray", 

88 showscale=False, 

89 zmin=0, 

90 zmax=255, 

91 ), 

92 row=row + 1, 

93 col=1, 

94 ) 

95 fig.add_trace( 

96 go.Heatmap( 

97 z=np.asarray(downscaled_im).astype(int), 

98 colorscale="gray", 

99 showscale=False, 

100 zmin=0, 

101 zmax=255, 

102 ), 

103 row=row + 1, 

104 col=2, 

105 ) 

106 fig.add_trace( 

107 go.Heatmap( 

108 z=np.asarray(upscaled).astype(int), 

109 colorscale="gray", 

110 showscale=False, 

111 zmin=0, 

112 zmax=255, 

113 ), 

114 row=row + 1, 

115 col=3, 

116 ) 

117 fig.add_trace( 

118 go.Heatmap(z=difference, coloraxis="coloraxis2"), row=row + 1, col=4 

119 ) 

120 

121 update_dict = { 

122 f"yaxis{1+row*4}_title": title, 

123 } 

124 fig.update_layout(**update_dict) 

125 fig.update_layout(plot_bgcolor="rgba(0,0,0,0)") 

126 fig.update_xaxes(showticklabels=False) 

127 fig.update_yaxes(showticklabels=False) 

128 fig.update_layout( 

129 height=150 + 240 * len(hr), 

130 width=1200, 

131 ) 

132 format_fig(fig) 

133 

134 # fig.update_layout(coloraxis1=dict(colorscale='gray'), showlegend=False) 

135 fig.update_layout(coloraxis2=dict(colorscale="Rainbow"), showlegend=False) 

136 fig.update_annotations(font_size=24) 

137 

138 if output is not None: 

139 print(f"Saving figure to {output}") 

140 if output.suffix.lower() in [".html", ".htm"]: 

141 fig.write_html(output) 

142 else: 

143 fig.write_image(output, scale=2) 

144 

145 return fig 

146 

147 

148if __name__ == "__main__": 

149 app()