Spaces:
Configuration error
Configuration error
| import json | |
| import os | |
| import pickle | |
| import nibabel as nib | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| import torch.nn.functional as F | |
| from PIL import Image | |
| from torch.utils.data import Dataset | |
| from utils import generate_click_prompt, random_box, random_click | |
| class Atlas(Dataset): | |
| def __init__(self, args, data_path , transform = None, transform_msk = None, mode = 'Training',prompt = 'click', plane = False): | |
| self.args = args | |
| self.data_path = os.path.join(data_path,'train') | |
| with open(os.path.join(self.data_path,'dataset.json'),'r') as file: | |
| data = json.load(file) | |
| self.name_list = data['training'] | |
| self.mode = mode | |
| self.prompt = prompt | |
| self.img_size = args.image_size | |
| self.transform = transform | |
| self.transform_msk = transform_msk | |
| def __len__(self): | |
| return len(self.name_list) | |
| def __getitem__(self, index): | |
| point_label = 1 | |
| label = 1 | |
| """Get the images""" | |
| img_name = self.name_list[index]['image'] | |
| mask_name = self.name_list[index]['label'] | |
| img = nib.load(os.path.join(self.data_path,img_name)).get_fdata() | |
| mask = nib.load(os.path.join(self.data_path,mask_name)).get_fdata() | |
| mask[mask!=label] = 0 | |
| mask[mask==label] = 1 | |
| # if self.mode == 'Training': | |
| # label = 0 if self.label_list[index] == 'benign' else 1 | |
| # else: | |
| # label = int(self.label_list[index]) | |
| img = np.transpose(img,(1,2,0)) | |
| mask = np.transpose(mask,(1,2,0)) | |
| # img = np.resize(mask,(self.args.image_size, self.args.image_size,128)) | |
| # mask = np.resize(mask,(self.args.out_size,self.args.out_size,128)) | |
| # # img = np.resize(img,(self.args.image_size, self.args.image_size,img.shape[-1])) | |
| # # mask = np.resize(mask,(self.args.out_size,self.args.out_size,mask.shape[-1])) | |
| img = torch.tensor(img).unsqueeze(0) | |
| mask = torch.tensor(mask).unsqueeze(0) | |
| if self.prompt == 'click': | |
| point_label, pt = random_click(np.array(mask), point_label) | |
| # if self.transform: | |
| # state = torch.get_rng_state() | |
| # img = self.transform(img) | |
| # torch.set_rng_state(state) | |
| # if self.transform_msk: | |
| # mask = self.transform_msk(mask) | |
| # # if (inout == 0 and point_label == 1) or (inout == 1 and point_label == 0): | |
| # # mask = 1 - mask | |
| name = img_name | |
| image_meta_dict = {'filename_or_obj':name} | |
| return { | |
| 'image':img, | |
| 'label': mask, | |
| 'p_label':point_label, | |
| 'pt':pt, | |
| 'image_meta_dict':image_meta_dict, | |
| } | |