| 49 | return {'image': image, 'depth': depth, 'dataset': "diode"} |
| 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 DIODE(Dataset): |