dataloader.py

import os
import csv
import torch
import random
import numpy as np
import pandas as pd
from collections import Counter
from torch.utils.data import Dataset
from sklearn.model_selection import StratifiedKFold

try:
    import h5py
except:
    print('no h5py')
    
os.environ['HDF5_USE_FILE_LOCKING'] = 'FALSE'

def readCSV(filename):
    lines = []
    with open(filename, "r") as f:
        csvreader = csv.reader(f)
        for line in csvreader:
            lines.append(line)
    return lines

def get_patient_label(csv_file):
    patients_list=[]
    labels_list=[]
    label_file = readCSV(csv_file)
    for i in range(0, len(label_file)):
        patients_list.append(label_file[i][0])
        labels_list.append(label_file[i][1])
    a=Counter(labels_list)
    print("patient_len:{} label_len:{}".format(len(patients_list), len(labels_list)))
    print("all_counter:{}".format(dict(a)))
    return np.array(patients_list,dtype=object), np.array(labels_list,dtype=object)

def data_split(full_list, ratio, shuffle=True,label=None,label_balance_val=True):
    """
    dataset split: split the full_list randomly into two sublist (val-set and train-set) based on the ratio
    :param full_list: 
    :param ratio:     
    :param shuffle:  
    """
    # select the val-set based on the label ratio
    if label_balance_val and label is not None:
        _label = label[full_list]
        _label_uni = np.unique(_label)
        sublist_1 = []
        sublist_2 = []

        for _l in _label_uni:
            _list = full_list[_label == _l]
            n_total = len(_list)
            offset = int(n_total * ratio)
            if shuffle:
                random.shuffle(_list)
            sublist_1.extend(_list[:offset])
            sublist_2.extend(_list[offset:])
    else:
        n_total = len(full_list)
        offset = int(n_total * ratio)
        if n_total == 0 or offset < 1:
            return [], full_list
        if shuffle:
            random.shuffle(full_list)
        val_set = full_list[:offset]
        train_set = full_list[offset:]

    return val_set, train_set


def get_kflod(k, patients_array, labels_array,val_ratio=False,label_balance_val=True):
    if k > 1:
        skf = StratifiedKFold(n_splits=k)
    else:
        raise NotImplementedError
    train_patients_list = []
    train_labels_list = []
    test_patients_list = []
    test_labels_list = []
    val_patients_list = []
    val_labels_list = []
    for train_index, test_index in skf.split(patients_array, labels_array):
        if val_ratio != 0.:
            val_index,train_index = data_split(train_index,val_ratio,True,labels_array,label_balance_val)
            x_val, y_val = patients_array[val_index], labels_array[val_index]
        else:
            x_val, y_val = [],[]
        x_train, x_test = patients_array[train_index], patients_array[test_index]
        y_train, y_test = labels_array[train_index], labels_array[test_index]

        train_patients_list.append(x_train)
        train_labels_list.append(y_train)
        test_patients_list.append(x_test)
        test_labels_list.append(y_test)
        val_patients_list.append(x_val)
        val_labels_list.append(y_val)
        
    # print("get_kflod.type:{}".format(type(np.array(train_patients_list))))
    return np.array(train_patients_list,dtype=object), np.array(train_labels_list,dtype=object), np.array(test_patients_list,dtype=object), np.array(test_labels_list,dtype=object),np.array(val_patients_list,dtype=object), np.array(val_labels_list,dtype=object)

def get_tcga_parser(root,cls_name,mini=False):
        x = []
        y = []

        for idx,_cls in enumerate(cls_name):
            _dir = 'mini_pt' if mini else 'pt_files'
            _files = os.listdir(os.path.join(root,_cls,'features',_dir))
            _files = [os.path.join(os.path.join(root,_cls,'features',_dir,_files[i])) for i in range(len(_files))]
            x.extend(_files)
            y.extend([idx for i in range(len(_files))])
            
        return np.array(x).flatten(),np.array(y).flatten()

class TCGADataset(Dataset):
    
    def __init__(self, file_name=None, file_label=None,max_patch=-1,root=None,persistence=True,keep_same_psize=0,is_train=False,_type='nsclc'):
        """
        Args
        :param images: 
        :param transform: optional transform to be applied on a sample
        """
        super(TCGADataset, self).__init__()

        self.patient_name = file_name
        self.patient_label = file_label
        self.max_patch = max_patch
        self.root = root
        self.all_pts = os.listdir(os.path.join(self.root,'h5_files')) if keep_same_psize else os.listdir(os.path.join(self.root,'pt_files'))
        self.slide_name = []
        self.slide_label = []
        self.persistence = persistence
        self.keep_same_psize = keep_same_psize
        self.is_train = is_train

        for i,_patient_name in enumerate(self.patient_name):
            _sides = np.array([ _slide if _patient_name in _slide else '0' for _slide in self.all_pts])
            _ids = np.where(_sides != '0')[0]
            for _idx in _ids:
                if persistence:
                    self.slide_name.append(torch.load(os.path.join(self.root,'pt_files',_sides[_idx])))
                else:
                    self.slide_name.append(_sides[_idx])
                self.slide_label.append(self.patient_label[i])
                
        if _type.lower() == 'nsclc':
            self.slide_label = [ 0 if _l == 'LUAD' else 1 for _l in self.slide_label]
        elif _type.lower() == 'brca':
            self.slide_label = [ 0 if _l == 'IDC' else 1 for _l in self.slide_label]

    def __len__(self):
        return len(self.slide_name)

    def __getitem__(self, idx):
        """
        Args
        :param idx: the index of item
        :return: image and its label
        """
        file_path = self.slide_name[idx]
        label = self.slide_label[idx]

        if self.persistence:
            features = file_path
        else:
            features = torch.load(os.path.join(self.root,'pt_files',file_path))
        return features , int(label)

class C16Dataset(Dataset):

    def __init__(self, file_name, file_label,root,persistence=False,keep_same_psize=0,is_train=False):
        """
        Args
        :param images: 
        :param transform: optional transform to be applied on a sample
        """
        super(C16Dataset, self).__init__()
        self.file_name = file_name
        self.slide_label = file_label
        self.slide_label = [int(_l) for _l in self.slide_label]
        self.size = len(self.file_name)
        self.root = root
        self.persistence = persistence
        self.keep_same_psize = keep_same_psize
        self.is_train = is_train

        if persistence:
            self.feats = [ torch.load(os.path.join(root,'pt', _f+'.pt')) for _f in file_name ]

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        """
        Args
        :param idx: the index of item
        :return: image and its label
        """
        if self.persistence:
            features = self.feats[idx]
        else:
            dir_path = os.path.join(self.root,"pt")

            file_path = os.path.join(dir_path, self.file_name[idx]+'.pt')
            features = torch.load(file_path)

        label = int(self.slide_label[idx])


        return features , label

class TCGADataset_SAM(Dataset):

    def __init__(self, file_name=None, file_label=None, max_patch=-1, root=None, persistence=True, keep_same_psize=0,
                 is_train=False, _type='NSCLC'):
        """
        Args
        :param images:
        :param transform: optional transform to be applied on a sample
        """
        super(TCGADataset_SAM, self).__init__()

        self.patient_name = file_name
        self.patient_label = file_label
        self.max_patch = max_patch
        self.root = root
        self.all_pts = os.listdir(os.path.join(self.root, 'h5_files')) if keep_same_psize else os.listdir(
            os.path.join(self.root, 'pt_files'))
        self.slide_name = []
        self.slide_label = []
        self.is_group_feat = []
        self.relative_area = []

        self.persistence = persistence
        self.keep_same_psize = keep_same_psize
        self.is_train = is_train

        for i, _patient_name in enumerate(self.patient_name):
            _sides = np.array([_slide if _patient_name in _slide else '0' for _slide in self.all_pts])
            _ids = np.where(_sides != '0')[0]
            for _idx in _ids:
                if persistence:
                    self.slide_name.append(torch.load(os.path.join(self.root, 'pt_files', _sides[_idx])))
                    self.is_group_feat.append(np.array(h5py.File(os.path.join(self.root, 'sam_info', os.path.splitext(_sides[_idx])[0] + '.h5'), 'r')['is_group_feat']))
                    self.relative_area.append(np.array(h5py.File(os.path.join(self.root, 'sam_info', os.path.splitext(_sides[_idx])[0] + '.h5'), 'r')['relative_area']))

                else:
                    self.slide_name.append(_sides[_idx])
                self.slide_label.append(self.patient_label[i])
        if _type.lower() == 'nsclc':
            self.slide_label = [0 if _l == 'LUAD' else 1 for _l in self.slide_label]
        elif _type.lower() == 'brca':
            self.slide_label = [0 if _l == 'IDC' else 1 for _l in self.slide_label]

    def __len__(self):
        return len(self.slide_name)

    def __getitem__(self, idx):
        """
        Args
        :param idx: the index of item
        :return: image and its label
        """
        file_path = self.slide_name[idx]
        # file_path = self.csv_file[idx] # ['1_1.png']
        # patient_path = file_path[1]
        label = self.slide_label[idx]

        if self.persistence:
            features = file_path
            is_group_feat = self.is_group_feat[idx]
            relative_area = self.relative_area[idx]
        else:
            features = torch.load(os.path.join(self.root, 'pt_files', file_path))

            with h5py.File(os.path.join(self.root, 'sam_info', os.path.splitext(self.slide_name[idx])[0] + '.h5'), 'r') as h5_file:
                is_group_feat = np.array(h5_file['is_group_feat'])
                relative_area = np.array(h5_file['relative_area'])

        return features, int(label), is_group_feat, relative_area



class C16Dataset_SAM(Dataset):
    def __init__(self, file_name, file_label,root,persistence=False,keep_same_psize=0,is_train=False,is_all=False):
        """
        Args
        :param images:
        :param transform: optional transform to be applied on a sample
        """
        super(C16Dataset_SAM, self).__init__()
        # self.csv_file = readCSV(csv_file)
        self.file_name = file_name
        self.slide_label = file_label
        self.slide_label = [int(_l) for _l in self.slide_label]
        self.size = len(self.file_name)
        self.root = root
        self.persistence = persistence
        self.keep_same_psize = keep_same_psize
        self.is_train = is_train
        self.is_all = is_all

        if persistence:
            # TODO: Not support to C16+C17
                self.feats = [ torch.load(os.path.join(root,'pt', _f+'.pt')) for _f in file_name ]
                self.is_group_feat = [ np.array(h5py.File(os.path.join(root, 'sam', _f + '.h5'), 'r')['is_group_feat']) for _f in file_name]
                self.relative_area = [ np.array(h5py.File(os.path.join(root, 'sam', _f + '.h5'), 'r')['relative_area']) for _f in file_name]

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        """
        Args
        :param idx: the index of item
        :return: image and its label
        """
        if self.persistence:
            features = self.feats[idx]
            is_group_feat = self.is_group_feat[idx]
            relative_area = self.relative_area[idx]
        else:
            dir_path = os.path.join(self.root,"pt")
            sam_path = os.path.join(self.root, 'sam')

            # TODO: Not support to C16+C17
            if self.is_all:
                file_path = os.path.join(dir_path, 'c16',self.file_name[idx]+'.pt') if self.file_name[idx].split('_')[0] != 'patient' else os.path.join(dir_path, 'c17',self.file_name[idx]+'.pt')
            else:
                file_path = os.path.join(dir_path, self.file_name[idx]+'.pt')
                sam_info_path = os.path.join(sam_path, self.file_name[idx]+'.h5')
            # file_path = self.csv_file[idx] # ['1_1.png']
            # patient_path = file_path[1]
            features = torch.load(file_path)

            with h5py.File(os.path.join(self.root, 'sam', self.file_name[idx] + '.h5'), 'r') as h5_file:
                is_group_feat = np.array(h5_file['is_group_feat'])
                relative_area = np.array(h5_file['relative_area'])

        label = int(self.slide_label[idx])
        # num = features.shape[0]
        # if num > 128:
        #     sample = np.random.choice(num, 128, replace=False)
        # else:
        #     sample = np.random.choice(num, num, replace=False)
        # sample = np.sort(sample)
        # new_feat = features[sample]
        # return new_feat, label

        return features, label, is_group_feat, relative_area