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

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_svg(xmlstr: str, size: gws.Size, mime=None) -> 'Image': 

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

111 

112 Args: 

113 xmlstr: XML String of the image. 

114 

115 size: `(width, height)` 

116 

117 mime: Mime type. 

118 

119 Returns: 

120 An image object. 

121 """ 

122 # @TODO rasterize svg 

123 raise NotImplemented 

124 

125 

126def qr_code( 

127 data: str, 

128 level='M', 

129 scale=4, 

130 border=True, 

131 color='black', 

132 background='white', 

133) -> 'Image': 

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

135 

136 Args: 

137 data: Data to encode. 

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

139 scale: Box size in pixels. 

140 border: Include a quiet zone of 4 boxes. 

141 color: Foreground color. 

142 background: Background color. 

143 

144 References: 

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

146 

147 """ 

148 

149 ec_map = { 

150 'L': qrcode.constants.ERROR_CORRECT_L, 

151 'M': qrcode.constants.ERROR_CORRECT_M, 

152 'Q': qrcode.constants.ERROR_CORRECT_Q, 

153 'H': qrcode.constants.ERROR_CORRECT_H, 

154 } 

155 

156 qr = qrcode.main.QRCode( 

157 version=None, 

158 error_correction=ec_map[level], 

159 box_size=scale, 

160 border=4 if border else 0, 

161 ) 

162 

163 qr.add_data(data) 

164 qr.make(fit=True) 

165 

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

167 return _new(img) 

168 

169 

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

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

172 

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

174 

175 

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

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

178 

179 Args: 

180 size: Font size. 

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

182 """ 

183 

184 if font: 

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

186 return PIL.ImageFont.load_default() 

187 

188 

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

190 try: 

191 img.load() 

192 except Exception as exc: 

193 raise Error from exc 

194 return Image(img) 

195 

196 

197class Image(gws.Image): 

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

199 

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

201 self.img: PIL.Image.Image = img 

202 

203 def mode(self): 

204 return self.img.mode 

205 

206 def size(self): 

207 return self.img.size 

208 

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

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

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

212 return self 

213 

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

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

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

217 return self 

218 

219 def crop(self, box): 

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

221 return self 

222 

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

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

225 return self 

226 

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

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

229 

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

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

232 

233 if opacity < 1: 

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

235 oth.putalpha(alpha) 

236 

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

238 return self 

239 

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

241 with io.BytesIO() as fp: 

242 self._save(fp, mime, options) 

243 return fp.getvalue() 

244 

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

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

247 return b.decode('ascii') 

248 

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

250 mime = mime or gws.lib.mime.PNG 

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

252 

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

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

255 self._save(fp, mime, options) 

256 return path 

257 

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

259 fmt = _mime_to_format(mime) 

260 opts = dict(options or {}) 

261 img = self.img 

262 

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

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

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

266 img.alpha_composite(self.img) 

267 img = img.convert('RGB') 

268 

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

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

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

272 

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

274 

275 def to_array(self): 

276 return np.array(self.img) 

277 

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

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

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

281 font = PIL.ImageFont.load_default() 

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

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

284 return self 

285 

286 def add_box(self, color=None): 

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

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

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

290 x, y = self.img.size 

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

292 return self 

293 

294 def compare_to(self, other): 

295 error = 0 

296 x, y = self.size() 

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

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

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

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

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

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

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

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

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

306 

307 

308_MIME_TO_FORMAT = { 

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

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

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

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

313} 

314 

315 

316def _mime_to_format(mime): 

317 if not mime: 

318 return 'PNG' 

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

320 if m in _MIME_TO_FORMAT: 

321 return _MIME_TO_FORMAT[m] 

322 m = m.split('/') 

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

324 return m[1].upper() 

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

326 

327 

328def _int_size(size: gws.Size): 

329 w, h = size 

330 return int(w), int(h) 

331 

332 

333_PIXELS = {} 

334_ERROR_COLOR = '#ffa1b4' 

335 

336 

337def empty_pixel(mime: str = None): 

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

339 

340 

341def error_pixel(mime: str = None): 

342 return pixel(mime, _ERROR_COLOR) 

343 

344 

345def pixel(mime, color): 

346 fmt = _mime_to_format(mime) 

347 key = fmt, str(color) 

348 

349 if key not in _PIXELS: 

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

351 with io.BytesIO() as fp: 

352 img.save(fp, fmt) 

353 _PIXELS[key] = fp.getvalue() 

354 

355 return _PIXELS[key]