WsiDataloader.py

import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
import random
import pandas as pd
import utils
from WsiDataset import DetWSIDataset, SegWSIDataset, SubWSIDataset, PatchDataset, TrainWSIDataset, AnnoPatchDataset
from WsiDataset import divide_data, divide_multiclass_data
import os



##### Dataset specific dataloader #####

class dataloader_generator(): #val_size not implemented
    def __init__(self, all_dir_path, train_dir_path, division_json, batch_size=256, shots=[1,5,10], dataset='BRCA'):
        all_h5_paths = utils.get_h5_paths(all_dir_path)
        self.train_h5_paths = utils.get_h5_paths(train_dir_path)
        self.division_json = division_json
        self.dataset = dataset

        self.batch_size = batch_size
        self.shots = shots
        
        all_h5_paths = all_h5_paths[:10]
        
        self.valset = SubWSIDataset(h5_path_lst=all_h5_paths, division_json=division_json, dataset=self.dataset)
        self.testset = SubWSIDataset(h5_path_lst=all_h5_paths, division_json=division_json, dataset=self.dataset)
        
        self.fewshot_paths = []
        
    def get_classnames(self):
        return self.testset.get_classnames()

    def gen_testloader(self):
        test_loader = DataLoader(self.testset, batch_size=1, shuffle=False) # wsidataset batch size is 1
        return test_loader
    
    def gen_valloader(self):
        val_loader = DataLoader(self.valset, batch_size=1, shuffle=False)
        return val_loader
    
    def gen_trainloaders(self):# 按shotsl倒序返回
        loaders = []
        shots_new = self.shots.copy()
        shots_new.sort(reverse=True)
        train_h5_paths = self.train_h5_paths.copy()
        while shots_new:
            shot = shots_new.pop(0)
            trainset = AnnoPatchDataset(h5_path_lst=train_h5_paths, division_json=self.division_json, shots=shot, dataset=self.dataset)
            trainloader = DataLoader(trainset, batch_size=self.batch_size, shuffle=True) #这里shuffle true就不用在dataset里shuffle了
            loaders.append(trainloader) 
            train_h5_paths = trainset.get_example_paths().copy()

        return loaders


class CM_det_dataloader_generator():#不需要固定test_ids，因为train和test file在不同的dir
    def __init__(self, train_dir_path, test_dir_path, train_label_csv, test_label_csv, batch_size=1024, shots=[1,2,5], val_size=32):
        self.train_h5_path_lst = utils.get_h5_paths(train_dir_path)
        self.test_h5_path_lst = utils.get_h5_paths(test_dir_path)

        self.train_label_csv = train_label_csv
        self.test_label_csv = test_label_csv
        self.batch_size = batch_size
        self.shots = shots

        self.valset = DetWSIDataset(h5_path_lst=self.test_h5_path_lst[:val_size], label_csv=self.test_label_csv)
        
        self.testset = DetWSIDataset(h5_path_lst=self.test_h5_path_lst, label_csv=self.test_label_csv)
    
    def gen_testloader(self):
        test_loader = DataLoader(self.testset, batch_size=1, shuffle=False) # wsidataset batch size is 1
        return test_loader
    
    def gen_valloader(self):
        val_loader = DataLoader(self.valset, batch_size=1, shuffle=False)
        return val_loader
    
    def gen_trainloaders(self):# 按shots倒序返回
        loaders = []
        shots_new = self.shots.copy()
        shots_new.sort(reverse=True)
        h5_path_lst = self.train_h5_path_lst.copy()
        max_len = 0
        while shots_new:
            shot = shots_new.pop(0)
            if max_len == 0:
                trainset = TrainWSIDataset(h5_path_lst=h5_path_lst, label_csv=self.train_label_csv, shots=shot, task='det')
                max_len = len(trainset)
            else:
                trainset = TrainWSIDataset(h5_path_lst=h5_path_lst, label_csv=self.train_label_csv, shots=shot, task='det', given_len=max_len)
            trainloader = DataLoader(trainset, batch_size=self.batch_size, shuffle=True)
            loaders.append(trainloader) 
            
            h5_paths = trainset.get_example_paths()
            h5_path_lst = sum([class_path_lst for class_path_lst in h5_paths], start=[])
        return loaders


class camelyon_det_dataloader_generator():
    def __init__(self, normal_dir_path, tumor_dir_path, batch_size=1024, shots=[1,2,5]):
        normal_h5_path_lst = utils.get_h5_paths(normal_dir_path)
        tumor_h5_path_lst = utils.get_h5_paths(tumor_dir_path)
        # Shuffle the lists to ensure randomness
        random.shuffle(normal_h5_path_lst)
        random.shuffle(tumor_h5_path_lst)
        half_normal_len, half_tumor_len = len(normal_h5_path_lst) // 2, len(tumor_h5_path_lst) // 2
        self.train_normal_h5_path_lst = normal_h5_path_lst[:half_normal_len]
        self.test_normal_h5_path_lst = normal_h5_path_lst[half_normal_len:]
        self.train_tumor_h5_path_lst = tumor_h5_path_lst[:half_tumor_len]
        self.test_tumor_h5_path_lst = tumor_h5_path_lst[half_tumor_len:]
        self.shots = shots
        self.batch_size = batch_size
        
        self.testset = DetWSIDataset(labeled_h5_paths=[self.test_normal_h5_path_lst, self.test_tumor_h5_path_lst])

    def gen_testloader(self):
        testloader = DataLoader(self.testset, batch_size=1, shuffle=False) # wsidataset batch size is 1
        return testloader

    def gen_trainloaders(self): # 按shots倒序返回
        loaders = []
        shots_new = self.shots.copy()
        shots_new.sort(reverse=True)
        train_normal_h5_path_lst = self.train_normal_h5_path_lst.copy()
        train_tumor_h5_path_lst = self.train_tumor_h5_path_lst.copy()

        while shots_new:
            shot = shots_new.pop(0)
            train_normal_h5_path_lst = random.sample(train_normal_h5_path_lst, shot)
            train_tumor_h5_path_lst = random.sample(train_tumor_h5_path_lst, shot)
            trainset = PatchDataset(labeled_h5_paths=[train_normal_h5_path_lst, train_tumor_h5_path_lst], shots=shot, task='det')
            trainloader = DataLoader(trainset, batch_size=self.batch_size, shuffle=True)
            loaders.append(trainloader)

        return loaders


class camelyon_seg_dataloader_generator():
    def __init__(self, normal_dir_path, tumor_dir_path, test_mask_dir, test_patch_mask_dir, batch_size=1024, shots=[1,2,5], val_size=16, testset_ids=None, device='cpu'):
        self.train_normal_h5_path_lst = utils.get_h5_paths(normal_dir_path) #segmentation only use tumor wsi to test
        tumor_h5_path_lst = utils.get_h5_paths(tumor_dir_path)

        if len(testset_ids) == 0:
            random.shuffle(tumor_h5_path_lst)
            half_tumor_len = len(tumor_h5_path_lst) // 2
            self.train_tumor_h5_path_lst = tumor_h5_path_lst[:half_tumor_len]
            self.test_tumor_h5_path_lst = tumor_h5_path_lst[half_tumor_len:]
        else:
            self.test_tumor_h5_path_lst = []
            self.train_tumor_h5_path_lst = []
            for wsi_path in tumor_h5_path_lst:
                wsi_id = os.path.basename(wsi_path)[:-3]
                if wsi_id in testset_ids:
                    self.test_tumor_h5_path_lst.append(wsi_path)
                else:
                    self.train_tumor_h5_path_lst.append(wsi_path)

        self.val_tumor_h5_path_lst = self.test_tumor_h5_path_lst[:val_size]
        self.test_mask_dir = test_mask_dir
        self.test_patch_mask_dir = test_patch_mask_dir
        self.val_mask_dir = test_mask_dir
        self.val_patch_mask_dir = test_patch_mask_dir
        self.batch_size = batch_size
        self.shots = shots
        self.device=device
        
        self.testset = SegWSIDataset(h5_path_lst=self.test_tumor_h5_path_lst, mask_dir=self.test_mask_dir, patch_mask_dir=self.test_patch_mask_dir)
        self.valset = SegWSIDataset(h5_path_lst=self.val_tumor_h5_path_lst, mask_dir=self.val_mask_dir, patch_mask_dir=self.val_patch_mask_dir)
        self.fewshot_paths = []

    def gen_testloader(self):
        # ### test ###
        # self.testset = SegWSIDataset(h5_path_lst=self.fewshot_paths[0][1], mask_dir=self.test_mask_dir, patch_mask_dir=self.test_patch_mask_dir)
        # ######
        testloader = DataLoader(self.testset, batch_size=1, shuffle=False) # wsidataset batch size is 1
        return testloader
        #return None

    def gen_valloader(self):
        # ### test ###
        # self.valset = SegWSIDataset(h5_path_lst=self.fewshot_paths[0][1], mask_dir=self.test_mask_dir, patch_mask_dir=self.test_patch_mask_dir)
        # ######
        valloader = DataLoader(self.valset, batch_size=1, shuffle=False) # wsidataset batch size is 1
        return valloader
    
    def gen_trainloaders(self):# 按shots倒序返回
        loaders = []
        
        shots_new = self.shots.copy()
        shots_new.sort(reverse=True)
        train_normal_h5_path_lst = self.train_normal_h5_path_lst.copy()
        train_tumor_h5_path_lst = self.train_tumor_h5_path_lst.copy()

        max_len = 0
        self.fewshot_paths = []
        while shots_new:
            shot = shots_new.pop(0)
            train_normal_h5_path_lst = random.sample(train_normal_h5_path_lst, shot)
            train_tumor_h5_path_lst = random.sample(train_tumor_h5_path_lst, shot)
            self.fewshot_paths.append([train_normal_h5_path_lst, train_tumor_h5_path_lst])

            # 统计真实的cancer ratio
            cancer_ratio, cancer_ratio_lst = utils.get_cancer_ratio(train_tumor_h5_path_lst, self.test_patch_mask_dir)
            print(f"shot: {shot}, gt cancer ratio: {cancer_ratio}")
            print(f'ratio list: {cancer_ratio_lst}')
            
            if max_len == 0: #此时shot最大
                trainset = TrainWSIDataset(labeled_h5_paths=[train_normal_h5_path_lst, train_tumor_h5_path_lst], shots=shot, task='seg', device=self.device)
                max_len = len(trainset)
            else:
                trainset = TrainWSIDataset(labeled_h5_paths=[train_normal_h5_path_lst, train_tumor_h5_path_lst], shots=shot, task='seg', device=self.device, given_len=max_len)
            trainloader = DataLoader(dataset=trainset, batch_size=self.batch_size, shuffle=False)
            loaders.append(trainloader)
        
        return loaders
    
    
class gbmlgg_sub_dataloader_generator():
    def __init__(self, gbm_dir_path, lgg_dir_path, label_csv, batch_size=256, shots=[1,2,5], val_size=32, test_ids=None):
        gbm_h5_paths = utils.get_h5_paths(gbm_dir_path)
        lgg_h5_paths = utils.get_h5_paths(lgg_dir_path)
        h5_path_lst = gbm_h5_paths + lgg_h5_paths
        
        paths, classnames = divide_multiclass_data(h5_path_lst, label_csv)
        classA_paths, classB_paths, classC_paths = paths
        
        random.shuffle(classA_paths)
        random.shuffle(classB_paths)
        random.shuffle(classC_paths)

        print("classA_paths:", len(classA_paths), "classB_paths:", len(classB_paths), "classC_paths:", len(classC_paths))
        if test_ids == None:
            half_classA_len, half_classB_len, half_classC_len = len(classA_paths) // 2, len(classB_paths) // 2, len(classC_paths) // 2
            # 注意这里的三个class的顺序不能换（train和test要一致），否则dataset load的时候的classnames就不一致了
            self.test_h5_path_lst = classA_paths[:half_classA_len] + classB_paths[:half_classB_len] + classC_paths[:half_classC_len]
            self.train_h5_paths = [classA_paths[half_classA_len:], classB_paths[half_classB_len:], classC_paths[half_classC_len:]]
            # self.train_h5_paths = [classA_paths[:half_classA_len], classB_paths[:half_classB_len], classC_paths[:half_classC_len]]

        else:
            self.test_h5_path_lst = []
            self.train_h5_paths = [[], [], []]
            for classI, classI_paths in enumerate(paths):
                for path in classI_paths:
                    h5_id = os.path.basename(path)[:-3] #去掉.h5后缀
                    if h5_id in test_ids: 
                        self.test_h5_path_lst.append(path)
                    else:
                        self.train_h5_paths[classI].append(path)
        
        self.label_csv = label_csv
        self.batch_size = batch_size
        self.shots = shots
        
        random.shuffle(self.test_h5_path_lst)
        self.valset = SubWSIDataset(h5_path_lst=self.test_h5_path_lst[:val_size], label_csv=self.label_csv)
        self.testset = SubWSIDataset(h5_path_lst=self.test_h5_path_lst, label_csv=self.label_csv)
        
        self.fewshot_paths = []
    
    def get_classnames(self):
        return self.testset.get_classnames()

    def gen_testloader(self):
        # ### test ###
        # test_paths = []
        # for class_path in self.fewshot_paths[0]:
        #     for path in class_path:
        #         test_paths.append(path)
        # self.testset = SubWSIDataset(h5_path_lst=test_paths, label_csv=self.label_csv)
        # ######
        test_loader = DataLoader(self.testset, batch_size=1, shuffle=False) # wsidataset batch size is 1
        return test_loader
    
    def gen_valloader(self):
        # ### test ###
        # test_paths = []
        # for class_path in self.fewshot_paths[0]:
        #     for path in class_path:
        #         test_paths.append(path)
        # self.valset = SubWSIDataset(h5_path_lst=test_paths, label_csv=self.label_csv)
        # ######
        val_loader = DataLoader(self.valset, batch_size=1, shuffle=False)
        return val_loader
    
    def gen_trainloaders(self):# 按shots倒序返回
        loaders = []
        shots_new = self.shots.copy()
        shots_new.sort(reverse=True)
        train_h5_paths = self.train_h5_paths.copy()
        max_len = 0
        self.fewshot_paths = []
        while shots_new:
            shot = shots_new.pop(0)
            train_h5_paths = [random.sample(class_path, shot) for class_path in train_h5_paths]
            self.fewshot_paths.append(train_h5_paths)
            if max_len == 0:
                #trainset = PatchDataset(labeled_h5_paths=train_h5_paths, classnames=self.testset.get_classnames(), task='sub')
                trainset = TrainWSIDataset(labeled_h5_paths=train_h5_paths, classnames=self.testset.get_classnames(), task='sub')
                max_len = len(trainset)
            else:
                #trainset = PatchDataset(labeled_h5_paths=train_h5_paths, classnames=self.testset.get_classnames(), task='sub', given_len=max_len)
                trainset = TrainWSIDataset(labeled_h5_paths=train_h5_paths, classnames=self.testset.get_classnames(), task='sub', given_len=max_len)
            trainloader = DataLoader(trainset, batch_size=self.batch_size, shuffle=False)
            loaders.append(trainloader) 
            train_h5_paths = trainset.get_example_paths().copy()

        return loaders




h5_path_lst = ['C:\\Users\\Administrator\\Code\\pathology\\features\\conch\\CM\\h5_files\\C3L-04331-24.h5', 'C:\\Users\\Administrator\\Code\\pathology\\features\\conch\\CM\\h5_files\\C3L-04331-25.h5']
label_csv = 'C:\\Users\\Administrator\\Code\\pathology\\data\\CM\\test_label.csv'

if __name__ == "__main__":
    dataset = PatchDataset(h5_path_lst, label_csv)
    test_loader = DataLoader(dataset, batch_size=128, shuffle=True)

    for (a, b) in test_loader:
        print(a.shape, b.shape)