Coverage for gws-app/gws/lib/image/__init__.py: 30%

1"""Wrapper for PIL objects""" 

2 

3from typing import Optional, Literal, cast 

4 

5import base64 

6import io 

7import re 

8 

9import PIL.Image 

10import PIL.ImageDraw 

11import PIL.ImageFont 

12import numpy as np 

13 

14import qrcode 

15 

16import gws 

17import gws.lib.mime 

18 

19# https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.open 

20# max 10k x 10k RGBA 

21PIL.Image.MAX_IMAGE_PIXELS = 10_000 * 10_000 * 4 

22 

23 

24class Error(gws.Error): 

25 pass 

26 

27 

28def from_size(size: gws.Size, color=None) -> 'Image': 

29 """Creates a monochrome image object. 

30 

31 Args: 

32 size: `(width, height)` 

33 color: `(red, green, blue, alpha)` 

34 

35 Returns: 

36 An image object. 

37 """ 

38 img = PIL.Image.new('RGBA', _int_size(size), color or (0, 0, 0, 0)) 

39 return _new(img) 

40 

41 

42def from_bytes(r: bytes) -> 'Image': 

43 """Creates an image object from bytes. 

44 

45 Args: 

46 r: Bytes encoding an image. 

47 

48 Returns: 

49 An image object. 

50 """ 

51 with io.BytesIO(r) as fp: 

52 return _new(PIL.Image.open(fp)) 

53 

54 

55def from_raw_data(r: bytes, mode: str, size: gws.Size) -> 'Image': 

56 """Creates an image object in a given mode from raw pixel data in arrays. 

57 

58 Args: 

59 r: Bytes encoding an image in arrays of pixels. 

60 mode: PIL image mode. 

61 size: `(width, height)` 

62 

63 Returns: 

64 An image object. 

65 """ 

66 return _new(PIL.Image.frombytes(mode, _int_size(size), r)) 

67 

68 

69def from_path(path: str) -> 'Image': 

70 """Creates an image object from a path. 

71 

72 Args: 

73 path: Path to an existing image. 

74 

75 Returns: 

76 An image object. 

77 """ 

78 with open(path, 'rb') as fp: 

79 return from_bytes(fp.read()) 

80 

81 

82_DATA_URL_RE = r'data:image/(png|gif|jpeg|jpg);base64,' 

83 

84 

85def from_data_url(url: str) -> Optional['Image']: 

86 """Creates an image object from a URL. 

87 

88 Args: 

89 url: URL encoding an image. 

90 

91 Returns: 

92 An image object. 

93 """ 

94 m = re.match(_DATA_URL_RE, url) 

95 if not m: 

96 raise Error(f'invalid data url') 

97 r = base64.standard_b64decode(url[m.end():]) 

98 return from_bytes(r) 

99 

100 

101def from_svg(xmlstr: str, size: gws.Size, mime=None) -> 'Image': 

102 """Not implemented yet. Should create an image object from a URL. 

103 

104 Args: 

105 xmlstr: XML String of the image. 

106 

107 size: `(width, height)` 

108 

109 mime: Mime type. 

110 

111 Returns: 

112 An image object. 

113 """ 

114 # @TODO rasterize svg 

115 raise NotImplemented 

116 

117 

118def qr_code( 

119 data: str, 

120 level='M', 

121 scale=4, 

122 border=True, 

123 color='black', 

124 background='white', 

125) -> 'Image': 

126 """Creates an Image with a QR code for the given data. 

127 

128 Args: 

129 data: Data to encode. 

130 level: Error correction level, one of L M Q H. 

131 scale: Box size in pixels. 

132 border: Include a quiet zone of 4 boxes. 

133 color: Foreground color. 

134 background: Background color. 

135 

136 References: 

137 - https://github.com/lincolnloop/python-qrcode/blob/main/README.rst#advanced-usage 

138 

139 """ 

140 

141 ec_map = { 

142 'L': qrcode.constants.ERROR_CORRECT_L, 

143 'M': qrcode.constants.ERROR_CORRECT_M, 

144 'Q': qrcode.constants.ERROR_CORRECT_Q, 

145 'H': qrcode.constants.ERROR_CORRECT_H, 

146 } 

147 

148 qr = qrcode.main.QRCode( 

149 version=None, 

150 error_correction=ec_map[level], 

151 box_size=scale, 

152 border=4 if border else 0, 

153 ) 

154 

155 qr.add_data(data) 

156 qr.make(fit=True) 

157 

158 img = qr.make_image(fill_color=color, back_color=background) 

159 return _new(img) 

160 

161 

162def _new(img: PIL.Image.Image): 

163 try: 

164 img.load() 

165 except Exception as exc: 

166 raise Error from exc 

167 return Image(img) 

168 

169 

170class Image(gws.Image): 

171 """Class to convert, save and do basic manipulations on images.""" 

172 

173 def __init__(self, img: PIL.Image.Image): 

174 self.img: PIL.Image.Image = img 

175 

176 def mode(self): 

177 return self.img.mode 

178 

179 def size(self): 

180 return self.img.size 

181 

182 def resize(self, size, **kwargs): 

183 kwargs.setdefault('resample', PIL.Image.BICUBIC) 

184 self.img = self.img.resize(_int_size(size), **kwargs) 

185 return self 

186 

187 def rotate(self, angle, **kwargs): 

188 kwargs.setdefault('resample', PIL.Image.BICUBIC) 

189 self.img = self.img.rotate(angle, **kwargs) 

190 return self 

191 

192 def crop(self, box): 

193 self.img = self.img.crop(box) 

194 return self 

195 

196 def paste(self, other, where=None): 

197 self.img.paste(cast('Image', other).img, where) 

198 return self 

199 

200 def compose(self, other, opacity=1): 

201 oth = cast('Image', other).img.convert('RGBA') 

202 

203 if oth.size != self.img.size: 

204 oth = oth.resize(size=self.img.size, resample=PIL.Image.BICUBIC) 

205 

206 if opacity < 1: 

207 alpha = oth.getchannel('A').point(lambda x: int(x * opacity)) 

208 oth.putalpha(alpha) 

209 

210 self.img = PIL.Image.alpha_composite(self.img, oth) 

211 return self 

212 

213 def to_bytes(self, mime=None, options=None): 

214 with io.BytesIO() as fp: 

215 self._save(fp, mime, options) 

216 return fp.getvalue() 

217 

218 def to_base64(self, mime=None, options=None): 

219 b = base64.standard_b64encode(self.to_bytes(mime, options)) 

220 return b.decode('ascii') 

221 

222 def to_data_url(self, mime=None, options=None): 

223 mime = mime or gws.lib.mime.PNG 

224 return f'data:{mime};base64,' + self.to_base64(mime, options) 

225 

226 def to_path(self, path, mime=None, options=None): 

227 with open(path, 'wb') as fp: 

228 self._save(fp, mime, options) 

229 return path 

230 

231 def _save(self, fp, mime: str, options: dict): 

232 fmt = _mime_to_format(mime) 

233 opts = dict(options or {}) 

234 img = self.img 

235 

236 if self.img.mode == 'RGBA' and fmt == 'JPEG': 

237 background = opts.pop('background', '#FFFFFF') 

238 img = PIL.Image.new('RGBA', self.img.size, background) 

239 img.alpha_composite(self.img) 

240 img = img.convert('RGB') 

241 

242 mode = opts.pop('mode', '') 

243 if mode and self.img.mode != mode: 

244 img = img.convert(mode, palette=PIL.Image.ADAPTIVE) 

245 

246 img.save(fp, fmt, **opts) 

247 

248 def to_array(self): 

249 return np.array(self.img) 

250 

251 def add_text(self, text, x=0, y=0, color=None): 

252 self.img = self.img.convert('RGBA') 

253 draw = PIL.ImageDraw.Draw(self.img) 

254 font = PIL.ImageFont.load_default() 

255 color = color or (0, 0, 0, 255) 

256 draw.multiline_text((x, y), text, font=font, fill=color) 

257 return self 

258 

259 def add_box(self, color=None): 

260 self.img = self.img.convert('RGBA') 

261 draw = PIL.ImageDraw.Draw(self.img) 

262 color = color or (0, 0, 0, 255) 

263 x, y = self.img.size 

264 draw.rectangle((0, 0) + (x - 1, y - 1), outline=color) # box goes around all edges 

265 return self 

266 

267 def compare_to(self, other): 

268 error = 0 

269 x, y = self.size() 

270 for i in range(int(x)): 

271 for j in range(int(y)): 

272 a_r, a_g, a_b, a_a = self.img.getpixel((i, j)) 

273 b_r, b_g, b_b, b_a = cast(Image, other).img.getpixel((i, j)) 

274 error += (a_r - b_r) ** 2 

275 error += (a_g - b_g) ** 2 

276 error += (a_b - b_b) ** 2 

277 error += (a_a - b_a) ** 2 

278 return error / (3 * x * y) 

279 

280 

281_MIME_TO_FORMAT = { 

282 gws.lib.mime.PNG: 'PNG', 

283 gws.lib.mime.JPEG: 'JPEG', 

284 gws.lib.mime.GIF: 'GIF', 

285 gws.lib.mime.WEBP: 'WEBP', 

286} 

287 

288 

289def _mime_to_format(mime): 

290 if not mime: 

291 return 'PNG' 

292 m = mime.split(';')[0].strip() 

293 if m in _MIME_TO_FORMAT: 

294 return _MIME_TO_FORMAT[m] 

295 m = m.split('/') 

296 if len(m) == 2 and m[0] == 'image': 

297 return m[1].upper() 

298 raise Error(f'unknown mime type {mime!r}') 

299 

300 

301def _int_size(size: gws.Size): 

302 w, h = size 

303 return int(w), int(h) 

304 

305 

306_PIXELS = {} 

307_ERROR_COLOR = '#ffa1b4' 

308 

309 

310def empty_pixel(mime: str = None): 

311 return pixel(mime, '#ffffff' if mime == gws.lib.mime.JPEG else None) 

312 

313 

314def error_pixel(mime: str = None): 

315 return pixel(mime, _ERROR_COLOR) 

316 

317 

318def pixel(mime, color): 

319 fmt = _mime_to_format(mime) 

320 key = fmt, str(color) 

321 

322 if key not in _PIXELS: 

323 img = PIL.Image.new('RGBA' if color is None else 'RGB', (1, 1), color) 

324 with io.BytesIO() as fp: 

325 img.save(fp, fmt) 

326 _PIXELS[key] = fp.getvalue() 

327 

328 return _PIXELS[key]