| 14 | |
| 15 | |
| 16 | class REFUGE(Dataset): |
| 17 | def __init__(self, args, data_path , transform = None, transform_msk = None, mode = 'Training',prompt = 'click', plane = False): |
| 18 | self.data_path = data_path |
| 19 | self.subfolders = [f.path for f in os.scandir(os.path.join(data_path, mode + '-400')) if f.is_dir()] |
| 20 | self.mode = mode |
| 21 | self.prompt = prompt |
| 22 | self.img_size = args.image_size |
| 23 | self.mask_size = args.out_size |
| 24 | |
| 25 | self.transform = transform |
| 26 | self.transform_msk = transform_msk |
| 27 | |
| 28 | def __len__(self): |
| 29 | return len(self.subfolders) |
| 30 | |
| 31 | def __getitem__(self, index): |
| 32 | |
| 33 | """Get the images""" |
| 34 | subfolder = self.subfolders[index] |
| 35 | name = subfolder.split('/')[-1] |
| 36 | |
| 37 | # raw image and raters path |
| 38 | img_path = os.path.join(subfolder, name + '_cropped.jpg') |
| 39 | multi_rater_cup_path = [os.path.join(subfolder, name + '_seg_cup_' + str(i) + '_cropped.jpg') for i in range(1, 8)] |
| 40 | |
| 41 | # img_path = os.path.join(subfolder, name + '.jpg') |
| 42 | # multi_rater_cup_path = [os.path.join(subfolder, name + '_seg_cup_' + str(i) + '.png') for i in range(1, 8)] |
| 43 | |
| 44 | # raw image and rater images |
| 45 | img = Image.open(img_path).convert('RGB') |
| 46 | multi_rater_cup = [Image.open(path).convert('L') for path in multi_rater_cup_path] |
| 47 | |
| 48 | # apply transform |
| 49 | if self.transform: |
| 50 | state = torch.get_rng_state() |
| 51 | img = self.transform(img) |
| 52 | multi_rater_cup = [torch.as_tensor((self.transform(single_rater) >=0.5).float(), dtype=torch.float32) for single_rater in multi_rater_cup] |
| 53 | multi_rater_cup = torch.stack(multi_rater_cup, dim=0) |
| 54 | |
| 55 | torch.set_rng_state(state) |
| 56 | |
| 57 | # find init click and apply majority vote |
| 58 | if self.prompt == 'click': |
| 59 | |
| 60 | point_label_cup, pt_cup = random_click(np.array((multi_rater_cup.mean(axis=0)).squeeze(0)), point_label = 1) |
| 61 | |
| 62 | selected_rater_mask_cup_ori = multi_rater_cup.mean(axis=0) |
| 63 | selected_rater_mask_cup_ori = (selected_rater_mask_cup_ori >= 0.5).float() |
| 64 | |
| 65 | |
| 66 | selected_rater_mask_cup = F.interpolate(selected_rater_mask_cup_ori.unsqueeze(0), size=(self.mask_size, self.mask_size), mode='bilinear', align_corners=False).mean(dim=0) # torch.Size([1, mask_size, mask_size]) |
| 67 | selected_rater_mask_cup = (selected_rater_mask_cup >= 0.5).float() |
| 68 | |
| 69 | |
| 70 | # # Or use any specific rater as GT |
| 71 | # point_label_cup, pt_cup = random_click(np.array(multi_rater_cup[0, :, :, :].squeeze(0)), point_label = 1) |
| 72 | # selected_rater_mask_cup_ori = multi_rater_cup[0,:,:,:] |
| 73 | # selected_rater_mask_cup_ori = (selected_rater_mask_cup_ori >= 0.5).float() |