Coverage for thumbhash / encode.py: 98%

1import math (empty)

2from base64 import b64encode (empty)

3from collections.abc import Sequence (empty)

4from itertools import chain (empty)

5from pathlib import Path (empty)

6from typing import BinaryIO (empty)

7 

8from PIL import Image (empty)

9from PIL.ImageOps import exif_transpose (empty)

10 

11from thumbhash.hash import Hash (empty)

12 

13 

14def image_to_thumbhash( (empty)

15 image: str | bytes | Path | BinaryIO | Image.Image, 

16) -> str: 

17 if not isinstance(image, Image.Image): (empty)

18 image = Image.open(image) (empty)

19 

20 m_image = exif_transpose(image).convert("RGBA") (empty)

21 

22 m_image.thumbnail((100, 100)) (empty)

23 

24 red_band = m_image.get_flattened_data(band=0) (empty)

25 green_band = m_image.get_flattened_data(band=1) (empty)

26 blue_band = m_image.get_flattened_data(band=2) (empty)

27 alpha_band = m_image.get_flattened_data(band=3) (empty)

28 rgb_data = list( (empty)

29 chain.from_iterable( 

30 zip(red_band, green_band, blue_band, alpha_band, strict=True) 

31 ) 

32 ) 

33 width, height = m_image.size (empty)

34 m_image.close() (empty)

35 

36 hash = rgba_to_thumbhash(width, height, rgb_data) # type: ignore (empty)

37 

38 return b64encode(hash).decode("utf-8") (empty)

39 

40 

41def rgba_to_thumbhash(w: int, h: int, rgba: Sequence[int]) -> bytes: (empty)

42 if w > 100 or h > 100: (empty)

43 raise ValueError(f"{w}x{h} doesn't fit in 100x100") 

44 if len(rgba) != w * h * 4: (empty)

45 raise ValueError("Image data does not fit the given size") 

46 

47 nb_pixels = w * h (empty)

48 is_landscape = w > h (empty)

49 avg_r = 0.0 (empty)

50 avg_g = 0.0 (empty)

51 avg_b = 0.0 (empty)

52 avg_a = 0.0 (empty)

53 

54 for i in range(nb_pixels): (empty)

55 alpha = float(rgba[i * 4 + 3]) / 255.0 (empty)

56 avg_r += alpha / 255.0 * float(rgba[i * 4]) (empty)

57 avg_g += alpha / 255.0 * float(rgba[i * 4 + 1]) (empty)

58 avg_b += alpha / 255.0 * float(rgba[i * 4 + 2]) (empty)

59 avg_a += alpha (empty)

60 

61 if avg_a > 0.0: (empty)

62 avg_r /= avg_a (empty)

63 avg_g /= avg_a (empty)

64 avg_b /= avg_a (empty)

65 

66 has_alpha = avg_a < float(nb_pixels) (empty)

67 l_limit = 5.0 if has_alpha else 7.0 (empty)

68 max_wh = max(w, h) (empty)

69 lx = max(1, round((l_limit * w) / max_wh)) (empty)

70 ly = max(1, round((l_limit * h) / max_wh)) (empty)

71 l_channel: list[float] = [0.0] * nb_pixels (empty)

72 p_channel: list[float] = [0.0] * nb_pixels (empty)

73 q_channel: list[float] = [0.0] * nb_pixels (empty)

74 a_channel: list[float] = [0.0] * nb_pixels (empty)

75 

76 for i in range(nb_pixels): (empty)

77 alpha = float(rgba[i * 4 + 3]) / 255.0 (empty)

78 

79 r = avg_r * (1.0 - alpha) + alpha / 255.0 * float(rgba[i * 4]) (empty)

80 g = avg_g * (1.0 - alpha) + alpha / 255.0 * float(rgba[i * 4 + 1]) (empty)

81 b = avg_b * (1.0 - alpha) + alpha / 255.0 * float(rgba[i * 4 + 2]) (empty)

82 

83 l_channel[i] = (r + g + b) / 3.0 (empty)

84 p_channel[i] = (r + g) / 2.0 - b (empty)

85 q_channel[i] = r - g (empty)

86 a_channel[i] = alpha (empty)

87 

88 def encode_channel( (empty)

89 channel: list[float], nx: int, ny: int 

90 ) -> tuple[float, list[float], float]: 

91 dc = 0.0 (empty)

92 ac: list[float] = [] (empty)

93 scale = 0.0 (empty)

94 fx: list[float] = [0.0] * w (empty)

95 

96 for cy in range(ny): (empty)

97 cyf = float(cy) (empty)

98 cx = 0 (empty)

99 

100 while cx * ny < nx * (ny - cy): (empty)

101 cxf = float(cx) (empty)

102 f = 0.0 (empty)

103 

104 for x in range(w): (empty)

105 fx[x] = math.cos(math.pi / w * cxf * (float(x) + 0.5)) (empty)

106 

107 for y in range(h): (empty)

108 fy = math.cos(math.pi / h * cyf * (float(y) + 0.5)) (empty)

109 

110 for x in range(w): (empty)

111 f += channel[x + y * w] * fx[x] * fy (empty)

112 

113 f /= float(nb_pixels) (empty)

114 

115 if cx > 0 or cy > 0: (empty)

116 ac.append(f) (empty)

117 scale = max(scale, abs(f)) (empty)

118 else: 

119 dc = f (empty)

120 

121 cx += 1 (empty)

122 

123 if scale > 0.0: (empty)

124 for i in range(len(ac)): (empty)

125 ac[i] = 0.5 + 0.5 / scale * ac[i] (empty)

126 

127 return (dc, ac, scale) (empty)

128 

129 l_dc, l_ac, l_scale = encode_channel(l_channel, max(lx, 3), max(ly, 3)) (empty)

130 p_dc, p_ac, p_scale = encode_channel(p_channel, 3, 3) (empty)

131 q_dc, q_ac, q_scale = encode_channel(q_channel, 3, 3) (empty)

132 a_dc, a_ac, a_scale = ( (empty)

133 encode_channel(a_channel, 5, 5) if has_alpha else (1.0, [], 1.0) 

134 ) 

135 

136 hash = Hash( (empty)

137 l_dc=l_dc, 

138 p_dc=p_dc, 

139 q_dc=q_dc, 

140 l_scale=l_scale, 

141 has_alpha=has_alpha, 

142 lx=lx, 

143 ly=ly, 

144 p_scale=p_scale, 

145 q_scale=q_scale, 

146 is_landscape=is_landscape, 

147 a_dc=a_dc, 

148 a_scale=a_scale, 

149 l_ac=l_ac, 

150 p_ac=p_ac, 

151 q_ac=q_ac, 

152 a_ac=a_ac, 

153 l_count=0, 

154 ) 

155 

156 return hash.encode() (empty)