Coverage for src/io_collection/load/load_image.py: 100%

1from __future__ import annotations 

2 

3from pathlib import Path 

4 

5import bioio_ome_tiff 

6import bioio_tifffile 

7from bioio import BioImage 

8 

9EXTENSIONS = (".tif", ".tiff", ".ome.tif", ".ome.tiff") 

10 

11 

12def load_image(location: str, key: str, dim_order: str | None = None) -> BioImage: 

13 """ 

14 Load key as BioIO image from specified location. 

15 

16 Method will load from the S3 bucket if the location begins with the 

17 **s3://** protocol, otherwise it assumes the location is a local path. 

18 

19 Method currently only supports `.tiff` and `ome.tiff` images. 

20 

21 Parameters 

22 ---------- 

23 location 

24 Object location (local path or S3 bucket). 

25 key 

26 Object key ending in `.tiff` or `ome.tiff`. 

27 dim_order 

28 Image dimensions order. 

29 

30 Returns 

31 ------- 

32 : 

33 Loaded image. 

34 """ 

35 

36 if not key.endswith(EXTENSIONS): 

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

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

39 raise ValueError(message) 

40 

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

42 return _load_image_from_s3(location[5:], key, dim_order) 

43 return _load_image_from_fs(location, key, dim_order) 

44 

45 

46def _load_image_from_fs(path: str, key: str, dim_order: str | None = None) -> BioImage: 

47 """ 

48 Load key as BioIO image from local file system. 

49 

50 Parameters 

51 ---------- 

52 path 

53 Local object path. 

54 key 

55 Object key ending in `.tiff` or `ome.tiff`. 

56 dim_order 

57 Image dimensions order. 

58 

59 Returns 

60 ------- 

61 : 

62 Loaded image. 

63 """ 

64 

65 full_path = Path(path) / key 

66 return BioImage( 

67 str(full_path), 

68 reader=bioio_ome_tiff.Reader if ".ome.tif" in key else bioio_tifffile.Reader, 

69 dim_order=dim_order, 

70 ) 

71 

72 

73def _load_image_from_s3(bucket: str, key: str, dim_order: str | None = None) -> BioImage: 

74 """ 

75 Load key as BioIO image from AWS S3 bucket. 

76 

77 Parameters 

78 ---------- 

79 bucket 

80 AWS S3 bucket name. 

81 key 

82 Object key ending in `.tiff` or `ome.tiff`. 

83 dim_order 

84 Image dimensions order. 

85 

86 Returns 

87 ------- 

88 : 

89 Loaded image. 

90 """ 

91 

92 full_key = f"s3://{bucket}/{key}" 

93 return BioImage( 

94 full_key, 

95 reader=bioio_ome_tiff.Reader if ".ome.tif" in key else bioio_tifffile.Reader, 

96 dim_order=dim_order, 

97 )