Coverage for tests/objects/test_draw.py: 100%

1import numpy as np 

2import shapely 

3from PIL import Image, ImageDraw 

4from qubalab.objects.draw import draw_geometry 

5 

6 

7def test_draw_point(): 

8 v = 60 

9 geometry = shapely.Point(1, 3) 

10 pixel_values = np.zeros((10, 5)) 

11 expected_pixel_values = np.array( 

12 [[0, 0, 0, 0, 0], 

13 [0, 0, 0, 0, 0], 

14 [0, 0, 0, 0, 0], 

15 [0, v, 0, 0, 0], 

16 [0, 0, 0, 0, 0], 

17 [0, 0, 0, 0, 0], 

18 [0, 0, 0, 0, 0], 

19 [0, 0, 0, 0, 0], 

20 [0, 0, 0, 0, 0], 

21 [0, 0, 0, 0, 0],] 

22 ) 

23 image = Image.fromarray(pixel_values) 

24 drawing_context = ImageDraw.Draw(image) 

25 

26 draw_geometry(image.size, drawing_context, geometry, v) 

27 

28 np.testing.assert_array_equal(np.asarray(image), expected_pixel_values) 

29 

30 

31def test_draw_line_string(): 

32 v = 32 

33 geometry = shapely.LineString([(1, 3), (2, 3), (4, 4), (1, 8)]) 

34 pixel_values = np.zeros((10, 5)) 

35 expected_pixel_values = np.array( 

36 [[0, 0, 0, 0, 0], 

37 [0, 0, 0, 0, 0], 

38 [0, 0, 0, 0, 0], 

39 [0, v, v, 0, 0], 

40 [0, 0, 0, v, v], 

41 [0, 0, 0, v, 0], 

42 [0, 0, v, 0, 0], 

43 [0, 0, v, 0, 0], 

44 [0, v, 0, 0, 0], 

45 [0, 0, 0, 0, 0],] 

46 ) 

47 image = Image.fromarray(pixel_values) 

48 drawing_context = ImageDraw.Draw(image) 

49 

50 draw_geometry(image.size, drawing_context, geometry, v) 

51 

52 np.testing.assert_array_equal(np.asarray(image), expected_pixel_values) 

53 

54 

55def test_draw_linear_ring(): 

56 v = 57 

57 geometry = shapely.LinearRing([(1, 1), (3, 4), (1, 7)]) 

58 pixel_values = np.zeros((10, 5)) 

59 expected_pixel_values = np.array( 

60 [[0, 0, 0, 0, 0], 

61 [0, v, 0, 0, 0], 

62 [0, v, v, 0, 0], 

63 [0, v, v, 0, 0], 

64 [0, v, 0, v, 0], 

65 [0, v, v, 0, 0], 

66 [0, v, v, 0, 0], 

67 [0, v, 0, 0, 0], 

68 [0, 0, 0, 0, 0], 

69 [0, 0, 0, 0, 0],] 

70 ) 

71 image = Image.fromarray(pixel_values) 

72 drawing_context = ImageDraw.Draw(image) 

73 

74 draw_geometry(image.size, drawing_context, geometry, v) 

75 

76 np.testing.assert_array_equal(np.asarray(image), expected_pixel_values) 

77 

78 

79def test_draw_polygon_without_holes(): 

80 v = 573 

81 geometry = shapely.Polygon([(0, 0), (3, 3), (3, 4), (0, 7)]) 

82 pixel_values = np.zeros((10, 5)) 

83 expected_pixel_values = np.array( 

84 [[v, 0, 0, 0, 0], 

85 [v, v, 0, 0, 0], 

86 [v, v, v, 0, 0], 

87 [v, v, v, v, 0], 

88 [v, v, v, v, 0], 

89 [v, v, v, 0, 0], 

90 [v, v, 0, 0, 0], 

91 [v, 0, 0, 0, 0], 

92 [0, 0, 0, 0, 0], 

93 [0, 0, 0, 0, 0],] 

94 ) 

95 image = Image.fromarray(pixel_values) 

96 drawing_context = ImageDraw.Draw(image) 

97 

98 draw_geometry(image.size, drawing_context, geometry, v) 

99 

100 np.testing.assert_array_equal(np.asarray(image), expected_pixel_values) 

101 

102 

103def test_draw_polygon_with_holes(): 

104 v = 9 

105 geometry = shapely.Polygon([(0, 0), (3, 3), (3, 4), (0, 7)], holes=[[(1, 3), (2, 3), (2, 4)]]) 

106 pixel_values = np.zeros((10, 5)) 

107 expected_pixel_values = np.array( 

108 [[v, 0, 0, 0, 0], 

109 [v, v, 0, 0, 0], 

110 [v, v, v, 0, 0], 

111 [v, 0, 0, v, 0], 

112 [v, v, 0, v, 0], 

113 [v, v, v, 0, 0], 

114 [v, v, 0, 0, 0], 

115 [v, 0, 0, 0, 0], 

116 [0, 0, 0, 0, 0], 

117 [0, 0, 0, 0, 0],] 

118 ) 

119 image = Image.fromarray(pixel_values) 

120 drawing_context = ImageDraw.Draw(image) 

121 

122 draw_geometry(image.size, drawing_context, geometry, v) 

123 

124 np.testing.assert_array_equal(np.asarray(image), expected_pixel_values) 

125 

126 

127def test_draw_multi_point(): 

128 v = 4 

129 geometry = shapely.MultiPoint([shapely.Point(1, 3), shapely.Point(3, 6)]) 

130 pixel_values = np.zeros((10, 5)) 

131 expected_pixel_values = np.array( 

132 [[0, 0, 0, 0, 0], 

133 [0, 0, 0, 0, 0], 

134 [0, 0, 0, 0, 0], 

135 [0, v, 0, 0, 0], 

136 [0, 0, 0, 0, 0], 

137 [0, 0, 0, 0, 0], 

138 [0, 0, 0, v, 0], 

139 [0, 0, 0, 0, 0], 

140 [0, 0, 0, 0, 0], 

141 [0, 0, 0, 0, 0],] 

142 ) 

143 image = Image.fromarray(pixel_values) 

144 drawing_context = ImageDraw.Draw(image) 

145 

146 draw_geometry(image.size, drawing_context, geometry, v) 

147 

148 np.testing.assert_array_equal(np.asarray(image), expected_pixel_values) 

149 

150 

151def test_draw_multi_line_string(): 

152 v = 32 

153 geometry = shapely.MultiLineString([shapely.LineString([(1, 3), (2, 3), (4, 4), (1, 8)]), shapely.LineString([(4, 2), (4, 7)])]) 

154 pixel_values = np.zeros((10, 5)) 

155 expected_pixel_values = np.array( 

156 [[0, 0, 0, 0, 0], 

157 [0, 0, 0, 0, 0], 

158 [0, 0, 0, 0, v], 

159 [0, v, v, 0, v], 

160 [0, 0, 0, v, v], 

161 [0, 0, 0, v, v], 

162 [0, 0, v, 0, v], 

163 [0, 0, v, 0, v], 

164 [0, v, 0, 0, 0], 

165 [0, 0, 0, 0, 0],] 

166 ) 

167 image = Image.fromarray(pixel_values) 

168 drawing_context = ImageDraw.Draw(image) 

169 

170 draw_geometry(image.size, drawing_context, geometry, v) 

171 

172 np.testing.assert_array_equal(np.asarray(image), expected_pixel_values) 

173 

174 

175def test_draw_multi_polygon(): 

176 v = 573 

177 geometry = shapely.MultiPolygon([ 

178 shapely.Polygon([(0, 0), (3, 3), (3, 4), (0, 7)]), 

179 shapely.Polygon([(2, 7), (4, 7), (4, 9), (2, 9)], holes=[[(3, 8), (3, 8), (3, 8)]]) 

180 ]) 

181 pixel_values = np.zeros((10, 5)) 

182 expected_pixel_values = np.array( 

183 [[v, 0, 0, 0, 0], 

184 [v, v, 0, 0, 0], 

185 [v, v, v, 0, 0], 

186 [v, v, v, v, 0], 

187 [v, v, v, v, 0], 

188 [v, v, v, 0, 0], 

189 [v, v, 0, 0, 0], 

190 [v, 0, v, v, v], 

191 [0, 0, v, 0, v], 

192 [0, 0, v, v, v],] 

193 ) 

194 image = Image.fromarray(pixel_values) 

195 drawing_context = ImageDraw.Draw(image) 

196 

197 draw_geometry(image.size, drawing_context, geometry, v) 

198 

199 np.testing.assert_array_equal(np.asarray(image), expected_pixel_values) 

200 

201 

202def test_draw_geometry_collection(): 

203 v = 7 

204 geometry = shapely.GeometryCollection([shapely.Point(1, 2), shapely.LineString([(1, 3), (2, 3), (4, 4), (1, 8)])]) 

205 pixel_values = np.zeros((10, 5)) 

206 expected_pixel_values = np.array( 

207 [[0, 0, 0, 0, 0], 

208 [0, 0, 0, 0, 0], 

209 [0, v, 0, 0, 0], 

210 [0, v, v, 0, 0], 

211 [0, 0, 0, v, v], 

212 [0, 0, 0, v, 0], 

213 [0, 0, v, 0, 0], 

214 [0, 0, v, 0, 0], 

215 [0, v, 0, 0, 0], 

216 [0, 0, 0, 0, 0],] 

217 ) 

218 image = Image.fromarray(pixel_values) 

219 drawing_context = ImageDraw.Draw(image) 

220 

221 draw_geometry(image.size, drawing_context, geometry, v) 

222 

223 np.testing.assert_array_equal(np.asarray(image), expected_pixel_values) 

224 

225 

226def test_draw_when_values_already_present(): 

227 i = 6 

228 v = 3 

229 geometry = shapely.Point(1, 3) 

230 pixel_values = np.full((10, 5), i, dtype=np.uint8) 

231 expected_pixel_values = np.array( 

232 [[i, i, i, i, i], 

233 [i, i, i, i, i], 

234 [i, i, i, i, i], 

235 [i, v, i, i, i], 

236 [i, i, i, i, i], 

237 [i, i, i, i, i], 

238 [i, i, i, i, i], 

239 [i, i, i, i, i], 

240 [i, i, i, i, i], 

241 [i, i, i, i, i],] 

242 ) 

243 image = Image.fromarray(pixel_values) 

244 drawing_context = ImageDraw.Draw(image) 

245 

246 draw_geometry(image.size, drawing_context, geometry, v) 

247 

248 np.testing.assert_array_equal(np.asarray(image), expected_pixel_values)