Coverage for gws-app/gws/lib/image/core.py: 89%

1"""Image object and factories. 

2 

3Provides ``Image``, the implementation of ``gws.Image`` on top of a PIL image, and the ``from_*`` 

4functions that create images from a size, encoded bytes, raw pixel data, a file, a data URL 

5or a numpy array. 

6 

7Image sizes are limited to ``MAX_PIXELS``. 

8""" 

9 

10import io 

11from typing import cast 

12 

13import PIL.Image 

14import PIL.ImageDraw 

15import PIL.ImageFont 

16import numpy as np 

17 

18import gws 

19import gws.lib.mime 

20import gws.lib.text 

21 

22__all__ = [ 

23 'MAX_PIXELS', 

24 'Error', 

25 'Image', 

26 'from_size', 

27 'from_bytes', 

28 'from_raw_data', 

29 'from_path', 

30 'from_data_url', 

31 'from_array', 

32 'from_svg', 

33] 

34 

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

36# up to ~4 GB RGBA images 

37MAX_PIXELS = 1_000_000_000 

38"""Maximum number of pixels of an image created from a size or raw data.""" 

39PIL.Image.MAX_IMAGE_PIXELS = MAX_PIXELS // 2 

40 

41 

42class Error(gws.Error): 

43 """Image error.""" 

44 

45 pass 

46 

47 

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

49 """Create an RGBA image filled with one color. 

50 

51 Args: 

52 size: Image size ``(width, height)``. 

53 color: Fill color ``(red, green, blue, alpha)``, transparent by default. 

54 

55 Returns: 

56 An image object. 

57 

58 Raises: 

59 ``Error``: If the image has more than ``MAX_PIXELS`` pixels. 

60 """ 

61 w, h = _int_size(size) 

62 if w * h > MAX_PIXELS: 

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

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

65 return _new(img) 

66 

67 

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

69 """Create an image from encoded bytes. 

70 

71 Args: 

72 r: Encoded image, in any format PIL can read. 

73 

74 Returns: 

75 An image object. 

76 

77 Raises: 

78 ``Error``: If the image data cannot be loaded. 

79 """ 

80 with io.BytesIO(r) as fp: 

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

82 

83 

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

85 """Create an image from raw pixel data. 

86 

87 Args: 

88 r: Raw pixel data. 

89 mode: PIL image mode of the data. 

90 size: Image size ``(width, height)``. 

91 

92 Returns: 

93 An image object. 

94 

95 Raises: 

96 ``Error``: If the image has more than ``MAX_PIXELS`` pixels. 

97 """ 

98 

99 w, h = _int_size(size) 

100 if w * h > MAX_PIXELS: 

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

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

103 

104 

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

106 """Create an image from a file. 

107 

108 Args: 

109 path: Path to an image file. 

110 

111 Returns: 

112 An image object. 

113 

114 Raises: 

115 ``Error``: If the image data cannot be loaded. 

116 """ 

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

118 return from_bytes(fp.read()) 

119 

120 

121_DATA_URL_MIME_TYPES = {gws.lib.mime.PNG, gws.lib.mime.GIF, gws.lib.mime.JPEG} 

122 

123 

124def from_data_url(url: str) -> 'Image': 

125 """Create an image from a base64 data URL. 

126 

127 Only PNG, GIF and JPEG data URLs are accepted. 

128 

129 Args: 

130 url: A ``data:image/...;base64,`` URL. 

131 

132 Returns: 

133 An image object. 

134 

135 Raises: 

136 ``Error``: If the URL is not a supported data URL or the image cannot be loaded. 

137 """ 

138 try: 

139 mime, content = gws.lib.text.parse_data_url(url) 

140 except gws.lib.text.Error as exc: 

141 raise Error('invalid data url') from exc 

142 if gws.lib.mime.get_base(mime) not in _DATA_URL_MIME_TYPES: 

143 raise Error('invalid data url') 

144 return from_bytes(content) 

145 

146 

147def from_array(arr: np.ndarray) -> 'Image': 

148 """Create an image from a numpy array. 

149 

150 Args: 

151 arr: Pixel array, as returned by ``Image.to_array``. 

152 

153 Returns: 

154 An image object. 

155 

156 Raises: 

157 ``Error``: If the image data cannot be loaded. 

158 """ 

159 img = PIL.Image.fromarray(arr) 

160 return _new(img) 

161 

162 

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

164 """Create an image from an SVG document. Not implemented. 

165 

166 Args: 

167 xmlstr: SVG source. 

168 size: Image size ``(width, height)``. 

169 mime: MIME type. 

170 

171 Returns: 

172 An image object. 

173 

174 Raises: 

175 ``NotImplementedError``: Always. 

176 """ 

177 # @TODO rasterize svg 

178 raise NotImplementedError 

179 

180 

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

182 """Load the PIL image data and wrap the image in an ``Image``.""" 

183 try: 

184 img.load() 

185 except Exception as exc: 

186 raise Error from exc 

187 return Image(img) 

188 

189 

190class Image(gws.Image): 

191 """Image object, a wrapper around a PIL image.""" 

192 

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

194 """Create an image object. 

195 

196 Args: 

197 img: A PIL image. 

198 """ 

199 

200 self.img: PIL.Image.Image = img 

201 

202 def mode(self): 

203 return self.img.mode 

204 

205 def size(self): 

206 return self.img.size 

207 

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

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

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

211 return self 

212 

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

214 w, h = self.img.size 

215 if width and height: 

216 sz = (width, height) 

217 elif width: 

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

219 elif height: 

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

221 else: 

222 return self 

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

224 

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

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

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

228 return self 

229 

230 def crop(self, box): 

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

232 return self 

233 

234 def convert(self, mode): 

235 if self.img.mode != mode: 

236 self.img = self.img.convert(mode) 

237 return self 

238 

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

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

241 return self 

242 

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

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

245 

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

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

248 

249 if opacity < 1: 

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

251 oth.putalpha(alpha) 

252 

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

254 return self 

255 

256 def to_bytes(self, mime='', options=None): 

257 with io.BytesIO() as fp: 

258 self._save(fp, mime, options or {}) 

259 return fp.getvalue() 

260 

261 def to_base64(self, mime='', options=None): 

262 return gws.lib.text.to_base64(self.to_bytes(mime, options)) 

263 

264 def to_data_url(self, mime='', options=None): 

265 mime = mime or gws.lib.mime.PNG 

266 return gws.lib.text.to_data_url(self.to_bytes(mime, options), mime) 

267 

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

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

270 self._save(fp, mime, options or {}) 

271 return path 

272 

273 def _save(self, fp, mime: gws.MimeType, options: dict): 

274 fmt = mime_to_pil_format(mime) 

275 if not fmt: 

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

277 opts = dict(options or {}) 

278 img = self.img 

279 

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

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

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

283 img.alpha_composite(self.img) 

284 img = img.convert('RGB') 

285 

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

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

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

289 

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

291 

292 def to_array(self): 

293 return np.array(self.img) 

294 

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

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

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

298 font = PIL.ImageFont.load_default() 

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

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

301 return self 

302 

303 def add_box(self, color=None): 

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

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

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

307 x, y = self.img.size 

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

309 return self 

310 

311 def compare_to(self, other): 

312 error = 0 

313 x, y = self.size() 

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

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

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

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

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

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

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

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

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

323 

324 

325_MIME_TO_PIL_FORMAT = { 

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

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

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

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

330 gws.lib.mime.TIFF: 'TIFF', 

331 gws.lib.mime.JPEG2000: 'JPEG2000', 

332} 

333 

334 

335def mime_to_pil_format(mime: gws.MimeType): 

336 """Return the PIL format name for a MIME type, PNG if no MIME type is given, an empty string if it is not supported.""" 

337 

338 if not mime: 

339 return 'PNG' 

340 return _MIME_TO_PIL_FORMAT.get(gws.lib.mime.get_base(mime), '') 

341 

342 

343def _int_size(size: gws.Size): 

344 w, h = size 

345 return int(w), int(h)