Coverage for src/lightmlrestapi/args/args_images.py: 79%

28 statements  

« prev     ^ index     » next       coverage.py v6.4.1, created at 2022-06-06 07:16 +0200

1""" 

2@file 

3@brief Playground with images. 

4""" 

5import io 

6import os 

7import base64 

8import numpy 

9from PIL import Image 

10 

11 

12def image2array(img): 

13 """ 

14 Converts a color imaged into an array. 

15 

16 @param img :epkg:`PIL:Image.Image` 

17 @return :epkg:`numpy:array` 

18 """ 

19 im_arr = numpy.frombuffer(img.tobytes(), dtype=numpy.uint8) 

20 return im_arr.reshape((img.size[1], img.size[0], 3)) # pylint: disable=E1121 

21 

22 

23def image2base64(path, fmt='png'): 

24 """ 

25 Encodes an image into :epkg:`*pyf:base64`. 

26 

27 @param path filename or an image 

28 @param fmt if the image is given as an image (:epkg:`Pillow`), 

29 it must be first saved in a specific format (png, jpg, ...) 

30 @return format, base64 

31 

32 The format is the file extension. 

33 """ 

34 if isinstance(path, bytes): 

35 content = path 

36 ext = None 

37 elif hasattr(path, 'convert'): 

38 # Most probably a Pillow object 

39 st = io.BytesIO() 

40 path.save(st, format=fmt) 

41 content = st.getvalue() 

42 ext = 'png' 

43 else: 

44 ext = os.path.splitext(path)[-1].lower().strip('.') 

45 with open(path, "rb") as f: 

46 content = f.read() 

47 return 'image/' + ext, base64.b64encode(content) 

48 

49 

50def base642image(encoded): 

51 """ 

52 Gets an encoded image and builds an 

53 :epkg:`PIL:Image.Image` from it. 

54 

55 @param encoded :epkg:`*pyf:base64` encoded image 

56 (see @see fn image2base64) 

57 @return :epkg:`PIL:Image.Image` 

58 """ 

59 cont = base64.b64decode(encoded) 

60 return Image.open(io.BytesIO(cont)) 

61 

62 

63def bytes2string(content): 

64 """ 

65 Converts bytes to string. 

66 """ 

67 return base64.b64encode(content) 

68 

69 

70def string2bytes(content): 

71 """ 

72 Converts string to bytes. 

73 """ 

74 return base64.b64decode(content)