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

1"""Wrapper for PIL objects""" 

2 

3import base64 

4import io 

5import re 

6from typing import Optional, cast 

7 

8import PIL.Image 

9import PIL.ImageDraw 

10import PIL.ImageFont 

11import numpy as np 

12import qrcode.main 

13import qrcode.constants 

14 

15import gws 

16import gws.lib.mime 

17 

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

19# up to ~4 GB RGBA images 

20MAX_PIXELS = 1_000_000_000 

21PIL.Image.MAX_IMAGE_PIXELS = MAX_PIXELS // 2 

22 

23 

24class Error(gws.Error): 

25 pass 

26 

27 

28class FormatConfig(gws.Config): 

29 """Image format configuration.""" 

30 

31 name: str = '' 

32 """Name of the format.""" 

33 mimeTypes: list[gws.MimeType] 

34 """Mime types for this format.""" 

35 options: Optional[dict] 

36 """Image options.""" 

37 

38 

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

40 """Creates a monochrome image object. 

41 

42 Args: 

43 size: `(width, height)` 

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

45 

46 Returns: 

47 An image object. 

48 """ 

49 w, h = _int_size(size) 

50 if w * h > MAX_PIXELS: 

51 raise Error(f'image too large: {w}x{h}') 

52 img = PIL.Image.new('RGBA', (w, h), color or (0, 0, 0, 0)) 

53 return _new(img) 

54 

55 

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

57 """Creates an image object from bytes. 

58 

59 Args: 

60 r: Bytes encoding an image. 

61 

62 Returns: 

63 An image object. 

64 """ 

65 with io.BytesIO(r) as fp: 

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

67 

68 

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

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

71 

72 Args: 

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

74 mode: PIL image mode. 

75 size: `(width, height)` 

76 

77 Returns: 

78 An image object. 

79 """ 

80 

81 w, h = _int_size(size) 

82 if w * h > MAX_PIXELS: 

83 raise Error(f'image too large: {w}x{h}') 

84 return _new(PIL.Image.frombytes(mode, (w, h), r)) 

85 

86 

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

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

89 

90 Args: 

91 path: Path to an existing image. 

92 

93 Returns: 

94 An image object. 

95 """ 

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

97 return from_bytes(fp.read()) 

98 

99 

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

101 

102 

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

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

105 

106 Args: 

107 url: URL encoding an image. 

108 

109 Returns: 

110 An image object. 

111 """ 

112 m = re.match(_DATA_URL_RE, url) 

113 if not m: 

114 raise Error(f'invalid data url') 

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

116 return from_bytes(r) 

117 

118 

119def from_array(arr: np.ndarray, mode: str = None) -> 'Image': 

120 """Creates an image object from a numpy array. 

121 

122 Args: 

123 arr: Numpy array encoding an image. 

124 Returns: 

125 An image object. 

126 """ 

127 img = PIL.Image.fromarray(arr) 

128 return _new(img) 

129 

130 

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

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

133 

134 Args: 

135 xmlstr: XML String of the image. 

136 

137 size: `(width, height)` 

138 

139 mime: Mime type. 

140 

141 Returns: 

142 An image object. 

143 """ 

144 # @TODO rasterize svg 

145 raise NotImplementedError 

146 

147 

148def thumbnail(r: bytes, size: gws.Size, max_pixels=0, mime=None, options=None) -> bytes: 

149 """Creates a thumbnail from image bytes. 

150 

151 The image is scaled to fit into ``size``, small images are not scaled up. 

152 

153 Args: 

154 r: Bytes encoding an image. 

155 size: Maximum thumbnail size `(width, height)` 

156 max_pixels: Maximum number of source pixels. 

157 mime: Mime type of the thumbnail. 

158 options: Image options. 

159 

160 Returns: 

161 Bytes encoding the thumbnail. 

162 """ 

163 

164 sz = _int_size(size) 

165 

166 try: 

167 with io.BytesIO(r) as fp: 

168 img = PIL.Image.open(fp) 

169 w, h = img.size 

170 if max_pixels and w * h > max_pixels: 

171 raise Error(f'image too big: {w}x{h}') 

172 img.draft(img.mode, sz) 

173 img.thumbnail(sz, resample=PIL.Image.Resampling.BICUBIC) 

174 if img.mode not in {'1', 'L', 'LA', 'P', 'RGB', 'RGBA'}: 

175 img = img.convert('RGB') 

176 return Image(img).to_bytes(mime, options) 

177 except Error: 

178 raise 

179 except Exception as exc: 

180 raise Error from exc 

181 

182 

183def qr_code( 

184 data: str, 

185 level='M', 

186 scale=4, 

187 border=True, 

188 color='black', 

189 background='white', 

190) -> 'Image': 

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

192 

193 Args: 

194 data: Data to encode. 

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

196 scale: Box size in pixels. 

197 border: Include a quiet zone of 4 boxes. 

198 color: Foreground color. 

199 background: Background color. 

200 

201 References: 

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

203 

204 """ 

205 

206 ec_map = { 

207 'L': qrcode.constants.ERROR_CORRECT_L, 

208 'M': qrcode.constants.ERROR_CORRECT_M, 

209 'Q': qrcode.constants.ERROR_CORRECT_Q, 

210 'H': qrcode.constants.ERROR_CORRECT_H, 

211 } 

212 

213 qr = qrcode.main.QRCode( 

214 version=None, 

215 error_correction=ec_map[level], 

216 box_size=scale, 

217 border=4 if border else 0, 

218 ) 

219 

220 qr.add_data(data) 

221 qr.make(fit=True) 

222 

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

224 return _new(img) 

225 

226 

227def get_draw(img: 'Image') -> PIL.ImageDraw.ImageDraw: 

228 """Returns a PIL ImageDraw object for the given image.""" 

229 

230 return PIL.ImageDraw.Draw(img.img) 

231 

232 

233def get_font(size: int = 12, font: Optional[str] = None) -> PIL.ImageFont.ImageFont | PIL.ImageFont.FreeTypeFont: 

234 """Returns a PIL ImageFont object for the given size and font. 

235 

236 Args: 

237 size: Font size. 

238 font: Path to a TTF font file or None for default font. 

239 """ 

240 

241 if font: 

242 return PIL.ImageFont.truetype(font, size) 

243 return PIL.ImageFont.load_default() 

244 

245 

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

247 try: 

248 img.load() 

249 except Exception as exc: 

250 raise Error from exc 

251 return Image(img) 

252 

253 

254class Image(gws.Image): 

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

256 

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

258 self.img: PIL.Image.Image = img 

259 

260 def mode(self): 

261 return self.img.mode 

262 

263 def size(self): 

264 return self.img.size 

265 

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

267 kwargs.setdefault('resample', PIL.Image.Resampling.BICUBIC) 

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

269 return self 

270 

271 def resize_to(self, width=0, height=0, **kwargs): 

272 w, h = self.img.size 

273 if width and height: 

274 sz = (width, height) 

275 elif width: 

276 sz = (width, int(h * width / w)) 

277 elif height: 

278 sz = (int(w * height / h), height) 

279 else: 

280 return self 

281 return self.resize(sz, **kwargs) 

282 

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

284 kwargs.setdefault('resample', PIL.Image.Resampling.BICUBIC) 

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

286 return self 

287 

288 def crop(self, box): 

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

290 return self 

291 

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

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

294 return self 

295 

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

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

298 

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

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

301 

302 if opacity < 1: 

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

304 oth.putalpha(alpha) 

305 

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

307 return self 

308 

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

310 with io.BytesIO() as fp: 

311 self._save(fp, mime, options) 

312 return fp.getvalue() 

313 

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

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

316 return b.decode('ascii') 

317 

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

319 mime = mime or gws.lib.mime.PNG 

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

321 

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

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

324 self._save(fp, mime, options) 

325 return path 

326 

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

328 fmt = _mime_to_format(mime) 

329 opts = dict(options or {}) 

330 img = self.img 

331 

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

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

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

335 img.alpha_composite(self.img) 

336 img = img.convert('RGB') 

337 

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

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

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

341 

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

343 

344 def to_array(self): 

345 return np.array(self.img) 

346 

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

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

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

350 font = PIL.ImageFont.load_default() 

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

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

353 return self 

354 

355 def add_box(self, color=None): 

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

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

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

359 x, y = self.img.size 

360 draw.rectangle((0, 0) + (x - 1, y - 1), outline=color) 

361 return self 

362 

363 def compare_to(self, other): 

364 error = 0 

365 x, y = self.size() 

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

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

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

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

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

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

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

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

374 return error / (4 * x * y * 255 * 255) 

375 

376 

377_MIME_TO_FORMAT = { 

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

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

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

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

382} 

383 

384 

385def _mime_to_format(mime): 

386 if not mime: 

387 return 'PNG' 

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

389 if m in _MIME_TO_FORMAT: 

390 return _MIME_TO_FORMAT[m] 

391 m = m.split('/') 

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

393 return m[1].upper() 

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

395 

396 

397def _int_size(size: gws.Size): 

398 w, h = size 

399 return int(w), int(h) 

400 

401 

402_PIXELS = {} 

403_ERROR_COLOR = '#ffa1b4' 

404 

405 

406def empty_pixel(mime: str = None): 

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

408 

409 

410def error_pixel(mime: str = None): 

411 return pixel(mime, _ERROR_COLOR) 

412 

413 

414def pixel(mime, color): 

415 fmt = _mime_to_format(mime) 

416 key = fmt, str(color) 

417 

418 if key not in _PIXELS: 

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

420 with io.BytesIO() as fp: 

421 img.save(fp, fmt) 

422 _PIXELS[key] = fp.getvalue() 

423 

424 return _PIXELS[key]