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

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# max 10k x 10k RGBA 

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

21 

22 

23class Error(gws.Error): 

24 pass 

25 

26 

27class FormatConfig(gws.Config): 

28 """Image format configuration.""" 

29 

30 mimeTypes: list[gws.MimeType] 

31 """Mime types for this format.""" 

32 options: Optional[dict] 

33 """Image options.""" 

34 

35 

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

37 """Creates a monochrome image object. 

38 

39 Args: 

40 size: `(width, height)` 

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

42 

43 Returns: 

44 An image object. 

45 """ 

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

47 return _new(img) 

48 

49 

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

51 """Creates an image object from bytes. 

52 

53 Args: 

54 r: Bytes encoding an image. 

55 

56 Returns: 

57 An image object. 

58 """ 

59 with io.BytesIO(r) as fp: 

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

61 

62 

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

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

65 

66 Args: 

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

68 mode: PIL image mode. 

69 size: `(width, height)` 

70 

71 Returns: 

72 An image object. 

73 """ 

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

75 

76 

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

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

79 

80 Args: 

81 path: Path to an existing image. 

82 

83 Returns: 

84 An image object. 

85 """ 

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

87 return from_bytes(fp.read()) 

88 

89 

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

91 

92 

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

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

95 

96 Args: 

97 url: URL encoding an image. 

98 

99 Returns: 

100 An image object. 

101 """ 

102 m = re.match(_DATA_URL_RE, url) 

103 if not m: 

104 raise Error(f'invalid data url') 

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

106 return from_bytes(r) 

107 

108 

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

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

111 

112 Args: 

113 arr: Numpy array encoding an image. 

114 Returns: 

115 An image object. 

116 """ 

117 img = PIL.Image.fromarray(arr) 

118 return _new(img) 

119 

120 

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

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

123 

124 Args: 

125 xmlstr: XML String of the image. 

126 

127 size: `(width, height)` 

128 

129 mime: Mime type. 

130 

131 Returns: 

132 An image object. 

133 """ 

134 # @TODO rasterize svg 

135 raise NotImplemented 

136 

137 

138def qr_code( 

139 data: str, 

140 level='M', 

141 scale=4, 

142 border=True, 

143 color='black', 

144 background='white', 

145) -> 'Image': 

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

147 

148 Args: 

149 data: Data to encode. 

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

151 scale: Box size in pixels. 

152 border: Include a quiet zone of 4 boxes. 

153 color: Foreground color. 

154 background: Background color. 

155 

156 References: 

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

158 

159 """ 

160 

161 ec_map = { 

162 'L': qrcode.constants.ERROR_CORRECT_L, 

163 'M': qrcode.constants.ERROR_CORRECT_M, 

164 'Q': qrcode.constants.ERROR_CORRECT_Q, 

165 'H': qrcode.constants.ERROR_CORRECT_H, 

166 } 

167 

168 qr = qrcode.main.QRCode( 

169 version=None, 

170 error_correction=ec_map[level], 

171 box_size=scale, 

172 border=4 if border else 0, 

173 ) 

174 

175 qr.add_data(data) 

176 qr.make(fit=True) 

177 

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

179 return _new(img) 

180 

181 

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

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

184 

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

186 

187 

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

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

190 

191 Args: 

192 size: Font size. 

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

194 """ 

195 

196 if font: 

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

198 return PIL.ImageFont.load_default() 

199 

200 

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

202 try: 

203 img.load() 

204 except Exception as exc: 

205 raise Error from exc 

206 return Image(img) 

207 

208 

209class Image(gws.Image): 

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

211 

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

213 self.img: PIL.Image.Image = img 

214 

215 def mode(self): 

216 return self.img.mode 

217 

218 def size(self): 

219 return self.img.size 

220 

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

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

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

224 return self 

225 

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

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

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

229 return self 

230 

231 def crop(self, box): 

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

233 return self 

234 

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

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

237 return self 

238 

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

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

241 

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

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

244 

245 if opacity < 1: 

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

247 oth.putalpha(alpha) 

248 

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

250 return self 

251 

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

253 with io.BytesIO() as fp: 

254 self._save(fp, mime, options) 

255 return fp.getvalue() 

256 

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

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

259 return b.decode('ascii') 

260 

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

262 mime = mime or gws.lib.mime.PNG 

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

264 

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

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

267 self._save(fp, mime, options) 

268 return path 

269 

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

271 fmt = _mime_to_format(mime) 

272 opts = dict(options or {}) 

273 img = self.img 

274 

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

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

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

278 img.alpha_composite(self.img) 

279 img = img.convert('RGB') 

280 

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

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

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

284 

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

286 

287 def to_array(self): 

288 return np.array(self.img) 

289 

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

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

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

293 font = PIL.ImageFont.load_default() 

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

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

296 return self 

297 

298 def add_box(self, color=None): 

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

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

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

302 x, y = self.img.size 

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

304 return self 

305 

306 def compare_to(self, other): 

307 error = 0 

308 x, y = self.size() 

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

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

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

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

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

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

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

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

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

318 

319 

320_MIME_TO_FORMAT = { 

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

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

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

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

325} 

326 

327 

328def _mime_to_format(mime): 

329 if not mime: 

330 return 'PNG' 

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

332 if m in _MIME_TO_FORMAT: 

333 return _MIME_TO_FORMAT[m] 

334 m = m.split('/') 

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

336 return m[1].upper() 

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

338 

339 

340def _int_size(size: gws.Size): 

341 w, h = size 

342 return int(w), int(h) 

343 

344 

345_PIXELS = {} 

346_ERROR_COLOR = '#ffa1b4' 

347 

348 

349def empty_pixel(mime: str = None): 

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

351 

352 

353def error_pixel(mime: str = None): 

354 return pixel(mime, _ERROR_COLOR) 

355 

356 

357def pixel(mime, color): 

358 fmt = _mime_to_format(mime) 

359 key = fmt, str(color) 

360 

361 if key not in _PIXELS: 

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

363 with io.BytesIO() as fp: 

364 img.save(fp, fmt) 

365 _PIXELS[key] = fp.getvalue() 

366 

367 return _PIXELS[key]