Coverage for src/deepdraw/data/loader.py: 73%

1# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch> 

2# 

3# SPDX-License-Identifier: GPL-3.0-or-later 

4 

5"""Data loading code.""" 

6 

7 

8import functools 

9import os 

10 

11import numpy 

12import PIL.Image 

13import PIL.ImageFile 

14import skimage.exposure 

15 

16from .sample import DelayedSample 

17 

18PIL.ImageFile.LOAD_TRUNCATED_IMAGES = True 

19 

20 

21def load_pil_raw_12bit_jsrt(path, width): 

22 """Loads a raw 16-bit sample data. 

23 

24 This method was designed to handle the raw images from the JSRT_ dataset. 

25 It reads the data file and applies a simple histogram equalization to the 

26 8-bit representation of the image to obtain something along the lines of 

27 the PNG (unofficial) version distributed at `JSRT-Kaggle`_. 

28 

29 

30 Parameters 

31 ---------- 

32 

33 path : str 

34 The full path leading to the image to be loaded 

35 

36 width : int 

37 The desired width of the output image 

38 

39 

40 Returns 

41 ------- 

42 

43 image : PIL.Image.Image 

44 A PIL image in RGB mode, with `width`x`width` pixels 

45 

46 

47 .. include:: ../../links.rst 

48 """ 

49 

50 raw_image = numpy.fromfile(path, numpy.dtype(">u2")).reshape(2048, 2048) 

51 raw_image[raw_image > 4095] = 4095 

52 raw_image = 4095 - raw_image # invert colors 

53 raw_image = (raw_image >> 4).astype(numpy.uint8) # 8-bit uint 

54 raw_image = skimage.exposure.equalize_hist(raw_image) 

55 return ( 

56 PIL.Image.fromarray((raw_image * 255).astype(numpy.uint8)) 

57 .resize((width, width)) 

58 .convert("RGB") 

59 ) 

60 

61 

62def load_pil_rgb(path): 

63 """Loads a sample data. 

64 

65 Parameters 

66 ---------- 

67 

68 path : str 

69 The full path leading to the image to be loaded 

70 

71 

72 Returns 

73 ------- 

74 

75 image : PIL.Image.Image 

76 A PIL image in RGB mode 

77 """ 

78 

79 return PIL.Image.open(path).convert("RGB") 

80 

81 

82def load_pil_1(path): 

83 """Loads a sample binary label or mask. 

84 

85 Parameters 

86 ---------- 

87 

88 path : str 

89 The full path leading to the image to be loaded 

90 

91 

92 Returns 

93 ------- 

94 

95 image : PIL.Image.Image 

96 A PIL image in mode "1" 

97 """ 

98 

99 return PIL.Image.open(path).convert(mode="1", dither=None) 

100 

101 

102def make_delayed(sample, loader, key=None): 

103 """Returns a delayed-loading Sample object. 

104 

105 Parameters 

106 ---------- 

107 

108 sample : dict 

109 A dictionary that maps field names to sample data values (e.g. paths) 

110 

111 loader : object 

112 A function that inputs ``sample`` dictionaries and returns the loaded 

113 data. 

114 

115 key : str 

116 A unique key identifier for this sample. If not provided, assumes 

117 ``sample`` is a dictionary with a ``data`` entry and uses its path as 

118 key. 

119 

120 

121 Returns 

122 ------- 

123 

124 sample : deepdraw.data.sample.DelayedSample 

125 In which ``key`` is as provided and ``data`` can be accessed to trigger 

126 sample loading. 

127 """ 

128 

129 return DelayedSample( 

130 functools.partial(loader, sample), 

131 key=key or os.path.splitext(sample["data"])[0], 

132 )