Commit db8b74eb authored by damiana s's avatar damiana s
Browse files

clean

parent dd7d067f
import torch
import numpy as np
import pandas as pd
import os
import SimpleITK as sitk
from dicom_utils import augmentation as aug, processing as dup
from torch.utils.data import Dataset
def augment_3D(image, mode, size):
N_CHANNELS = image.shape[0]
seq_list = []
for ID_SEQ in range(N_CHANNELS):
image_seq = image[ID_SEQ,:,:,:]
image_seq = sitk.GetImageFromArray(image_seq)
if mode == 'train':
#morphological augmentation
image_seq = aug.augment_morph([image_seq])[0]
fg = aug.get_gauss_noise()
image_seq = fg.Execute(image_seq)
image_seq = dup.resample(image_seq, size = (size, size, size)) #qst è numpy
seq_list.append(sitk.GetArrayFromImage(image_seq))
image = np.stack(seq_list, axis=0)
return(image)
class NumpyCSVDataset(Dataset):
def __init__(self, data_dir, label_file, label_col, size, augmentation_function=augment_3D, mode='train'):
super(NumpyCSVDataset, self).__init__()
self.data_dir = data_dir
self.size = size
self.augmentation = augmentation_function
self.mode = mode
files = os.listdir(data_dir)
self._samples = [x.split('.')[0] for x in files]
labels = pd.read_csv(label_file, sep=',', header=0, index_col=0)[label_col].to_frame()
labels = labels.loc[self._samples]
self._labels = labels
self._indexes = np.arange(len(self._samples))
def __getitem__(self, idx, no_data=False):
sample = self._samples[self._indexes[idx]]
label = self._labels.values[self._indexes[idx]][0]
data_file = f'{self.data_dir}/{sample}.npy'
data = np.load(data_file)
data = self.augmentation(data, self.mode, self.size) #qst è numpy
data = torch.Tensor(data) #qst è tensor - la network vuole tensor
output = {'data': data, 'target': label, 'sample': sample}
return(output)
def __len__(self):
return(len(self._indexes))
import numpy as np
import pandas as pd
import os
import pickle
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import gc
from .dapmodel import DAPModel
from ..lr_schedulers import ConstantScheduler
'''
Example of compatible Network class
class CompatibleNet(nn.Module):
def __init__(self, **params):
#! network parameters should be passed as a dictionary
# set parameters with default values if the parameter is not in the **params dict:
self.parameter1 = params.get(key, default)
#define modules
def forward(self, x):
# as usual
def get_params(self):
return a dictionary with the parameters
def initialize_weights(self):
#initialize network weights (needed to reset the network for every cycle/fold)
'''
class DeepLearningModel(DAPModel):
def __init__(self, network, params=None, name = ''):
self._other_params = {'random_labels':False,
'epochs':100,
'batch_size': 16,
'optimizer': torch.optim.Adam,
'criterion' : nn.CrossEntropyLoss,
'lr_scheduler': ConstantScheduler,
'lr_start': 0.001,
'device': 'cpu',
'n_gpu':1, #TODO: implement
'n_layers': None # how many layers to train (starting from end)
}
self.model_params = dict()
self.network = network
# TODO: check that network has get_params and initialize_weights
self.model = None
self.init_state = None
self.name = name
if params is not None:
self.set_params(params)
def set_params(self, params):
'''
Set the parameters of the model.
At this point we need to initialize a new model
'''
for P in params:
if P in self._other_params.keys():
self._other_params[P] = params[P]
else:
self.model_params[P] = params[P]
device = torch.device(self._other_params['device'])
self.model = self.network(**self.model_params).to(device)
self.model.initialize_weights()
if self._other_params['device'] == 'cuda' and self._other_params['n_gpu']>1:
self.model = nn.DataParallel(self.model)
def get_params(self):
return(self.model.get_params(), self._other_params)
def fit(self, dataset_train):
if self._other_params['random_labels']:
dataset_train.__shuffle_labels__()
dataset_train.__shuffle__() #this step is required otherwise the samples are ordered by class
EPOCHS = self._other_params['epochs']
BATCH = self._other_params['batch_size']
OPTIM = self._other_params['optimizer']
LR_START = self._other_params['lr_start']
criterion = self._other_params['criterion']() #TODO: add weights
loader_train = DataLoader(dataset_train, batch_size=BATCH, drop_last=True) # TODO: add other params here
lr_scheduler = self._other_params['lr_scheduler'](LR_START, len(loader_train), EPOCHS)
device = torch.device(self._other_params['device'])
if self._other_params['device'] == 'cuda':
self.model.module.initialize_weights()
else:
self.model.initialize_weights()
self.model.train()
losses_tr = []
iteration = 0
for epoch in range(EPOCHS):
for j, data in enumerate(loader_train):
#set the optimizer
LR = lr_scheduler.get_lr(epoch, j, iteration)
optimizer = OPTIM(list(self.model.parameters()), lr = LR) # TODO: add n_layers to optim
optimizer.zero_grad()
#get the data and targets
input_tr = data['data'].float().to(device)
labels_tr = data['target'].to(device)
#forward, compute loss, backward, optimize weights
outputs_tr = self.model(input_tr)
loss = criterion(outputs_tr, labels_tr)
loss.backward()
optimizer.step()
del input_tr, labels_tr
gc.collect()
losses_tr.append(loss.item())
iteration +=1
def predict(self, dataset_valid):
loader_valid = DataLoader(dataset_valid, batch_size=1) # TODO: add other params here
device = torch.device(self._other_params['device'])
self.model.eval()
y_pred = []
y_true = []
with torch.no_grad():
for j, data in enumerate(loader_valid):
images_vl = data['data'].float().to(device)
labels_vl = data['target'].to(device)
outputs_vl = self.model.forward(images_vl)
_, pred = torch.max(outputs_vl, 1)
y_pred.append(pred.item())
y_true.append(labels_vl.item())
y_pred = np.array(y_pred)
y_true = np.array(y_true)
return(y_true, y_pred)
def fit_best(self, dataset, metrics):
best_point = metrics.get_best_param_set()
self.set_params(best_point)
self.fit(dataset)
def save(self, component_dir):
filename = os.path.join(component_dir, 'DAP_model.pkl')
with open(filename, 'wb') as f:
pickle.dump(self, f, pickle.HIGHEST_PROTOCOL)
torch.save(self.model, os.path.join(component_dir, 'torch_model.th'))
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Sep 5 11:31:52 2018
@author: andrea
Su modello di : radiomics_pipeline/deep_features
"""
import torch.nn as nn
import torch
class CNN3D_NET(nn.Module):
def __init__(self):
super(CNN3D_NET, self).__init__()
self.n_classes = 2
self.n_channels = 2
dropout = 0.5
self.CT_branch = nn.Sequential(
nn.Conv3d(1, 20, 3), #CT branch
nn.BatchNorm3d(20),
nn.ReLU(True),
nn.MaxPool3d(2,stride=2),
nn.Dropout3d(dropout),
nn.Conv3d(20, 40, 3),
nn.BatchNorm3d(40),
nn.ReLU(True),
nn.MaxPool3d(2,stride=2),
nn.Dropout3d(dropout/2),
nn.Conv3d(40, 80, 3),
nn.BatchNorm3d(80),
nn.ReLU(True),
nn.MaxPool3d(2,stride=2),
nn.Dropout3d(dropout/4)
)
self.PT_branch = nn.Sequential(
nn.Conv3d(1, 20, 3), #PT branch
nn.BatchNorm3d(20),
nn.ReLU(True),
nn.MaxPool3d(2,stride=2),
nn.Dropout3d(dropout),
nn.Conv3d(20, 40, 3),
nn.BatchNorm3d(40),
nn.ReLU(True),
nn.MaxPool3d(2,stride=2),
nn.Dropout3d(dropout/2),
nn.Conv3d(40, 80, 3),
nn.BatchNorm3d(80),
nn.ReLU(True),
nn.MaxPool3d(2,stride=2),
nn.Dropout3d(dropout/4))
self.linear = nn.Sequential(nn.Linear(80*2*8, 50),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(50, self.n_classes),
nn.Softmax(1))
def forward(self, x): #x = N_batch x N_ch * s x s x s
x_CT = x[:, 0, :,:,:].unsqueeze(1)#x = N_batch * s x s x s
features_CT = self.CT_branch(x_CT)
features_CT = features_CT.view(x.shape[0], -1) #n_batxh * X
x_PT = x[:, 1, :,:,:].unsqueeze(1)
features_PT = self.PT_branch(x_PT)
features_PT = features_PT.view(x.shape[0], -1)
out_merged = torch.cat([features_CT, features_PT], dim=1)
out = self.linear(out_merged)
return(out)
def extract_features(self, x):
x_CT = x[:, 0, :,:,:].unsqueeze(1)
features_CT = self.CT_branch(x_CT)
features_CT = features_CT.view(x.shape[0], -1)
x_PT = x[:, 1, :,:,:].unsqueeze(1)
features_PT = self.PT_branch(x_PT)
features_PT = features_PT.view(x.shape[0], -1)
out_merged = torch.cat([x_CT, x_PT], dim=1)
return(out_merged)
def initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Conv3d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, nn.BatchNorm3d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment