#!/usr/bin/env python
# coding: utf-8
# In[]:
########################### IMPORTS #############################################
import torch
import torch.nn as nn
import torch.nn.functional as F
import monai
from monai.networks.nets import UNet,UNet_RBF
from monai.networks.utils import one_hot
from monai.transforms import (
AsDiscrete,
AddChanneld,
AsChannelFirstd,
Compose,
LoadNiftid,
RandCropByPosNegLabeld,
RandRotate90d,
ScaleIntensityd,
ToTensord,
)
from monai.visualize import plot_2d_or_3d_image
from monai.data.utils import list_data_collate, worker_init_fn
from monai.inferers import sliding_window_inference
from monai.metrics import DiceMetric
from monai.metrics import compute_meandice
from torch.autograd import Variable
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
import torch.backends.cudnn as cudnn
import torchvision
from torchvision import datasets, models, transforms
import csv
import time
import SimpleITK as sitk
from os.path import splitext,basename
import random
from glob import glob
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from copy import copy
import os
import numpy as np
from torch.utils.tensorboard import SummaryWriter
#from global_tools.tools import display_loading_bar
from class_modalities.transforms import LoadNifti, Roi2Mask, ResampleReshapeAlign, Sitk2Numpy, ConcatModality
from monai.utils import first, set_determinism
##################
train_transforms = Compose(
[ # read img + meta info
LoadNifti(keys=["pet_img", "ct_img", "mask_img"]),
Sitk2Numpy(keys=['pet_img', 'ct_img', 'mask_img']),
# user can also add other random transforms
# Prepare for neural network
ConcatModality(keys=['pet_img', 'ct_img']),
AddChanneld(keys=["mask_img"]), # Add channel to the first axis
ToTensord(keys=["image", "mask_img"]),
])
# without data augmentation for validation
val_transforms = Compose(
[ # read img + meta info
LoadNifti(keys=["pet_img", "ct_img", "mask_img"]),
Sitk2Numpy(keys=['pet_img', 'ct_img', 'mask_img']),
# Prepare for neural network
ConcatModality(keys=['pet_img', 'ct_img']),
AddChanneld(keys=["mask_img"]), # Add channel to the first axis
ToTensord(keys=["image", "mask_img"]),
])
##################loading data###############################
device = torch.device("cuda:4" if torch.cuda.is_available() else "cpu")
base_path="/home/lab/hualing/2.5_SUV_dilation" #####the path you put the pre_processed data.
pet_path = base_path + '/' + 'pet_test'
ct_path = base_path + '/' + 'ct_test'
mask_path = base_path + '/' + 'pet_test_mask'
PET_ids = sorted(glob(os.path.join(pet_path, '*pet.nii')))
CT_ids = sorted(glob(os.path.join(ct_path, '*ct.nii')))
MASK_ids = sorted(glob(os.path.join(mask_path, '*mask.nii')))
data_dicts= zip(PET_ids, CT_ids, MASK_ids)
files=list(data_dicts)
train_files = [{"pet_img": PET, "ct_img": CT, 'mask_img': MASK} for PET, CT, MASK in files[:138]]
val_files = [{"pet_img": PET, "ct_img": CT, 'mask_img': MASK} for PET, CT, MASK in files[138:156]]
test_files = [{"pet_img": PET, "ct_img": CT, 'mask_img': MASK} for PET, CT, MASK in files[156:]]
train_ds = monai.data.Dataset(data=train_files,transform=train_transforms)
val_ds = monai.data.Dataset(data=val_files,transform=val_transforms)
test_ds = monai.data.Dataset(data=test_files,transform=val_transforms)
train_loader = DataLoader(
train_ds,
batch_size=3,
shuffle=True,
num_workers=4,
collate_fn=list_data_collate,
pin_memory=torch.cuda.is_available(),)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=4, collate_fn=list_data_collate)
test_loader = DataLoader(test_ds, batch_size=1, num_workers=4, collate_fn=list_data_collate)
###################################### define model#######################################
trained_model_path="./pre-trained_model/best_metric_model_unet.pth" ####path to the pre-trained UNET model
model = UNet_RBF(
dimensions=3, # 3D
in_channels=2,
out_channels=2,
kernel_size=5,
channels=(8,16, 32, 64,128),
strides=(2, 2, 2, 2),
num_res_units=2,).to(device)
model_dict = model.state_dict()
pre_dict = torch.load(trained_model_path)
pre_dict = {k: v for k, v in pre_dict.items() if k in model_dict}
model_dict.update(pre_dict)
model.load_state_dict(model_dict)
####code to make sure only the parameters from ENN are optimized
for name, param in model.named_parameters():
if param.requires_grad==True:
print(name)
####code to make sure only the parameters from ENN are optimized
params = filter(lambda p: p.requires_grad, model.parameters())
optimizer = torch.optim.Adam(params, 1e-2)
dice_metric = monai.metrics.DiceMetric( include_background=False,reduction="mean")
scheduler=torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,'min',patience=10)
loss_function = monai.losses.DiceLoss(include_background=False,softmax=False,squared_pred=True,to_onehot_y=True)
# TODO : generate a learning rate scheduler
val_interval = 1
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = list()
metric_values = list()
writer = SummaryWriter()
post_pred = AsDiscrete(argmax=True, to_onehot=True, n_classes=2)
post_label = AsDiscrete(to_onehot=True, n_classes=2)
############################################# training and validation#############################################
for epoch in range(100):
print("-" * 10)
print(f"epoch {epoch + 1}/{100}")
model.train()
epoch_loss = 0
step = 0
for batch_data in train_loader:
step += 1
inputs, labels = batch_data["image"].to(device), batch_data["mask_img"].to(device)
optimizer.zero_grad()
pm,mass = model(inputs)
loss=loss_function(pm, labels)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
epoch_len = len(train_ds) // train_loader.batch_size
print(f"{step}/{epoch_len}, train_loss: {loss.item():.4f}")
writer.add_scalar("train_loss", loss.item(), epoch_len * epoch + step)
epoch_loss /= step
epoch_loss_values.append(epoch_loss)
print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
scheduler.step(epoch_loss)
if (epoch + 1) % val_interval == 0:
model.eval()
with torch.no_grad():
metric_sum = 0.0
metric_count = 0
for val_data in val_loader:
val_images, val_labels = val_data["image"].to(device), val_data["mask_img"].to(device)
pm,mass = model(val_images)
val_outputs=pm
output=pm
value = dice_metric(y_pred=val_outputs, y=val_labels)
metric_count += len(value)
metric_sum += value.item() * len(value)
metric = metric_sum / metric_count
metric_values.append(metric)
if metric > best_metric:
best_metric = metric
best_metric_epoch = epoch + 1
torch.save(model.state_dict(), "RBF_best_metric_model_segmentation3d_dict.pth")
print("saved new best metric model")
print(
"current epoch: {} current mean dice: {:.4f} best mean dice: {:.4f} at epoch {}".format(
epoch + 1, metric, best_metric, best_metric_epoch
)
)
print(f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}")
writer.close()
Datasets
Models
