| 12 | |
| 13 | |
| 14 | class LIDCMultiInputMultiResTypes(Dataset): |
| 15 | def __init__(self, paths, opt: dict, size=None, labels=None, num_ctslice_per_item=1): |
| 16 | assert num_ctslice_per_item in [1, 3] |
| 17 | |
| 18 | self.opt = opt |
| 19 | self.ct_size = opt['ct_size'] # ct resolution, number of CT slice |
| 20 | self.xray_size = opt['xray_size'] # xray_resolution |
| 21 | self.input_types = opt['input_type'] |
| 22 | |
| 23 | self.labels = dict() if labels is None else labels |
| 24 | self.labels["file_path_"] = paths |
| 25 | self._length = len(paths) |
| 26 | self.num_ctslice_per_item = num_ctslice_per_item |
| 27 | |
| 28 | self.CT_MIN_MAX = opt["CT_MIN_MAX"] |
| 29 | self.XRAY_MIN_MAX = opt["XRAY_MIN_MAX"] |
| 30 | |
| 31 | self.set_preprocessing() |
| 32 | self.mapping_camera_type2pose = { |
| 33 | "PA": torch.tensor([0, 0]), |
| 34 | "Lateral": torch.tensor([math.pi / 2, math.pi / 2]), |
| 35 | } |
| 36 | |
| 37 | def __len__(self): |
| 38 | return self._length |
| 39 | |
| 40 | def set_preprocessing(self): |
| 41 | dict_augment_list = {} |
| 42 | for input_type in self.input_types: |
| 43 | i_type = 'ct' if input_type in ['ct', 'ctslice'] else 'xray' |
| 44 | dict_augment_list[input_type] = self.opt[f"{i_type}_augment_list"] |
| 45 | |
| 46 | self.dict_preprocessing = {} |
| 47 | for input_type in self.input_types: |
| 48 | augment_list = [] |
| 49 | if input_type in ['ct', 'ctslice']: |
| 50 | if 'min_max_th' in dict_augment_list[input_type]: |
| 51 | augment_list.append((Limit_Min_Max_Threshold(self.CT_MIN_MAX[0], self.CT_MIN_MAX[1]),)) |
| 52 | if 'normalization' in dict_augment_list[input_type]: |
| 53 | augment_list.append((Normalization(self.CT_MIN_MAX[0], self.CT_MIN_MAX[1]),)) |
| 54 | elif input_type in ['PA', 'Lateral']: |
| 55 | if 'normalization' in dict_augment_list[input_type]: |
| 56 | augment_list.append((Normalization(self.XRAY_MIN_MAX[0], self.XRAY_MIN_MAX[1]),)) |
| 57 | augment_list.append((ToTensor(),)) |
| 58 | self.dict_preprocessing[input_type] = List_Compose(augment_list) |
| 59 | |
| 60 | def get_image(self, image_path, data_type='ct'): |
| 61 | ext = image_path.split(".")[-1] |
| 62 | assert ext in ['png', 'h5'] |
| 63 | |
| 64 | if ext in ['png']: |
| 65 | image = imageio.imread(image_path) |
| 66 | image = np.asarray(image) |
| 67 | elif ext in ["h5"]: |
| 68 | with h5py.File(image_path, 'r') as f: |
| 69 | image = np.asarray(f[data_type]) # 128 x 128 |
| 70 | |
| 71 | return image |
nothing calls this directly
no outgoing calls
no test coverage detected