| 33 | |
| 34 | class iBims(Dataset): |
| 35 | def __init__(self, config): |
| 36 | root_folder = config.ibims_root |
| 37 | with open(os.path.join(root_folder, "imagelist.txt"), 'r') as f: |
| 38 | imglist = f.read().split() |
| 39 | |
| 40 | samples = [] |
| 41 | for basename in imglist: |
| 42 | img_path = os.path.join(root_folder, 'rgb', basename + ".png") |
| 43 | depth_path = os.path.join(root_folder, 'depth', basename + ".png") |
| 44 | valid_mask_path = os.path.join( |
| 45 | root_folder, 'mask_invalid', basename+".png") |
| 46 | transp_mask_path = os.path.join( |
| 47 | root_folder, 'mask_transp', basename+".png") |
| 48 | |
| 49 | samples.append( |
| 50 | (img_path, depth_path, valid_mask_path, transp_mask_path)) |
| 51 | |
| 52 | self.samples = samples |
| 53 | # self.normalize = T.Normalize( |
| 54 | # mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) |
| 55 | self.normalize = lambda x : x |
| 56 | |
| 57 | def __getitem__(self, idx): |
| 58 | img_path, depth_path, valid_mask_path, transp_mask_path = self.samples[idx] |