MicroMIL/dataset.py at main · bryanwong17/MicroMIL · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
import pandas as pd
import torch
from torch.utils.data import Dataset
import os

class Dataset_image(Dataset):
    def __init__(self, data_path, train=True):
        self.train = train
        self.data_path = data_path
        self.file_path = "D:/data/" + self.data_path + ("_train_grouped.pkl" if self.train else "_test_grouped.pkl")
        if not os.path.exists(self.file_path):
            self._preprocess_data()
        self.df = pd.read_pickle(self.file_path)

    def _preprocess_data(self):
        if self.train:
            df = pd.read_pickle("D:/data/" + self.data_path + ".pkl")[:123203]
        else:
            df = pd.read_pickle("D:/data/" + self.data_path + ".pkl")[123203:]

        grouped_data = df.groupby('patient').apply(lambda x: (x['label'].tolist(), x['idx'].tolist()))
        pd.to_pickle(grouped_data, self.file_path)

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

    def __getitem__(self, idx):
        patient = self.df.index[idx]
        labels, idxs = self.df[idx]
        imgs = []
        for idx in idxs:
            try:
                imgs.append(pd.read_pickle(f"D:/feature/image_{int(idx)}.pkl"))
            except:
                pass
        imgs = torch.stack(imgs)
        return imgs, labels[0],patient.replace(" ",'_')


class Dataset_BreakHis(Dataset):
    def __init__(self, data_path, train=True):
        self.train = train
        self.data_path = data_path
        self.file_path = "D:/data/" + self.data_path + ("_train_grouped.pkl" if self.train else "_test_grouped.pkl")
        if not os.path.exists(self.file_path):
            self._preprocess_data()
        self.df = pd.read_pickle(self.file_path)

    def _preprocess_data(self):
        if self.train:
            df = pd.read_pickle("D:/data/" + self.data_path + "_train.pkl")
        else:
            df = pd.read_pickle("D:/data/" + self.data_path + "_test.pkl")

        grouped_data = df.groupby('patient').apply(lambda x: (x['label'].tolist(), x['idx'].tolist()))
        pd.to_pickle(grouped_data, self.file_path)

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

    def __getitem__(self, idx):
        labels, idxs = self.df[idx]
        imgs = []

        for idx in idxs:
            try:
                if self.train:
                    imgs.append(pd.read_pickle(f"D:/feature_breakhis_train/image_{int(idx)}.pkl"))
                else:
                    imgs.append(pd.read_pickle(f"D:/feature_breakhis_test/image_{int(idx)}.pkl"))
            except:
                pass
        imgs = torch.stack(imgs)
        return imgs, labels[0]