Coverage for src/io_collection/save/save_image.py: 100%

1import io 

2import tempfile 

3from pathlib import Path 

4 

5import boto3 

6import numpy as np 

7from bioio.writers import OmeTiffWriter 

8from PIL import Image 

9 

10from io_collection.save.save_buffer import _save_buffer_to_s3 

11 

12EXTENSIONS = (".ome.tif", ".ome.tiff", ".png") 

13 

14 

15def save_image(location: str, key: str, image: np.ndarray) -> None: 

16 """ 

17 Save image array to key at specified location. 

18 

19 Method will save to the S3 bucket if the location begins with the **s3://** 

20 protocol, otherwise it assumes the location is a local path. 

21 

22 Parameters 

23 ---------- 

24 location 

25 Object location (local path or S3 bucket). 

26 key 

27 Object key ending in `.ome.tiff` or `.png`. 

28 image 

29 Image array. 

30 """ 

31 

32 if not key.endswith(EXTENSIONS): 

33 extensions = " | ".join([ext[1:] for ext in EXTENSIONS]) 

34 message = f"key [ {key} ] must have [ {extensions} ] extension" 

35 raise ValueError(message) 

36 

37 if location[:5] == "s3://": 

38 _save_image_to_s3(location[5:], key, image) 

39 else: 

40 _save_image_to_fs(location, key, image) 

41 

42 

43def _save_image_to_fs(path: str, key: str, image: np.ndarray) -> None: 

44 """ 

45 Save image array to key on local file system. 

46 

47 Parameters 

48 ---------- 

49 path 

50 Local object path. 

51 key 

52 Object key ending in `.ome.tiff` or `.png`. 

53 image 

54 Image array. 

55 """ 

56 

57 full_path = Path(path) / key 

58 full_path.parent.mkdir(parents=True, exist_ok=True) 

59 

60 if key.endswith((".ome.tiff", ".ome.tif")): 

61 OmeTiffWriter.save(image, full_path) 

62 elif key.endswith(".png"): 

63 Image.fromarray(image).save(full_path) # type: ignore[no-untyped-call] 

64 

65 

66def _save_image_to_s3(bucket: str, key: str, image: np.ndarray) -> None: 

67 """ 

68 Save image array to key in AWS S3 bucket. 

69 

70 Parameters 

71 ---------- 

72 bucket 

73 AWS S3 bucket name. 

74 key 

75 Object key ending in `.ome.tiff` or `.png`. 

76 image 

77 Image array. 

78 """ 

79 

80 s3_client = boto3.client("s3") 

81 

82 if key.endswith((".ome.tiff", ".ome.tif")): 

83 with tempfile.TemporaryDirectory() as temp_dir: 

84 temp_path = Path(temp_dir) / "temp.ome.tiff" 

85 OmeTiffWriter.save(image, temp_path) 

86 

87 with temp_path.open("rb") as fileobj: 

88 s3_client.upload_fileobj(fileobj, bucket, key) 

89 elif key.endswith(".png"): 

90 with io.BytesIO() as buffer: 

91 Image.fromarray(image).save(buffer, format="png") # type: ignore[no-untyped-call] 

92 _save_buffer_to_s3(bucket, key, buffer, "image/png")