import os
import logging
from collections import defaultdict
import gc
import click
import resource
import numpy as np
import cv2
import torch
import torch.nn.functional as torch_fn
from torch import nn
from torch.utils.data.dataloader import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from rocaseg.datasets import (DatasetOAIiMoSagittal2d,
DatasetOKOASagittal2d,
DatasetMAKNEESagittal2d,
sources_from_path)
from rocaseg.models import dict_models
from rocaseg.components import (dict_losses, confusion_matrix, dice_score_from_cm,
dict_optimizers, CheckpointHandler)
from rocaseg.preproc import *
from rocaseg.repro import set_ultimate_seed
from rocaseg.components.mixup import mixup_criterion, mixup_data
rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (4096, rlimit[1]))
cv2.ocl.setUseOpenCL(False)
cv2.setNumThreads(0)
logging.basicConfig()
logger = logging.getLogger('train')
logger.setLevel(logging.DEBUG)
set_ultimate_seed()
if torch.cuda.is_available():
maybe_gpu = 'cuda'
else:
maybe_gpu = 'cpu'
class ModelTrainer:
def __init__(self, config, fold_idx=None):
self.config = config
self.fold_idx = fold_idx
self.paths_weights_fold = dict()
self.paths_weights_fold['segm'] = \
os.path.join(config['path_weights'], 'segm', f'fold_{self.fold_idx}')
os.makedirs(self.paths_weights_fold['segm'], exist_ok=True)
self.paths_weights_fold['discr'] = \
os.path.join(config['path_weights'], 'discr', f'fold_{self.fold_idx}')
os.makedirs(self.paths_weights_fold['discr'], exist_ok=True)
self.path_logs_fold = \
os.path.join(config['path_logs'], f'fold_{self.fold_idx}')
os.makedirs(self.path_logs_fold, exist_ok=True)
self.handlers_ckpt = dict()
self.handlers_ckpt['segm'] = CheckpointHandler(self.paths_weights_fold['segm'])
self.handlers_ckpt['discr'] = CheckpointHandler(self.paths_weights_fold['discr'])
paths_ckpt_sel = dict()
paths_ckpt_sel['segm'] = self.handlers_ckpt['segm'].get_last_ckpt()
paths_ckpt_sel['discr'] = self.handlers_ckpt['discr'].get_last_ckpt()
# Initialize and configure the models
self.models = dict()
self.models['segm'] = (dict_models[config['model_segm']]
(input_channels=self.config['input_channels'],
output_channels=self.config['output_channels'],
center_depth=self.config['center_depth'],
pretrained=self.config['pretrained'],
path_pretrained=self.config['path_pretrained_segm'],
restore_weights=self.config['restore_weights'],
path_weights=paths_ckpt_sel['segm']))
self.models['segm'] = nn.DataParallel(self.models['segm'])
self.models['segm'] = self.models['segm'].to(maybe_gpu)
self.models['discr'] = (dict_models[config['model_discr']]
(input_channels=self.config['output_channels'],
output_channels=1,
pretrained=self.config['pretrained'],
restore_weights=self.config['restore_weights'],
path_weights=paths_ckpt_sel['discr']))
self.models['discr'] = nn.DataParallel(self.models['discr'])
self.models['discr'] = self.models['discr'].to(maybe_gpu)
# Configure the training
self.optimizers = dict()
self.optimizers['segm'] = (dict_optimizers['adam'](
self.models['segm'].parameters(),
lr=self.config['lr_segm'],
weight_decay=self.config['wd_segm']))
self.optimizers['discr'] = (dict_optimizers['adam'](
self.models['discr'].parameters(),
lr=self.config['lr_discr'],
weight_decay=self.config['wd_discr']))
self.lr_update_rule = {25: 0.1}
self.losses = dict()
self.losses['segm'] = dict_losses[self.config['loss_segm']](
num_classes=self.config['output_channels'],
)
self.losses['advers'] = dict_losses['bce_loss']()
self.losses['discr'] = dict_losses['bce_loss']()
self.losses['segm'] = self.losses['segm'].to(maybe_gpu)
self.losses['advers'] = self.losses['advers'].to(maybe_gpu)
self.losses['discr'] = self.losses['discr'].to(maybe_gpu)
self.tensorboard = SummaryWriter(self.path_logs_fold)
def run_one_epoch(self, epoch_idx, loaders):
COEFF_DISCR = 1
COEFF_SEGM = 1
COEFF_ADVERS = 0.001
fnames_acc = defaultdict(list)
metrics_acc = dict()
metrics_acc['samplew'] = defaultdict(list)
metrics_acc['batchw'] = defaultdict(list)
metrics_acc['datasetw'] = defaultdict(list)
metrics_acc['datasetw']['cm_oai'] = \
np.zeros((self.config['output_channels'],) * 2, dtype=np.uint32)
metrics_acc['datasetw']['cm_okoa'] = \
np.zeros((self.config['output_channels'],) * 2, dtype=np.uint32)
prog_bar_params = {'postfix': {'epoch': epoch_idx}, }
if self.models['segm'].training and self.models['discr'].training:
# ------------------------ Training regime ------------------------
loader_oai = loaders['oai_imo']['train']
loader_maknee = loaders['maknee']['train']
steps_oai, steps_maknee = len(loader_oai), len(loader_maknee)
steps_total = steps_oai
prog_bar_params.update({'total': steps_total,
'desc': f'Train, epoch {epoch_idx}'})
loader_oai_iter = iter(loader_oai)
loader_maknee_iter = iter(loader_maknee)
loader_oai_iter_old = None
loader_maknee_iter_old = None
with tqdm(**prog_bar_params) as prog_bar:
for step_idx in range(steps_total):
self.optimizers['segm'].zero_grad()
self.optimizers['discr'].zero_grad()
metrics_acc['batchw']['loss_total'].append(0)
try:
data_batch_oai = next(loader_oai_iter)
except StopIteration:
loader_oai_iter_old = loader_oai_iter
loader_oai_iter = iter(loader_oai)
data_batch_oai = next(loader_oai_iter)
try:
data_batch_maknee = next(loader_maknee_iter)
except StopIteration:
loader_maknee_iter_old = loader_maknee_iter
loader_maknee_iter = iter(loader_maknee)
data_batch_maknee = next(loader_maknee_iter)
xs_oai, ys_true_oai = data_batch_oai['xs'], data_batch_oai['ys']
fnames_acc['oai'].extend(data_batch_oai['path_image'])
xs_oai = xs_oai.to(maybe_gpu)
ys_true_arg_oai = torch.argmax(ys_true_oai.long().to(maybe_gpu), dim=1)
xs_maknee, _ = data_batch_maknee['xs'], data_batch_maknee['ys']
fnames_acc['maknee'].extend(data_batch_maknee['path_image'])
xs_maknee = xs_maknee.to(maybe_gpu)
# -------------- Train discriminator network -------------
# With source
ys_pred_oai = self.models['segm'](xs_oai)
ys_pred_softmax_oai = torch_fn.softmax(ys_pred_oai, dim=1)
zs_pred_oai = self.models['discr'](ys_pred_softmax_oai)
# Use 0 as a label for the source domain
loss_discr_0 = self.losses['discr'](
input=zs_pred_oai,
target=torch.zeros_like(zs_pred_oai, device=maybe_gpu))
loss_discr_0 = loss_discr_0 / 2 * COEFF_DISCR
loss_discr_0.backward(retain_graph=True)
metrics_acc['batchw']['loss_discr_0'].append(loss_discr_0.item())
metrics_acc['batchw']['loss_total'][-1] += \
metrics_acc['batchw']['loss_discr_0'][-1]
# With target
self.models['segm'] = self.models['segm'].eval()
ys_pred_maknee = self.models['segm'](xs_maknee)
self.models['segm'] = self.models['segm'].train()
ys_pred_softmax_maknee = torch_fn.softmax(ys_pred_maknee, dim=1)
zs_pred_maknee = self.models['discr'](ys_pred_softmax_maknee)
# Use 1 as a label for the target domain
loss_discr_1 = self.losses['discr'](
input=zs_pred_maknee,
target=torch.ones_like(zs_pred_maknee, device=maybe_gpu))
loss_discr_1 = loss_discr_1 / 2 * COEFF_DISCR
loss_discr_1.backward()
metrics_acc['batchw']['loss_discr_1'].append(loss_discr_1.item())
metrics_acc['batchw']['loss_total'][-1] += \
metrics_acc['batchw']['loss_discr_1'][-1]
self.models['segm'].zero_grad()
self.optimizers['discr'].step()
self.models['discr'].zero_grad()
# ---------------- Train segmentation network ------------
# With source
if not self.config['with_mixup']:
ys_pred_oai = self.models['segm'](xs_oai)
loss_segm = self.losses['segm'](input_=ys_pred_oai,
target=ys_true_arg_oai)
else:
xs_mixup, ys_mixup_a, ys_mixup_b, lambda_mixup = mixup_data(
x=xs_oai, y=ys_true_arg_oai,
alpha=self.config['mixup_alpha'], device=maybe_gpu)
ys_pred_oai = self.models['segm'](xs_mixup)
loss_segm = mixup_criterion(criterion=self.losses['segm'],
pred=ys_pred_oai,
y_a=ys_mixup_a,
y_b=ys_mixup_b,
lam=lambda_mixup)
loss_segm.backward(retain_graph=True)
loss_segm = loss_segm * COEFF_SEGM
metrics_acc['batchw']['loss_segm'].append(loss_segm.item())
metrics_acc['batchw']['loss_total'][-1] += \
metrics_acc['batchw']['loss_segm'][-1]
# With target
self.models['segm'] = self.models['segm'].eval()
ys_pred_maknee = self.models['segm'](xs_maknee)
self.models['segm'] = self.models['segm'].train()
ys_pred_softmax_maknee = torch_fn.softmax(ys_pred_maknee, dim=1)
zs_pred_maknee = self.models['discr'](ys_pred_softmax_maknee)
# Use 0 as a label for the source domain
loss_advers = self.losses['advers'](
input=zs_pred_maknee,
target=torch.zeros_like(zs_pred_maknee, device=maybe_gpu))
loss_advers = loss_advers * COEFF_ADVERS
loss_advers.backward()
metrics_acc['batchw']['loss_advers'].append(loss_advers.item())
metrics_acc['batchw']['loss_total'][-1] += \
metrics_acc['batchw']['loss_advers'][-1]
self.models['discr'].zero_grad()
self.optimizers['segm'].step()
if step_idx % 10 == 0:
self.tensorboard.add_scalars(
f'fold_{self.fold_idx}/losses_train',
{'discr_0_batchw': metrics_acc['batchw']['loss_discr_0'][-1],
'discr_1_batchw': metrics_acc['batchw']['loss_discr_1'][-1],
'discr_sum_batchw':
(metrics_acc['batchw']['loss_discr_0'][-1] +
metrics_acc['batchw']['loss_discr_1'][-1]),
'segm_batchw': metrics_acc['batchw']['loss_segm'][-1],
'advers_batchw':
metrics_acc['batchw']['loss_advers'][-1],
'total_batchw': metrics_acc['batchw']['loss_total'][-1],
}, global_step=(epoch_idx * steps_total + step_idx))
prog_bar.update(1)
del [loader_oai_iter_old, loader_maknee_iter_old]
gc.collect()
else:
# ----------------------- Validation regime -----------------------
loader_oai = loaders['oai_imo']['val']
loader_okoa = loaders['okoa']['val']
loader_maknee = loaders['maknee']['val']
steps_oai, steps_okoa, steps_maknee = len(loader_oai), len(loader_okoa), len(loader_maknee)
steps_total = steps_oai
prog_bar_params.update({'total': steps_total,
'desc': f'Validate, epoch {epoch_idx}'})
loader_oai_iter = iter(loader_oai)
loader_okoa_iter = iter(loader_okoa)
loader_maknee_iter = iter(loader_maknee)
loader_oai_iter_old = None
loader_okoa_iter_old = None
loader_maknee_iter_old = None
with torch.no_grad(), tqdm(**prog_bar_params) as prog_bar:
for step_idx in range(steps_total):
metrics_acc['batchw']['loss_total'].append(0)
try:
data_batch_oai = next(loader_oai_iter)
except StopIteration:
loader_oai_iter_old = loader_oai_iter
loader_oai_iter = iter(loader_oai)
data_batch_oai = next(loader_oai_iter)
try:
data_batch_okoa = next(loader_okoa_iter)
except StopIteration:
loader_okoa_iter_old = loader_okoa_iter
loader_okoa_iter = iter(loader_okoa)
data_batch_okoa = next(loader_okoa_iter)
try:
data_batch_maknee = next(loader_maknee_iter)
except StopIteration:
loader_maknee_iter_old = loader_maknee_iter
loader_maknee_iter = iter(loader_maknee)
data_batch_maknee = next(loader_maknee_iter)
xs_oai, ys_true_oai = data_batch_oai['xs'], data_batch_oai['ys']
fnames_acc['oai'].extend(data_batch_oai['path_image'])
xs_oai = xs_oai.to(maybe_gpu)
ys_true_arg_oai = torch.argmax(ys_true_oai.long().to(maybe_gpu), dim=1)
xs_maknee, _ = data_batch_maknee['xs'], data_batch_maknee['ys']
fnames_acc['maknee'].extend(data_batch_maknee['path_image'])
xs_maknee = xs_maknee.to(maybe_gpu)
# -------------- Validate discriminator network -------------
# With source
ys_pred_oai = self.models['segm'](xs_oai)
ys_pred_softmax_oai = torch_fn.softmax(ys_pred_oai, dim=1)
zs_pred_oai = self.models['discr'](ys_pred_softmax_oai)
# Use 0 as a label for the source domain
loss_discr_0 = self.losses['discr'](
input=zs_pred_oai,
target=torch.zeros_like(zs_pred_oai, device=maybe_gpu))
loss_discr_0 = loss_discr_0 / 2 * COEFF_DISCR
metrics_acc['batchw']['loss_discr_0'].append(loss_discr_0.item())
metrics_acc['batchw']['loss_total'][-1] += \
metrics_acc['batchw']['loss_discr_0'][-1]
# With target
ys_pred_maknee = self.models['segm'](xs_maknee)
ys_pred_softmax_maknee = torch_fn.softmax(ys_pred_maknee, dim=1)
zs_pred_maknee = self.models['discr'](ys_pred_softmax_maknee)
# Use 1 as a label for the target domain
loss_discr_1 = self.losses['discr'](
input=zs_pred_maknee,
target=torch.ones_like(zs_pred_oai, device=maybe_gpu))
loss_discr_1 = loss_discr_1 / 2 * COEFF_DISCR
metrics_acc['batchw']['loss_discr_1'].append(loss_discr_1.item())
metrics_acc['batchw']['loss_total'][-1] += \
metrics_acc['batchw']['loss_discr_1'][-1]
# ---------------- Validate segmentation network ------------
# With source
if not self.config['with_mixup']:
ys_pred_oai = self.models['segm'](xs_oai)
loss_segm = self.losses['segm'](input_=ys_pred_oai,
target=ys_true_arg_oai)
else:
xs_mixup, ys_mixup_a, ys_mixup_b, lambda_mixup = mixup_data(
x=xs_oai, y=ys_true_arg_oai,
alpha=self.config['mixup_alpha'], device=maybe_gpu)
ys_pred_oai = self.models['segm'](xs_mixup)
loss_segm = mixup_criterion(criterion=self.losses['segm'],
pred=ys_pred_oai,
y_a=ys_mixup_a,
y_b=ys_mixup_b,
lam=lambda_mixup)
loss_segm = loss_segm * COEFF_SEGM
metrics_acc['batchw']['loss_segm'].append(loss_segm.item())
metrics_acc['batchw']['loss_total'][-1] += \
metrics_acc['batchw']['loss_segm'][-1]
# With target
ys_pred_maknee = self.models['segm'](xs_maknee)
ys_pred_softmax_maknee = torch_fn.softmax(ys_pred_maknee, dim=1)
zs_pred_maknee = self.models['discr'](ys_pred_softmax_maknee)
# Use 0 as a label for the source domain
loss_advers = self.losses['advers'](
input=zs_pred_maknee,
target=torch.zeros_like(zs_pred_maknee, device=maybe_gpu))
loss_advers = loss_advers * COEFF_ADVERS
metrics_acc['batchw']['loss_advers'].append(loss_advers.item())
metrics_acc['batchw']['loss_total'][-1] += \
metrics_acc['batchw']['loss_advers'][-1]
if step_idx % 10 == 0:
self.tensorboard.add_scalars(
f'fold_{self.fold_idx}/losses_val',
{'discr_0_batchw': metrics_acc['batchw']['loss_discr_0'][-1],
'discr_1_batchw': metrics_acc['batchw']['loss_discr_1'][-1],
'discr_sum_batchw':
(metrics_acc['batchw']['loss_discr_0'][-1] +
metrics_acc['batchw']['loss_discr_1'][-1]),
'segm_batchw': metrics_acc['batchw']['loss_segm'][-1],
'advers_batchw':
metrics_acc['batchw']['loss_advers'][-1],
'total_batchw': metrics_acc['batchw']['loss_total'][-1],
}, global_step=(epoch_idx * steps_total + step_idx))
# ------------------ Calculate metrics -------------------
ys_pred_arg_np_oai = torch.argmax(ys_pred_softmax_oai, 1).to('cpu').numpy()
ys_true_arg_np_oai = ys_true_arg_oai.to('cpu').numpy()
metrics_acc['datasetw']['cm_oai'] += confusion_matrix(
ys_pred_arg_np_oai, ys_true_arg_np_oai,
self.config['output_channels'])
# Don't consider repeating entries for the metrics calculation
if step_idx < steps_okoa:
xs_okoa, ys_true_okoa = data_batch_okoa['xs'], data_batch_okoa['ys']
fnames_acc['okoa'].extend(data_batch_okoa['path_image'])
xs_okoa = xs_okoa.to(maybe_gpu)
ys_pred_okoa = self.models['segm'](xs_okoa)
ys_true_arg_okoa = torch.argmax(ys_true_okoa.long().to(maybe_gpu), dim=1)
ys_pred_softmax_okoa = torch_fn.softmax(ys_pred_okoa, dim=1)
ys_pred_arg_np_okoa = torch.argmax(ys_pred_softmax_okoa, 1).to('cpu').numpy()
ys_true_arg_np_okoa = ys_true_arg_okoa.to('cpu').numpy()
metrics_acc['datasetw']['cm_okoa'] += confusion_matrix(
ys_pred_arg_np_okoa, ys_true_arg_np_okoa,
self.config['output_channels'])
prog_bar.update(1)
del [loader_oai_iter_old, loader_okoa_iter_old, loader_maknee_iter_old]
gc.collect()
for k, v in metrics_acc['samplew'].items():
metrics_acc['samplew'][k] = np.asarray(v)
metrics_acc['datasetw']['dice_score_oai'] = np.asarray(
dice_score_from_cm(metrics_acc['datasetw']['cm_oai']))
metrics_acc['datasetw']['dice_score_okoa'] = np.asarray(
dice_score_from_cm(metrics_acc['datasetw']['cm_okoa']))
return metrics_acc, fnames_acc
def fit(self, loaders):
epoch_idx_best = -1
loss_best = float('inf')
metrics_train_best = dict()
fnames_train_best = []
metrics_val_best = dict()
fnames_val_best = []
for epoch_idx in range(self.config['epoch_num']):
self.models = {n: m.train() for n, m in self.models.items()}
metrics_train, fnames_train = \
self.run_one_epoch(epoch_idx, loaders)
# Process the accumulated metrics
for k, v in metrics_train['batchw'].items():
if k.startswith('loss'):
metrics_train['datasetw'][k] = np.mean(np.asarray(v))
else:
logger.warning(f'Non-processed batch-wise entry: {k}')
self.models = {n: m.eval() for n, m in self.models.items()}
metrics_val, fnames_val = \
self.run_one_epoch(epoch_idx, loaders)
# Process the accumulated metrics
for k, v in metrics_val['batchw'].items():
if k.startswith('loss'):
metrics_val['datasetw'][k] = np.mean(np.asarray(v))
else:
logger.warning(f'Non-processed batch-wise entry: {k}')
# Learning rate update
for s, m in self.lr_update_rule.items():
if epoch_idx == s:
for name, optim in self.optimizers.items():
for param_group in optim.param_groups:
param_group['lr'] *= m
# Add console logging
logger.info(f'Epoch: {epoch_idx}')
for subset, metrics in (('train', metrics_train),
('val', metrics_val)):
logger.info(f'{subset} metrics:')
for k, v in metrics['datasetw'].items():
logger.info(f'{k}: \n{v}')
# Add TensorBoard logging
for subset, metrics in (('train', metrics_train),
('val', metrics_val)):
# Log only dataset-reduced metrics
for k, v in metrics['datasetw'].items():
if isinstance(v, np.ndarray):
self.tensorboard.add_scalars(
f'fold_{self.fold_idx}/{k}_{subset}',
{f'class{i}': e for i, e in enumerate(v.ravel().tolist())},
global_step=epoch_idx)
elif isinstance(v, (str, int, float)):
self.tensorboard.add_scalar(
f'fold_{self.fold_idx}/{k}_{subset}',
float(v),
global_step=epoch_idx)
else:
logger.warning(f'{k} is of unsupported dtype {v}')
for name, optim in self.optimizers.items():
for param_group in optim.param_groups:
self.tensorboard.add_scalar(
f'fold_{self.fold_idx}/learning_rate/{name}',
param_group['lr'],
global_step=epoch_idx)
# Save the model
loss_curr = metrics_val['datasetw']['loss_total']
if loss_curr < loss_best:
loss_best = loss_curr
epoch_idx_best = epoch_idx
metrics_train_best = metrics_train
metrics_val_best = metrics_val
fnames_train_best = fnames_train
fnames_val_best = fnames_val
self.handlers_ckpt['segm'].save_new_ckpt(
model=self.models['segm'],
model_name=self.config['model_segm'],
fold_idx=self.fold_idx,
epoch_idx=epoch_idx)
self.handlers_ckpt['discr'].save_new_ckpt(
model=self.models['discr'],
model_name=self.config['model_discr'],
fold_idx=self.fold_idx,
epoch_idx=epoch_idx)
msg = (f'Finished fold {self.fold_idx} '
f'with the best loss {loss_best:.5f} '
f'on epoch {epoch_idx_best}, '
f'weights: ({self.paths_weights_fold})')
logger.info(msg)
return (metrics_train_best, fnames_train_best,
metrics_val_best, fnames_val_best)
@click.command()
@click.option('--path_data_root', default='../../data')
@click.option('--path_experiment_root', default='../../results/temporary')
@click.option('--model_segm', default='unet_lext')
@click.option('--center_depth', default=1, type=int)
@click.option('--model_discr', default='discriminator_a')
@click.option('--pretrained', is_flag=True)
@click.option('--path_pretrained_segm', type=str, help='Path to .pth file')
@click.option('--restore_weights', is_flag=True)
@click.option('--input_channels', default=1, type=int)
@click.option('--output_channels', default=1, type=int)
@click.option('--mask_mode', default='all_unitibial_unimeniscus', type=str)
@click.option('--sample_mode', default='x_y', type=str)
@click.option('--loss_segm', default='multi_ce_loss')
@click.option('--lr_segm', default=0.0001, type=float)
@click.option('--lr_discr', default=0.0001, type=float)
@click.option('--wd_segm', default=5e-5, type=float)
@click.option('--wd_discr', default=5e-5, type=float)
@click.option('--optimizer_segm', default='adam')
@click.option('--optimizer_discr', default='adam')
@click.option('--batch_size', default=64, type=int)
@click.option('--epoch_size', default=1.0, type=float)
@click.option('--epoch_num', default=2, type=int)
@click.option('--fold_num', default=5, type=int)
@click.option('--fold_idx', default=-1, type=int)
@click.option('--fold_idx_ignore', multiple=True, type=int)
@click.option('--num_workers', default=1, type=int)
@click.option('--seed_trainval_test', default=0, type=int)
@click.option('--with_mixup', is_flag=True)
@click.option('--mixup_alpha', default=1, type=float)
def main(**config):
config['path_data_root'] = os.path.abspath(config['path_data_root'])
config['path_experiment_root'] = os.path.abspath(config['path_experiment_root'])
config['path_weights'] = os.path.join(config['path_experiment_root'], 'weights')
config['path_logs'] = os.path.join(config['path_experiment_root'], 'logs_train')
os.makedirs(config['path_weights'], exist_ok=True)
os.makedirs(config['path_logs'], exist_ok=True)
logging_fh = logging.FileHandler(
os.path.join(config['path_logs'], 'main_{}.log'.format(config['fold_idx'])))
logging_fh.setLevel(logging.DEBUG)
logger.addHandler(logging_fh)
# Collect the available and specified sources
sources = sources_from_path(path_data_root=config['path_data_root'],
selection=('oai_imo', 'okoa', 'maknee'),
with_folds=True,
fold_num=config['fold_num'],
seed_trainval_test=config['seed_trainval_test'])
# Build a list of folds to run on
if config['fold_idx'] == -1:
fold_idcs = list(range(config['fold_num']))
else:
fold_idcs = [config['fold_idx'], ]
for g in config['fold_idx_ignore']:
fold_idcs = [i for i in fold_idcs if i != g]
# Train each fold separately
fold_scores = dict()
# Use straightforward fold allocation strategy
folds = list(zip(sources['oai_imo']['trainval_folds'],
sources['okoa']['trainval_folds'],
sources['maknee']['trainval_folds']))
for fold_idx, idcs_subsets in enumerate(folds):
if fold_idx not in fold_idcs:
continue
logger.info(f'Training fold {fold_idx}')
(sources['oai_imo']['train_idcs'], sources['oai_imo']['val_idcs']) = idcs_subsets[0]
(sources['okoa']['train_idcs'], sources['okoa']['val_idcs']) = idcs_subsets[1]
(sources['maknee']['train_idcs'], sources['maknee']['val_idcs']) = idcs_subsets[2]
sources['oai_imo']['train_df'] = sources['oai_imo']['trainval_df'].iloc[sources['oai_imo']['train_idcs']]
sources['oai_imo']['val_df'] = sources['oai_imo']['trainval_df'].iloc[sources['oai_imo']['val_idcs']]
sources['okoa']['train_df'] = sources['okoa']['trainval_df'].iloc[sources['okoa']['train_idcs']]
sources['okoa']['val_df'] = sources['okoa']['trainval_df'].iloc[sources['okoa']['val_idcs']]
sources['maknee']['train_df'] = sources['maknee']['trainval_df'].iloc[sources['maknee']['train_idcs']]
sources['maknee']['val_df'] = sources['maknee']['trainval_df'].iloc[sources['maknee']['val_idcs']]
for n, s in sources.items():
logger.info('Made {} train-val split, number of samples: {}, {}'
.format(n, len(s['train_df']), len(s['val_df'])))
datasets = defaultdict(dict)
datasets['oai_imo']['train'] = DatasetOAIiMoSagittal2d(
df_meta=sources['oai_imo']['train_df'],
mask_mode=config['mask_mode'],
sample_mode=config['sample_mode'],
transforms=[
PercentileClippingAndToFloat(cut_min=10, cut_max=99),
CenterCrop(height=300, width=300),
HorizontalFlip(prob=.5),
GammaCorrection(gamma_range=(0.5, 1.5), prob=.5),
OneOf([
DualCompose([
Scale(ratio_range=(0.7, 0.8), prob=1.),
Scale(ratio_range=(1.5, 1.6), prob=1.),
]),
NoTransform()
]),
Crop(output_size=(300, 300)),
BilateralFilter(d=5, sigma_color=50, sigma_space=50, prob=.3),
Normalize(mean=0.252699, std=0.251142),
ToTensor(),
])
datasets['okoa']['train'] = DatasetOKOASagittal2d(
df_meta=sources['okoa']['train_df'],
mask_mode='background_femoral_unitibial',
sample_mode=config['sample_mode'],
transforms=[
PercentileClippingAndToFloat(cut_min=10, cut_max=99),
CenterCrop(height=300, width=300),
HorizontalFlip(prob=.5),
GammaCorrection(gamma_range=(0.5, 1.5), prob=.5),
OneOf([
DualCompose([
Scale(ratio_range=(0.7, 0.8), prob=1.),
Scale(ratio_range=(1.5, 1.6), prob=1.),
]),
NoTransform()
]),
Crop(output_size=(300, 300)),
BilateralFilter(d=5, sigma_color=50, sigma_space=50, prob=.3),
Normalize(mean=0.252699, std=0.251142),
ToTensor(),
])
datasets['maknee']['train'] = DatasetMAKNEESagittal2d(
df_meta=sources['maknee']['train_df'],
mask_mode='',
sample_mode=config['sample_mode'],
transforms=[
PercentileClippingAndToFloat(cut_min=10, cut_max=99),
CenterCrop(height=300, width=300),
HorizontalFlip(prob=.5),
GammaCorrection(gamma_range=(0.5, 1.5), prob=.5),
OneOf([
DualCompose([
Scale(ratio_range=(0.7, 0.8), prob=1.),
Scale(ratio_range=(1.5, 1.6), prob=1.),
]),
NoTransform()
]),
Crop(output_size=(300, 300)),
BilateralFilter(d=5, sigma_color=50, sigma_space=50, prob=.3),
Normalize(mean=0.252699, std=0.251142),
ToTensor(),
])
datasets['oai_imo']['val'] = DatasetOAIiMoSagittal2d(
df_meta=sources['oai_imo']['val_df'],
mask_mode=config['mask_mode'],
sample_mode=config['sample_mode'],
transforms=[
PercentileClippingAndToFloat(cut_min=10, cut_max=99),
CenterCrop(height=300, width=300),
Normalize(mean=0.252699, std=0.251142),
ToTensor()
])
datasets['okoa']['val'] = DatasetOKOASagittal2d(
df_meta=sources['okoa']['val_df'],
mask_mode='background_femoral_unitibial',
sample_mode=config['sample_mode'],
transforms=[
PercentileClippingAndToFloat(cut_min=10, cut_max=99),
CenterCrop(height=300, width=300),
Normalize(mean=0.252699, std=0.251142),
ToTensor()
])
datasets['maknee']['val'] = DatasetMAKNEESagittal2d(
df_meta=sources['maknee']['val_df'],
mask_mode='',
sample_mode=config['sample_mode'],
transforms=[
PercentileClippingAndToFloat(cut_min=10, cut_max=99),
CenterCrop(height=300, width=300),
Normalize(mean=0.252699, std=0.251142),
ToTensor()
])
loaders = defaultdict(dict)
loaders['oai_imo']['train'] = DataLoader(
datasets['oai_imo']['train'],
batch_size=int(config['batch_size'] / 2),
shuffle=True,
num_workers=config['num_workers'],
drop_last=True)
loaders['oai_imo']['val'] = DataLoader(
datasets['oai_imo']['val'],
batch_size=int(config['batch_size'] / 2),
shuffle=False,
num_workers=config['num_workers'],
drop_last=True)
loaders['okoa']['train'] = DataLoader(
datasets['okoa']['train'],
batch_size=int(config['batch_size'] / 2),
shuffle=True,
num_workers=config['num_workers'],
drop_last=True)
loaders['okoa']['val'] = DataLoader(
datasets['okoa']['val'],
batch_size=int(config['batch_size'] / 2),
shuffle=False,
num_workers=config['num_workers'],
drop_last=True)
loaders['maknee']['train'] = DataLoader(
datasets['maknee']['train'],
batch_size=int(config['batch_size'] / 2),
shuffle=True,
num_workers=config['num_workers'],
drop_last=True)
loaders['maknee']['val'] = DataLoader(
datasets['maknee']['val'],
batch_size=int(config['batch_size'] / 2),
shuffle=False,
num_workers=config['num_workers'],
drop_last=True)
trainer = ModelTrainer(config=config, fold_idx=fold_idx)
tmp = trainer.fit(loaders=loaders)
metrics_train, fnames_train, metrics_val, fnames_val = tmp
fold_scores[fold_idx] = (metrics_val['datasetw']['dice_score_oai'],
metrics_val['datasetw']['dice_score_okoa'])
trainer.tensorboard.close()
logger.info(f'Fold scores:\n{repr(fold_scores)}')
if __name__ == '__main__':
main()
