| 511 | |
| 512 | |
| 513 | class ToTensor(object): |
| 514 | def __init__(self, mode, do_normalize=False, size=None): |
| 515 | self.mode = mode |
| 516 | self.normalize = transforms.Normalize( |
| 517 | mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) if do_normalize else nn.Identity() |
| 518 | self.size = size |
| 519 | if size is not None: |
| 520 | self.resize = transforms.Resize(size=size) |
| 521 | else: |
| 522 | self.resize = nn.Identity() |
| 523 | |
| 524 | def __call__(self, sample): |
| 525 | image, focal = sample['image'], sample['focal'] |
| 526 | image = self.to_tensor(image) |
| 527 | image = self.normalize(image) |
| 528 | image = self.resize(image) |
| 529 | |
| 530 | if self.mode == 'test': |
| 531 | return {'image': image, 'focal': focal} |
| 532 | |
| 533 | depth = sample['depth'] |
| 534 | if self.mode == 'train': |
| 535 | depth = self.to_tensor(depth) |
| 536 | return {**sample, 'image': image, 'depth': depth, 'focal': focal} |
| 537 | else: |
| 538 | has_valid_depth = sample['has_valid_depth'] |
| 539 | image = self.resize(image) |
| 540 | return {**sample, 'image': image, 'depth': depth, 'focal': focal, 'has_valid_depth': has_valid_depth, |
| 541 | 'image_path': sample['image_path'], 'depth_path': sample['depth_path']} |
| 542 | |
| 543 | def to_tensor(self, pic): |
| 544 | if not (_is_pil_image(pic) or _is_numpy_image(pic)): |
| 545 | raise TypeError( |
| 546 | 'pic should be PIL Image or ndarray. Got {}'.format(type(pic))) |
| 547 | |
| 548 | if isinstance(pic, np.ndarray): |
| 549 | img = torch.from_numpy(pic.transpose((2, 0, 1))) |
| 550 | return img |
| 551 | |
| 552 | # handle PIL Image |
| 553 | if pic.mode == 'I': |
| 554 | img = torch.from_numpy(np.array(pic, np.int32, copy=False)) |
| 555 | elif pic.mode == 'I;16': |
| 556 | img = torch.from_numpy(np.array(pic, np.int16, copy=False)) |
| 557 | else: |
| 558 | img = torch.ByteTensor( |
| 559 | torch.ByteStorage.from_buffer(pic.tobytes())) |
| 560 | # PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK |
| 561 | if pic.mode == 'YCbCr': |
| 562 | nchannel = 3 |
| 563 | elif pic.mode == 'I;16': |
| 564 | nchannel = 1 |
| 565 | else: |
| 566 | nchannel = len(pic.mode) |
| 567 | img = img.view(pic.size[1], pic.size[0], nchannel) |
| 568 | |
| 569 | img = img.transpose(0, 1).transpose(0, 2).contiguous() |
| 570 | if isinstance(img, torch.ByteTensor): |
no outgoing calls
no test coverage detected