WsiDataset.py

import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
import random
import h5py
import pandas as pd
import openslide
import os
import json



def read_h5(path):
    with h5py.File(path, 'r') as f:
        coords = f['coords'][:]
        features = f['features'][:]
        if 'labels' in f:
            labels = f['labels'][:]
            return torch.tensor(coords, dtype=torch.uint8), torch.tensor(features, dtype=torch.float32), torch.tensor(labels, dtype=torch.long)
        else:
            return torch.tensor(coords, dtype=torch.uint8), torch.tensor(features, dtype=torch.float32)

def divide_data(h5_path_lst, label_csv, shots=None):
    assert len(h5_path_lst) >= 2*shots, "Too less files for the given number of shots"
    random.shuffle(h5_path_lst)

    # Read labels from the CSV file
    label_df = pd.read_csv(label_csv)
    
    # Initialize lists for normal and cancer h5 paths
    normal_h5_paths = []
    cancer_h5_paths = []
    
    # Iterate over each h5 file path in the list
    for h5_path in h5_path_lst:
        # slide_id = h5_path.split('\\')[-1][:-3]  # Extract the slide ID by removing the .h5 extension
        slide_id = os.path.basename(h5_path)[:-3]
        # Find the corresponding label from the CSV file
        row = label_df[label_df['Slide_ID'] == slide_id]
        if not row.empty:
            diagnosis = row.iloc[0]['Specimen_Type']
            
            # Add to the appropriate list based on the diagnosis
            if diagnosis == 'tumor_tissue' and len(cancer_h5_paths) < shots:
                cancer_h5_paths.append(h5_path)
            elif diagnosis == 'normal_tissue' and len(normal_h5_paths) < shots:
                normal_h5_paths.append(h5_path)
        else:
            # print(f"Warning: Label of {slide_id} not found in the CSV file.")
            continue  # Skip this h5 file if no corresponding row is found in the CSV file
    
    return normal_h5_paths, cancer_h5_paths

def divide_multiclass_data(h5_path_lst, label_csv, shots=None):
    label_df = pd.read_csv(label_csv)
    paths = []
    classes = []
    for path in h5_path_lst:
        slide_id = os.path.basename(path)[:-3]  # Extract the slide ID by removing the .h5 extension
        row = label_df[label_df['slide_id'] == slide_id]
        if not row.empty:
            diagnosis = row.iloc[0]['Diagnosis']
            if diagnosis not in classes:
                classes.append(diagnosis)
                paths.append([path])
            else:
                idx = classes.index(diagnosis)
                if shots:
                    if len(paths[idx]) < shots:
                        paths[idx].append(path)
                else:
                    paths[idx].append(path)
        else:
            print(f"Warning: Label of {slide_id} not found in the CSV file.")
            continue  # Skip this h5 file if no corresponding row is found in the CSV file
    return paths, classes




##### Datasets definition #####

class PatchDataset(Dataset):
    def __init__(self, h5_path_lst=None, label_csv=None, labeled_h5_paths=None, classnames=None, shots=5, task="det", given_len=0): #给label path时不用csv和shot数据，path数量（shot）应提前切好
        self.datas = [] # 三类的patch feat，每一类一个子list # (n_cls, n_patch, dim)
        self.coords = [] # 三类的patch coord，每一类一个子list # (n_cls, n_patch, 2)
        self.paths = []
        self.classes = [] #e.g. 3类
        if task in ["det", "seg"]:
            if h5_path_lst:
                assert len(h5_path_lst) >= 2*shots, "Too less files for the given number of shots"
                random.shuffle(h5_path_lst)
                # Initialize lists for normal and cancer h5 paths
                normal_h5_paths, cancer_h5_paths = divide_data(h5_path_lst, label_csv, shots)
            else:
                normal_h5_paths, cancer_h5_paths = labeled_h5_paths[0], labeled_h5_paths[1]
            print(f"few-shot normal data: {normal_h5_paths}")
            print(f"few-shot cancer data:{cancer_h5_paths}")
            
            self.paths.append(normal_h5_paths)
            self.paths.append(cancer_h5_paths)
            self.classes = ["Normal", "Cancer"]

        elif task in ["sub"]: #subtyping dataset里是没有normal wsi的
            if h5_path_lst:
                self.paths, self.classes = divide_multiclass_data(h5_path_lst, label_csv, shots)
            else:
                self.paths, self.classes = labeled_h5_paths, classnames
            for i in range(len(self.paths)):
                print(f"few-shot classnames: {self.classes[i]}")
                for path in self.paths[i]:
                    print(path)
        
        else:
            print('task not defined')

        # 目前是随机采patch
        for path_lst in self.paths:
            coord, data = read_h5(path_lst[0])
            if len(path_lst) > 1:
                for path in path_lst[1:]:
                    coord2, data2 = read_h5(path)
                    data = torch.cat((data, data2), dim=0)
                    coord = torch.cat((coord, coord2), dim=0)
            self.datas.append(data)
            self.coords.append(coord)
        #print(self.coords)
        

        # Ensure the normal and cancer datasets have the same number of samples by filling in missing values randomly
        self.length = max([len(d) for d in self.datas])
        for i in range(len(self.datas)):
            num_missing = self.length - len(self.datas[i])
            if num_missing > 0:
                indices = torch.randint(0, len(self.datas[i]), (num_missing,))
                self.datas[i] = torch.cat((self.datas[i], self.datas[i][indices]), dim=0)
                self.coords[i] = torch.cat((self.coords[i], self.coords[i][indices]), dim=0)

        # Adjust the length of each data segment to given_len
        if given_len > 0:
            for i in range(len(self.datas)):
                current_len = len(self.datas[i])
                if current_len < given_len:
                    # Randomly copy and append data to reach given_len
                    num_missing = given_len - current_len
                    indices = torch.randint(0, current_len, (num_missing,))
                    self.datas[i] = torch.cat((self.datas[i], self.datas[i][indices]), dim=0)
                    self.coords[i] = torch.cat((self.coords[i], self.coords[i][indices]), dim=0)
                elif current_len > given_len:
                    # Randomly remove data to reach given_len
                    indices = torch.randperm(current_len)[:given_len]
                    self.datas[i] = self.datas[i][indices]
                    self.coords[i] = self.coords[i][indices]

        self.length = len(self.datas[0])

        print(f'task: {task}, patch data num: {self.length}')


    def __len__(self):
        return self.length

    def __getitem__(self, idx): #每一个class都返回一份数据和一个坐标
        data = []
        coord = []
        for idx2 in range(len(self.datas)):
            #random.shuffle(d) #对tensor用random.shuffle可能会导致某些数据点消失
            data.append(self.datas[idx2][idx])
            coord.append(self.coords[idx2][idx])
        data = torch.stack(data)
        coord = torch.stack(coord)
        return data, coord

    # def __getitem__(self, idx): #每一个class都返回一份数据和一个坐标
    #     data = []
    #     for idx2 in range(len(self.datas)):
    #         #random.shuffle(d) #对tensor用random.shuffle可能会导致某些数据点消失
    #         data.append(self.datas[idx2][idx])
    #     data = torch.stack(data)
    #     return data
    
    def shuffle_data(self):
        for i in range(len(self.datas)):
            rand_idxs = torch.randperm(len(self.datas[i]))
            self.datas[i] = self.datas[i][rand_idxs]
            self.coords[i] = self.coords[i][rand_idxs]

    def get_example_paths(self):
        return self.paths
    
    # def get_classes(self):
    #     return self.classes


class TrainWSIDataset(Dataset): #每一次只返回一张wsi里所有的patch而不是所有wsi打乱
    def __init__(self, h5_path_lst=None, label_csv=None, labeled_h5_paths=None, classnames=None, shots=5, task="det", given_len=0, device='cpu'): #给label path时不用csv和shot数据，path数量（shot）应提前切好
        self.datas = [] #三类的feat，每个子list里装的是该类的wsi，子list里元素len可能不一致（已补充到一致） # (n_cls=3, shot=10, n_patches=4w, dim=512)
        self.coords = []  # (n_cls=3, shot=10, n_patches=4w, dim=2)
        self.paths = []
        self.classes = []
        if task in ["det", "seg"]:
            if h5_path_lst:
                assert len(h5_path_lst) >= 2*shots, "Too less files for the given number of shots"
                random.shuffle(h5_path_lst)
                # Initialize lists for normal and cancer h5 paths
                normal_h5_paths, cancer_h5_paths = divide_data(h5_path_lst, label_csv, shots)
            else:
                normal_h5_paths, cancer_h5_paths = labeled_h5_paths[0], labeled_h5_paths[1]
            self.paths.append(normal_h5_paths)
            self.paths.append(cancer_h5_paths)
            self.classes = ["Normal", "Cancer"]

        elif task in ["sub"]: #subtyping dataset里是没有normal wsi的
            if h5_path_lst:
                self.paths, self.classes = divide_multiclass_data(h5_path_lst, label_csv, shots)
            else:
                self.paths, self.classes = labeled_h5_paths, classnames

        else:
            print('task not defined')
            
        for i in range(len(self.paths)):
            print(f"few-shot classnames: {self.classes[i]}")
            for path in self.paths[i]:
                print(path)

        max_len = 0

        for path_lst in self.paths:
            data_lst = []
            coord_lst = []
            for path in path_lst:
                coord, data = read_h5(path)
                data = data.to(device)
                coord = coord.to(device)
                if len(data) > max_len:
                    max_len = len(data)
                data_lst.append(data)
                coord_lst.append(coord)
            self.datas.append(data_lst)
            self.coords.append(coord_lst)
        

        # 保证每张wsi的patch数一致
        for i in range(len(self.datas)):
            for j in range(len(self.datas[i])):
                num_missing = max_len - len(self.datas[i][j])
                if num_missing > 0:
                    indices = torch.randint(0, len(self.datas[i][j]), (num_missing,))
                    self.datas[i][j] = torch.cat((self.datas[i][j], self.datas[i][j][indices]), dim=0)
                    self.coords[i][j] = torch.cat((self.coords[i][j], self.coords[i][j][indices]), dim=0)

        self.length = max_len
        # self.datas = torch.stack([torch.stack(cls_wsi) for cls_wsi in self.datas]) # (n_cls=3, shot=10, n_patches=4w, dim=512)
        # self.coords = torch.stack([torch.stack(cls_wsi) for cls_wsi in self.coords]) # (n_cls=3, shot=10, n_patches=4w, dim=2)

        print(f'task: {task}, patch each wsi: {self.length}')

    def __len__(self):
        return self.length

    def __getitem__(self, idx): #返回(batch=1024, n_cls=3, shot=10, dim=512)
        #return self.datas[:, :, idx*self.aug_times: (idx+1)*self.aug_times, :], self.coords[:, :, idx*self.aug_times: (idx+1)*self.aug_times, :]
        data = []
        coord = []
        for i in range(len(self.datas)):
            cls_data = []
            cls_coord = []
            for j in range(len(self.datas[i])):
                cls_data.append(self.datas[i][j][idx])
                cls_coord.append(self.coords[i][j][idx])
            data.append(torch.stack(cls_data))
            coord.append(torch.stack(cls_coord))
        return torch.stack(data), torch.stack(coord)
                
    
    def shuffle_data(self):
        patch_num = len(self.datas[0][0])
        shot = len(self.datas[0])
        for i in range(len(self.datas)): #datas[i]是一个class的所有wsi
            for j in range(len(self.datas[i])): #datas[i][j]是一个wsi
                rand_patch_idxs = torch.randperm(patch_num)
                # self.datas[i,j] = self.datas[i,j, rand_patch_idxs]
                # self.coords[i,j] = self.coords[i,j, rand_patch_idxs]
                self.datas[i][j] = self.datas[i][j][rand_patch_idxs]
                self.coords[i][j] = self.coords[i][j][rand_patch_idxs]
            rand_shot_idxs = torch.randperm(shot)
            # self.datas[i] = self.datas[i, rand_shot_idxs]#先shuffle里层，保证索引不会丢失
            # self.coords[i] = self.coords[i, rand_shot_idxs]
            self.datas[i] = [self.datas[i][k] for k in rand_shot_idxs]
            self.coords[i] = [self.coords[i][k] for k in rand_shot_idxs]

    def get_example_paths(self):
        return self.paths


class AnnoPatchDataset(Dataset):
    def __init__(self, h5_path_lst, division_json, shots=10, dataset='BRCA'):
        self.datas = []
        self.coords = []
        self.labels = []
        self.paths = []
        path_label_count = {}
        
        random.shuffle(h5_path_lst)

        with open(division_json, 'r', encoding='utf-8') as f:
            meta = json.load(f)[dataset]
        
        for label in meta['train_IDs'].keys():
            path_label_count[label] = 0
        
        for h5_path in h5_path_lst:
            slide_ID = os.path.splitext(os.path.basename(h5_path))[0]
            for label, ID_list in meta['train_IDs'].items():
                if slide_ID in ID_list and path_label_count[label] < shots:
                    self.paths.append(h5_path)
                    path_label_count[label] += 1
        
        for count in path_label_count.values():
            assert count == shots
        
        for path in self.paths:
            print(f'load {os.path.basename(path)}')
            coord, data, label = read_h5(path)
            self.datas.append(data)
            self.coords.append(coord)
            self.labels.append(label)
            
        self.datas = torch.cat(self.datas)
        self.coords = torch.cat(self.coords)
        self.labels = torch.cat(self.labels)

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

    def __getitem__(self, idx):
        return self.datas[idx], self.coords[idx], self.labels[idx]
    
    def get_example_paths(self):
        return self.paths
    
    def shuffle_data(self):
        pass



class DetWSIDataset(Dataset): # Only for test data
    def __init__(self, h5_path_lst=None, label_csv=None, labeled_h5_paths=None):
        self.data = []
        self.label = []
        print("loading detection dataset")

        if labeled_h5_paths:
            print(f"loading given h5s: {labeled_h5_paths}")
            for path in labeled_h5_paths[0]:
                _, data = read_h5(path)
                self.data.append(data)
                self.label.append(0)  # Normal is labeled as 0
            for path in labeled_h5_paths[1]:
                _, data = read_h5(path)
                self.data.append(data)
                self.label.append(1)  # Cancer is labeled as 1
                
                
        else:
            label_df = pd.read_csv(label_csv)
            # Iterate over each h5 file path in the list
            for h5_path in h5_path_lst:
                #slide_id = h5_path.split('/')[-1][:-3]  # Extract the slide ID by removing the .h5 extension
                slide_id = os.path.basename(h5_path)[:-3]
                print(f"loading {slide_id} in {h5_path}")
                # Find the corresponding label from the CSV file
                row = label_df[label_df['slide_id'] == slide_id]
                if not row.empty:
                    _, data = read_h5(h5_path)
                    self.data.append(data)
                    diagnosis = row.iloc[0]['Diagnosis']
                    
                    # Add to the appropriate list based on the diagnosis
                    if diagnosis == 'Tumor':
                        self.label.append(1)  # Tumor label is 1
                    elif diagnosis == 'Normal':
                        self.label.append(0)  # Normal label is 0
                else:
                    print(f"Warning: Label of {slide_id} not found in the CSV file.")
                    continue  # Skip this h5 file if no corresponding row is found in the CSV file
        print("detection dataset done")
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        data = self.data[idx]
        label = self.label[idx]
        return data, label


class SegWSIDataset(Dataset): # Only for test data
    def __init__(self, h5_path_lst, mask_dir, patch_mask_dir):
        self.h5_path_lst = h5_path_lst
        self.data = []
        self.coords = []
        self.wsi_masks = []
        self.patch_masks = []
        print("loading segmentation dataset")

        # Iterate over each h5 file path in the list
        for h5_path in h5_path_lst:
            #slide_id = h5_path.split('\\')[-1][:-3]  # Extract the slide ID by removing the .h5 extension
            slide_id = os.path.basename(h5_path)[:-3]
            print(f"loading {slide_id} in {h5_path}")
            coords, data = read_h5(h5_path)
            self.data.append(data)
            self.coords.append(coords)
            mask_file = mask_dir + "/" + slide_id + "_mask.tif"
            mask_wsi = openslide.open_slide(mask_file)
            wsi_mask = np.array(mask_wsi.read_region([0,0], 4, mask_wsi.level_dimensions[4]).convert('L')) ## gray image for 1/16
            self.wsi_masks.append(wsi_mask)

            patch_label_path = os.path.join(patch_mask_dir, slide_id + "_patch_mask.npy")
            patch_label = np.load(patch_label_path)
            self.patch_masks.append(patch_label)
        print("segmentation dataset done")
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        data = self.data[idx]
        coords = self.coords[idx]
        mask_img = self.wsi_masks[idx]
        patch_mask = self.patch_masks[idx]
        return data, coords, mask_img, patch_mask
    

class SubWSIDataset(Dataset): # Only for test data #label是从1开始的
    def __init__(self, h5_path_lst=None, label_csv=None, division_json=None, dataset=None):
        self.data = []
        self.label = []
        print("loading subtyping test/validate dataset")
        if division_json:
            with open(division_json, 'r') as f:
                meta = json.load(f)[dataset]
            for h5_path in h5_path_lst:
                for label, labeled_train_list in meta['test_IDs'].items():
                    if os.path.splitext(os.path.basename(h5_path))[0] in labeled_train_list:
                        print(f'loading h5: {os.path.basename(h5_path)} of label {label}')
                        _, data = read_h5(h5_path)
                        self.data.append(data)
                        self.label.append(int(label))
            name2label = meta['name2label']
            self.classnames = sorted(name2label.keys(), key=lambda x: name2label[x])
        else:
            self.paths, self.classnames = divide_multiclass_data(h5_path_lst, label_csv)
            for i in range(len(self.paths)):
                label = i + 1 #label是从1开始的
                print(f'loading class {self.classnames[i]}')
                for j in range(len(self.paths[i])):
                    wsi_path = self.paths[i][j]
                    print(f'loading h5: {os.path.basename(wsi_path)}')
                    _, data = read_h5(wsi_path)
                    self.data.append(data)
                    self.label.append(label)

    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        data = self.data[idx]
        label = self.label[idx]
        return data, label
    
    def get_classnames(self): #注意这里的classnames里没有normal
        return self.classnames