1#!/usr/bin/env python 

2# -*- coding: utf-8 -*- 

3 

4import os 

5import random 

6 

7import numpy 

8import PIL.Image 

9import PIL.ImageDraw 

10import PIL.ImageOps 

11import pkg_resources 

12import torch 

13import torchvision.transforms.functional 

14 

15from ..data.transforms import ( 

16 CenterCrop, 

17 ColorJitter, 

18 Compose, 

19 Crop, 

20 Pad, 

21 RandomHorizontalFlip, 

22 RandomRotation, 

23 RandomVerticalFlip, 

24 Resize, 

25 ResizeCrop, 

26 SingleAutoLevel16to8, 

27 ToTensor, 

28) 

29 

30 

31def _create_img(size): 

32 t = torch.randn(size) 

33 pil = torchvision.transforms.functional.to_pil_image(t) 

34 return pil 

35 

36 

37def test_center_crop(): 

38 

39 # parameters 

40 im_size = (3, 22, 20) # (planes, height, width) 

41 crop_size = (10, 12) # (height, width) 

42 

43 # test 

44 bh = (im_size[1] - crop_size[0]) // 2 

45 bw = (im_size[2] - crop_size[1]) // 2 

46 idx = (slice(bh, -bh), slice(bw, -bw), slice(0, im_size[0])) 

47 transforms = CenterCrop(crop_size) 

48 img, gt, mask = [_create_img(im_size) for i in range(3)] 

49 assert img.size == (im_size[2], im_size[1]) # confirms the above 

50 img_t, gt_t, mask_t = transforms(img, gt, mask) 

51 assert img_t.size == (crop_size[1], crop_size[0]) # confirms the above 

52 # notice that PIL->array does array.transpose(1, 2, 0) 

53 # so it creates an array that is (height, width, planes) 

54 assert numpy.all(numpy.array(img_t) == numpy.array(img)[idx]) 

55 assert numpy.all(numpy.array(gt_t) == numpy.array(gt)[idx]) 

56 assert numpy.all(numpy.array(mask_t) == numpy.array(mask)[idx]) 

57 

58 

59def test_center_crop_uneven(): 

60 

61 # parameters - WARNING: notice that modern implementations of centercrop in 

62 # torchvision (noticed with version 0.10.x) may provide slightly different 

63 # outputs depending on the multiplicity of pixels that need to be cropped 

64 # from both sides. In the current configuration, the output is such that 

65 # the centercrop is 0.5 pixel to the right/bottom from an actual center 

66 # crop. With other configurations, it may be that the crop is 0.5 pixel to 

67 # the left/top part of the image, which makes testing this function a bit 

68 # hard. If you have issues with this function, try different 

69 # configurations here, or to change the setting of "idx" a few lines down. 

70 im_size = (3, 21, 20) # (planes, height, width) 

71 crop_size = (10, 13) # (height, width) 

72 

73 # test 

74 bh = (im_size[1] - crop_size[0]) // 2 

75 bw = (im_size[2] - crop_size[1]) // 2 

76 # when the crop size is uneven, this is what happens - notice here that the 

77 # image height is uneven, and the crop width as well - the attributions of 

78 # extra pixels will depend on what is uneven (original image or crop) 

79 # crop is 0.5 pixel to the right/bottom: 

80 idx = (slice(bh + 1, -bh), slice(bw + 1, -bw), slice(0, im_size[0])) 

81 # crop is 0.5 pixel to the left/top: 

82 # idx = (slice(bh, -(bh + 1)), slice(bw, -(bw + 1)), slice(0, im_size[0])) 

83 transforms = CenterCrop(crop_size) 

84 img, gt, mask = [_create_img(im_size) for i in range(3)] 

85 assert img.size == (im_size[2], im_size[1]) # confirms the above 

86 img_t, gt_t, mask_t = transforms(img, gt, mask) 

87 assert img_t.size == (crop_size[1], crop_size[0]) # confirms the above 

88 # notice that PIL->array does array.transpose(1, 2, 0) 

89 # so it creates an array that is (height, width, planes) 

90 assert numpy.all(numpy.array(img_t) == numpy.array(img)[idx]) 

91 assert numpy.all(numpy.array(gt_t) == numpy.array(gt)[idx]) 

92 assert numpy.all(numpy.array(mask_t) == numpy.array(mask)[idx]) 

93 

94 

95def test_pad_default(): 

96 

97 # parameters 

98 im_size = (3, 22, 20) # (planes, height, width) 

99 pad_size = 2 

100 

101 # test 

102 idx = ( 

103 slice(pad_size, -pad_size), 

104 slice(pad_size, -pad_size), 

105 slice(0, im_size[0]), 

106 ) 

107 transforms = Pad(pad_size) 

108 img, gt, mask = [_create_img(im_size) for i in range(3)] 

109 assert img.size == (im_size[2], im_size[1]) # confirms the above 

110 img_t, gt_t, mask_t = transforms(img, gt, mask) 

111 # notice that PIL->array does array.transpose(1, 2, 0) 

112 # so it creates an array that is (height, width, planes) 

113 assert numpy.all(numpy.array(img_t)[idx] == numpy.array(img)) 

114 assert numpy.all(numpy.array(gt_t)[idx] == numpy.array(gt)) 

115 assert numpy.all(numpy.array(mask_t)[idx] == numpy.array(mask)) 

116 

117 # checks that the border introduced with padding is all about "fill" 

118 img_t = numpy.array(img_t) 

119 img_t[idx] = 0 

120 # border_size_plane = img_t[:, :, 0].size - numpy.array(img)[:, :, 0].size 

121 assert img_t.sum() == 0 

122 

123 gt_t = numpy.array(gt_t) 

124 gt_t[idx] = 0 

125 assert gt_t.sum() == 0 

126 

127 mask_t = numpy.array(mask_t) 

128 mask_t[idx] = 0 

129 assert mask_t.sum() == 0 

130 

131 

132def test_pad_2tuple(): 

133 

134 # parameters 

135 im_size = (3, 22, 20) # (planes, height, width) 

136 pad_size = (1, 2) # left/right, top/bottom 

137 fill = (3, 4, 5) 

138 

139 # test 

140 idx = ( 

141 slice(pad_size[1], -pad_size[1]), 

142 slice(pad_size[0], -pad_size[0]), 

143 slice(0, im_size[0]), 

144 ) 

145 transforms = Pad(pad_size, fill) 

146 img, gt, mask = [_create_img(im_size) for i in range(3)] 

147 assert img.size == (im_size[2], im_size[1]) # confirms the above 

148 img_t, gt_t, mask_t = transforms(img, gt, mask) 

149 # notice that PIL->array does array.transpose(1, 2, 0) 

150 # so it creates an array that is (height, width, planes) 

151 assert numpy.all(numpy.array(img_t)[idx] == numpy.array(img)) 

152 assert numpy.all(numpy.array(gt_t)[idx] == numpy.array(gt)) 

153 assert numpy.all(numpy.array(mask_t)[idx] == numpy.array(mask)) 

154 

155 # checks that the border introduced with padding is all about "fill" 

156 img_t = numpy.array(img_t) 

157 img_t[idx] = 0 

158 border_size_plane = img_t[:, :, 0].size - numpy.array(img)[:, :, 0].size 

159 expected_sum = sum((fill[k] * border_size_plane) for k in range(3)) 

160 assert img_t.sum() == expected_sum 

161 

162 gt_t = numpy.array(gt_t) 

163 gt_t[idx] = 0 

164 assert gt_t.sum() == expected_sum 

165 

166 mask_t = numpy.array(mask_t) 

167 mask_t[idx] = 0 

168 assert mask_t.sum() == expected_sum 

169 

170 

171def test_pad_4tuple(): 

172 

173 # parameters 

174 im_size = (3, 22, 20) # (planes, height, width) 

175 pad_size = (1, 2, 3, 4) # left, top, right, bottom 

176 fill = (3, 4, 5) 

177 

178 # test 

179 idx = ( 

180 slice(pad_size[1], -pad_size[3]), 

181 slice(pad_size[0], -pad_size[2]), 

182 slice(0, im_size[0]), 

183 ) 

184 transforms = Pad(pad_size, fill) 

185 img, gt, mask = [_create_img(im_size) for i in range(3)] 

186 assert img.size == (im_size[2], im_size[1]) # confirms the above 

187 img_t, gt_t, mask_t = transforms(img, gt, mask) 

188 # notice that PIL->array does array.transpose(1, 2, 0) 

189 # so it creates an array that is (height, width, planes) 

190 assert numpy.all(numpy.array(img_t)[idx] == numpy.array(img)) 

191 assert numpy.all(numpy.array(gt_t)[idx] == numpy.array(gt)) 

192 assert numpy.all(numpy.array(mask_t)[idx] == numpy.array(mask)) 

193 

194 # checks that the border introduced with padding is all about "fill" 

195 img_t = numpy.array(img_t) 

196 img_t[idx] = 0 

197 border_size_plane = img_t[:, :, 0].size - numpy.array(img)[:, :, 0].size 

198 expected_sum = sum((fill[k] * border_size_plane) for k in range(3)) 

199 assert img_t.sum() == expected_sum 

200 

201 gt_t = numpy.array(gt_t) 

202 gt_t[idx] = 0 

203 assert gt_t.sum() == expected_sum 

204 

205 mask_t = numpy.array(mask_t) 

206 mask_t[idx] = 0 

207 assert mask_t.sum() == expected_sum 

208 

209 

210def test_resize_downscale_w(): 

211 

212 # parameters 

213 im_size = (3, 22, 20) # (planes, height, width) 

214 new_size = 10 # (smallest edge) 

215 

216 # test 

217 transforms = Resize(new_size) 

218 img, gt, mask = [_create_img(im_size) for i in range(3)] 

219 assert img.size == (im_size[2], im_size[1]) # confirms the above 

220 img_t, gt_t, mask_t = transforms(img, gt, mask) 

221 new_size = (new_size, (new_size * im_size[1]) / im_size[2]) 

222 assert img_t.size == new_size 

223 assert gt_t.size == new_size 

224 assert mask_t.size == new_size 

225 

226 

227def test_resize_downscale_hw(): 

228 

229 # parameters 

230 im_size = (3, 22, 20) # (planes, height, width) 

231 new_size = (10, 12) # (height, width) 

232 

233 # test 

234 transforms = Resize(new_size) 

235 img, gt, mask = [_create_img(im_size) for i in range(3)] 

236 assert img.size == (im_size[2], im_size[1]) # confirms the above 

237 img_t, gt_t, mask_t = transforms(img, gt, mask) 

238 assert img_t.size == (new_size[1], new_size[0]) 

239 assert gt_t.size == (new_size[1], new_size[0]) 

240 assert mask_t.size == (new_size[1], new_size[0]) 

241 

242 

243def test_crop(): 

244 

245 # parameters 

246 im_size = (3, 22, 20) # (planes, height, width) 

247 crop_size = (3, 2, 10, 12) # (upper, left, height, width) 

248 

249 # test 

250 idx = ( 

251 slice(crop_size[0], crop_size[0] + crop_size[2]), 

252 slice(crop_size[1], crop_size[1] + crop_size[3]), 

253 slice(0, im_size[0]), 

254 ) 

255 transforms = Crop(*crop_size) 

256 img, gt, mask = [_create_img(im_size) for i in range(3)] 

257 assert img.size == (im_size[2], im_size[1]) # confirms the above 

258 img_t, gt_t, mask_t = transforms(img, gt, mask) 

259 # notice that PIL->array does array.transpose(1, 2, 0) 

260 # so it creates an array that is (height, width, planes) 

261 assert numpy.all(numpy.array(img_t) == numpy.array(img)[idx]) 

262 assert numpy.all(numpy.array(gt_t) == numpy.array(gt)[idx]) 

263 assert numpy.all(numpy.array(mask_t) == numpy.array(mask)[idx]) 

264 

265 

266def test_to_tensor(): 

267 

268 transforms = ToTensor() 

269 img, gt, mask = [_create_img((3, 5, 5)) for i in range(3)] 

270 gt = gt.convert("1", dither=None) 

271 mask = mask.convert("1", dither=None) 

272 img_t, gt_t, mask_t = transforms(img, gt, mask) 

273 assert img_t.dtype == torch.float32 

274 assert gt_t.dtype == torch.float32 

275 assert mask_t.dtype == torch.float32 

276 

277 

278def test_horizontal_flip(): 

279 

280 transforms = RandomHorizontalFlip(p=1) 

281 

282 im_size = (3, 24, 42) # (planes, height, width) 

283 img, gt, mask = [_create_img(im_size) for i in range(3)] 

284 img_t, gt_t, mask_t = transforms(img, gt, mask) 

285 

286 # notice that PIL->array does array.transpose(1, 2, 0) 

287 # so it creates an array that is (height, width, planes) 

288 assert numpy.all(numpy.flip(img_t, axis=1) == numpy.array(img)) 

289 assert numpy.all(numpy.flip(gt_t, axis=1) == numpy.array(gt)) 

290 assert numpy.all(numpy.flip(mask_t, axis=1) == numpy.array(mask)) 

291 

292 

293def test_vertical_flip(): 

294 

295 transforms = RandomVerticalFlip(p=1) 

296 

297 im_size = (3, 24, 42) # (planes, height, width) 

298 img, gt, mask = [_create_img(im_size) for i in range(3)] 

299 img_t, gt_t, mask_t = transforms(img, gt, mask) 

300 

301 # notice that PIL->array does array.transpose(1, 2, 0) 

302 # so it creates an array that is (height, width, planes) 

303 assert numpy.all(numpy.flip(img_t, axis=0) == numpy.array(img)) 

304 assert numpy.all(numpy.flip(gt_t, axis=0) == numpy.array(gt)) 

305 assert numpy.all(numpy.flip(mask_t, axis=0) == numpy.array(mask)) 

306 

307 

308def test_rotation(): 

309 

310 im_size = (3, 24, 42) # (planes, height, width) 

311 transforms = RandomRotation(degrees=90, p=1) 

312 img = _create_img(im_size) 

313 

314 # asserts all images are rotated the same 

315 # and they are different from the original 

316 random.seed(42) 

317 img1_t, img2_t, img3_t = transforms(img, img, img) 

318 assert img1_t.size == (im_size[2], im_size[1]) 

319 assert numpy.all(numpy.array(img1_t) == numpy.array(img2_t)) 

320 assert numpy.all(numpy.array(img1_t) == numpy.array(img3_t)) 

321 assert numpy.any(numpy.array(img1_t) != numpy.array(img)) 

322 

323 # asserts two random transforms are not the same 

324 (img_t2,) = transforms(img) 

325 assert numpy.any(numpy.array(img_t2) != numpy.array(img1_t)) 

326 

327 

328def test_color_jitter(): 

329 

330 im_size = (3, 24, 42) # (planes, height, width) 

331 transforms = ColorJitter(p=1) 

332 img = _create_img(im_size) 

333 

334 # asserts only the first image is jittered 

335 # and it is different from the original 

336 # all others match the input data 

337 random.seed(42) 

338 img1_t, img2_t, img3_t = transforms(img, img, img) 

339 assert img1_t.size == (im_size[2], im_size[1]) 

340 assert numpy.any(numpy.array(img1_t) != numpy.array(img)) 

341 assert numpy.any(numpy.array(img1_t) != numpy.array(img2_t)) 

342 assert numpy.all(numpy.array(img2_t) == numpy.array(img3_t)) 

343 assert numpy.all(numpy.array(img2_t) == numpy.array(img)) 

344 

345 # asserts two random transforms are not the same 

346 img1_t2, img2_t2, img3_t2 = transforms(img, img, img) 

347 assert numpy.any(numpy.array(img1_t2) != numpy.array(img1_t)) 

348 assert numpy.all(numpy.array(img2_t2) == numpy.array(img)) 

349 assert numpy.all(numpy.array(img3_t2) == numpy.array(img)) 

350 

351 

352def test_compose(): 

353 

354 transforms = Compose( 

355 [ 

356 RandomVerticalFlip(p=1), 

357 RandomHorizontalFlip(p=1), 

358 RandomVerticalFlip(p=1), 

359 RandomHorizontalFlip(p=1), 

360 ] 

361 ) 

362 

363 img, gt, mask = [_create_img((3, 24, 42)) for i in range(3)] 

364 img_t, gt_t, mask_t = transforms(img, gt, mask) 

365 assert numpy.all(numpy.array(img_t) == numpy.array(img)) 

366 assert numpy.all(numpy.array(gt_t) == numpy.array(gt)) 

367 assert numpy.all(numpy.array(mask_t) == numpy.array(mask)) 

368 

369 

370def test_16bit_autolevel(): 

371 

372 path = pkg_resources.resource_filename( 

373 __name__, os.path.join("data", "img-16bit.png") 

374 ) 

375 # the way to load a 16-bit PNG image correctly, according to: 

376 # https://stackoverflow.com/questions/32622658/read-16-bit-png-image-file-using-python 

377 # https://github.com/python-pillow/Pillow/issues/3011 

378 img = PIL.Image.fromarray( 

379 numpy.array(PIL.Image.open(path)).astype("uint16") 

380 ) 

381 assert img.mode == "I;16" 

382 assert img.getextrema() == (0, 65281) 

383 

384 timg = SingleAutoLevel16to8()(img) 

385 assert timg.mode == "L" 

386 assert timg.getextrema() == (0, 255) 

387 # timg.show() 

388 # import ipdb; ipdb.set_trace() 

389 

390 

391def test_ResizeCrop(): 

392 

393 # parameters 

394 im_size = (3, 128, 140) # (planes, height, width) 

395 mask_gt_size = (1, 128, 140) # (planes, height, width) 

396 crop_size = (30, 30, 91, 91) # (left,up , right ,down) 

397 size_after_crop = (61, 61) 

398 

399 idx = (slice(crop_size[0], crop_size[2]), slice(crop_size[1], crop_size[3])) 

400 

401 # Create random image and a mask with a circle inside 

402 img, gt = _create_img(im_size), _create_img(mask_gt_size) 

403 mask = PIL.Image.new("L", (140, 128), "black") 

404 dr = PIL.ImageDraw.Draw(mask) 

405 dr.ellipse((30, 30, 90, 90), "white") 

406 

407 # Test 

408 transform = ResizeCrop() 

409 img_, gt_, mask_ = transform(img, gt, mask) 

410 

411 assert img_.size == size_after_crop 

412 assert gt_.size == size_after_crop 

413 assert mask_.size == size_after_crop 

414 

415 assert img_.mode == "RGB" 

416 assert gt_.mode == "L" 

417 assert mask_.mode == "L" 

418 

419 assert numpy.all(numpy.array(img_) == numpy.array(img)[idx]) 

420 assert numpy.all(numpy.array(gt_) == numpy.array(gt)[idx]) 

421 assert numpy.all(numpy.array(mask_) == numpy.array(mask)[idx])