| 256 | |
| 257 | |
| 258 | class TestDataset(BaseDataset): |
| 259 | def __init__(self, odgt, opt, **kwargs): |
| 260 | super(TestDataset, self).__init__(odgt, opt, **kwargs) |
| 261 | |
| 262 | def __getitem__(self, index): |
| 263 | this_record = self.list_sample[index] |
| 264 | # load image |
| 265 | image_path = this_record['fpath_img'] |
| 266 | img = Image.open(image_path).convert('RGB') |
| 267 | |
| 268 | ori_width, ori_height = img.size |
| 269 | |
| 270 | img_resized_list = [] |
| 271 | for this_short_size in self.imgSizes: |
| 272 | # calculate target height and width |
| 273 | scale = min(this_short_size / float(min(ori_height, ori_width)), |
| 274 | self.imgMaxSize / float(max(ori_height, ori_width))) |
| 275 | target_height, target_width = int(ori_height * scale), int(ori_width * scale) |
| 276 | |
| 277 | # to avoid rounding in network |
| 278 | target_width = self.round2nearest_multiple(target_width, self.padding_constant) |
| 279 | target_height = self.round2nearest_multiple(target_height, self.padding_constant) |
| 280 | |
| 281 | # resize images |
| 282 | img_resized = imresize(img, (target_width, target_height), interp='bilinear') |
| 283 | |
| 284 | # image transform, to torch float tensor 3xHxW |
| 285 | img_resized = self.img_transform(img_resized) |
| 286 | img_resized = torch.unsqueeze(img_resized, 0) |
| 287 | img_resized_list.append(img_resized) |
| 288 | |
| 289 | output = dict() |
| 290 | output['img_ori'] = np.array(img) |
| 291 | output['img_data'] = [x.contiguous() for x in img_resized_list] |
| 292 | output['info'] = this_record['fpath_img'] |
| 293 | return output |
| 294 | |
| 295 | def __len__(self): |
| 296 | return self.num_sample |