forked from etetteh/OoD_Gen-Chest_Xray
-
Notifications
You must be signed in to change notification settings - Fork 0
/
merger.py
59 lines (50 loc) · 2.28 KB
/
merger.py
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
import numpy as np
import pandas as pd
import os, sys
import torch
from torchxrayvision.datasets import Dataset
class Merge_Dataset(Dataset):
def __init__(self, datasets, seed=0, num_samples=None, label_concat=False):
super(Merge_Dataset, self).__init__()
np.random.seed(seed) # Reset the seed so all runs are the same.
self.datasets = datasets
self.length = 0
self.pathologies = datasets[0].pathologies
self.which_dataset = np.zeros(0)
self.offset = np.zeros(0)
currentoffset = 0
for i, dataset in enumerate(datasets):
self.which_dataset = np.concatenate([self.which_dataset, np.zeros(num_samples)+i])
self.length += num_samples
self.offset = np.concatenate([self.offset, np.zeros(num_samples)+currentoffset])
currentoffset += num_samples
if dataset.pathologies != self.pathologies:
raise Exception("incorrect pathology alignment")
if hasattr(datasets[0], 'labels'):
self.labels = np.concatenate([d.labels for d in datasets])
else:
print("WARN: not adding .labels")
self.which_dataset = self.which_dataset.astype(int)
if label_concat:
new_labels = np.zeros([self.labels.shape[0], self.labels.shape[1]*num_samples])*np.nan
for i, shift in enumerate(self.which_dataset):
size = self.labels.shape[1]
new_labels[i,shift*size:shift*size+size] = self.labels[i]
self.labels = new_labels
try:
self.csv = pd.concat([d.csv for d in datasets])
except:
print("Could not merge dataframes (.csv not available):", sys.exc_info()[0])
self.csv = self.csv.reset_index()
def string(self):
s = self.__class__.__name__ + " num_samples={}\n".format(len(self))
for d in self.datasets:
s += "└ " + d.string().replace("\n","\n ") + "\n"
return s
def __len__(self):
return self.length
def __getitem__(self, idx):
item = self.datasets[int(self.which_dataset[idx])][idx - int(self.offset[idx])]
item["lab"] = self.labels[idx]
item["source"] = self.which_dataset[idx]
return item