| 32 | |
| 33 | |
| 34 | class ToTensor(object): |
| 35 | def __init__(self, resize_shape): |
| 36 | # self.normalize = transforms.Normalize( |
| 37 | # mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) |
| 38 | self.normalize = lambda x : x |
| 39 | self.resize = transforms.Resize(resize_shape) |
| 40 | |
| 41 | def __call__(self, sample): |
| 42 | image, depth = sample['image'], sample['depth'] |
| 43 | image = self.to_tensor(image) |
| 44 | image = self.normalize(image) |
| 45 | depth = self.to_tensor(depth) |
| 46 | |
| 47 | image = self.resize(image) |
| 48 | |
| 49 | return {'image': image, 'depth': depth, 'dataset': "ddad"} |
| 50 | |
| 51 | def to_tensor(self, pic): |
| 52 | |
| 53 | if isinstance(pic, np.ndarray): |
| 54 | img = torch.from_numpy(pic.transpose((2, 0, 1))) |
| 55 | return img |
| 56 | |
| 57 | # # handle PIL Image |
| 58 | if pic.mode == 'I': |
| 59 | img = torch.from_numpy(np.array(pic, np.int32, copy=False)) |
| 60 | elif pic.mode == 'I;16': |
| 61 | img = torch.from_numpy(np.array(pic, np.int16, copy=False)) |
| 62 | else: |
| 63 | img = torch.ByteTensor( |
| 64 | torch.ByteStorage.from_buffer(pic.tobytes())) |
| 65 | # PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK |
| 66 | if pic.mode == 'YCbCr': |
| 67 | nchannel = 3 |
| 68 | elif pic.mode == 'I;16': |
| 69 | nchannel = 1 |
| 70 | else: |
| 71 | nchannel = len(pic.mode) |
| 72 | img = img.view(pic.size[1], pic.size[0], nchannel) |
| 73 | |
| 74 | img = img.transpose(0, 1).transpose(0, 2).contiguous() |
| 75 | |
| 76 | if isinstance(img, torch.ByteTensor): |
| 77 | return img.float() |
| 78 | else: |
| 79 | return img |
| 80 | |
| 81 | |
| 82 | class DDAD(Dataset): |
no outgoing calls
no test coverage detected