def main():
args = docopt(docstr, version='v0.1')
print(args)
gpu0 = int(args['--gpu0'])
im_path = args['--testIMpath']
gt_path = args['--testGTpath']
model = deeplab_resnet.Res_Deeplab(int(args['--NoLabels']), args['--dgf'],
4, 1e-2)
model.eval().cuda(gpu0)
img_list = open('data/list/val.txt').readlines()
saved_state_dict = torch.load(args['--snapshots'])
model.load_state_dict(saved_state_dict)
save_path = os.path.join('data', args['--exp'])
if not os.path.isdir(save_path):
os.makedirs(save_path)
max_label = int(args['--NoLabels']) - 1 # labels from 0,1, ... 20(for VOC)
hist = np.zeros((max_label + 1, max_label + 1))
for idx, i in enumerate(img_list):
print('{}/{} ...'.format(idx + 1, len(img_list)))
img = cv2.imread(os.path.join(im_path, i[:-1] + '.jpg')).astype(float)
img_original = img.copy() / 255.0
img[:, :, 0] = img[:, :, 0] - 104.008
img[:, :, 1] = img[:, :, 1] - 116.669
img[:, :, 2] = img[:, :, 2] - 122.675
if args['--dgf']:
inputs = [img, img_original]
else:
inputs = [np.zeros((513, 513, 3))]
inputs[0][:img.shape[0], :img.shape[1], :] = img
with torch.no_grad():
output = model(*[
torch.from_numpy(i[np.newaxis, :].transpose(
0, 3, 1, 2)).float().cuda(gpu0) for i in inputs
])
if not args['--dgf']:
interp = nn.Upsample(size=(513, 513),
mode='bilinear',
align_corners=True)
output = interp(output)
output = output[:, :, :img.shape[0], :img.shape[1]]
output = output.cpu().data[0].numpy().transpose(1, 2, 0)
output = np.argmax(output, axis=2)
vis_output = decode_labels(output)
imsave(os.path.join(save_path, i[:-1] + '.png'), vis_output)
gt = cv2.imread(os.path.join(gt_path, i[:-1] + '.png'), 0)
hist += fast_hist(gt.flatten(), output.flatten(), max_label + 1)
miou = np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist))
print("Mean iou = ", np.sum(miou) / len(miou))
def __init__(self, num_classes):
super(PSPNet,self).__init__()
print("initializing model")
init_net=deeplab_resnet.Res_Deeplab()
state=torch.load("models/MS_DeepLab_resnet_trained_VOC.pth")
init_net.load_state_dict(state)
self.resnet=init_net
self.layer5a = PyramidPool(2048, 512, 1)
self.layer5b = PyramidPool(2048, 512, 2)
self.layer5c = PyramidPool(2048, 512, 3)
self.layer5d = PyramidPool(2048, 512, 6)
self.final = nn.Sequential(
nn.Conv2d(4096, 512, 3, padding=1, bias=False),
nn.BatchNorm2d(512, momentum=.95),
nn.ReLU(inplace=True),
nn.Dropout(.1),
nn.Conv2d(512, num_classes, 1),
)
initialize_weights(self.layer5a,self.layer5b,self.layer5c,self.layer5d,self.final)
def main():
args = docopt(docstr, version='v0.1')
print(args)
gpu0 = int(args['--gpu0'])
model = deeplab_resnet.Res_Deeplab(21, True, 4, 1e-2)
model.load_state_dict(torch.load(args['--snapshots']))
model.eval().cuda(gpu0)
im_path = args['--img_path']
img = cv2.imread(im_path).astype(float)
img_original = img.copy() / 255.0
img[:, :, 0] = img[:, :, 0] - 104.008
img[:, :, 1] = img[:, :, 1] - 116.669
img[:, :, 2] = img[:, :, 2] - 122.675
with torch.no_grad():
output = model(*[torch.from_numpy(i[np.newaxis, :].transpose(0, 3, 1, 2)).float().cuda(gpu0) for i in [img, img_original]])
output = output.cpu().data[0].numpy().transpose(1, 2, 0)
output = np.argmax(output, axis=2)
vis_output = decode_labels(output)
output_directory = os.path.dirname(im_path)
output_name = os.path.splitext(os.path.basename(im_path))[0]
save_path = os.path.join(output_directory, '{}_labels.png'.format(output_name))
imsave(save_path, vis_output)
def __init__(self, num_classes):
super(PSPNet, self).__init__()
init_net = deeplab_resnet.Res_Deeplab()
#resnet = models.resnet101(pretrained=True)
state = torch.load("../models/MS_DeepLab_resnet_trained_VOC.pth")
init_net.load_state_dict(state)
self.resnet = init_net
for m in self.modules():
if isinstance(m, nn.Conv2d):
#m.stride = 1
m.requires_grad = False
if isinstance(m, nn.BatchNorm2d):
m.requires_grad = False
self.layer5a = PSPDec(21, 5, 18)
self.layer5b = PSPDec(21, 5, 9)
self.layer5c = PSPDec(21, 5, 6)
self.layer5d = PSPDec(21, 5, 3)
self.final = nn.Sequential(
nn.Conv2d(41, 25, 3, padding=1, bias=False),
nn.BatchNorm2d(25, momentum=.95),
nn.ReLU(inplace=True),
nn.Dropout(.1),
nn.Conv2d(25, num_classes, 1),
)
def convert(img_p, layers):
caffe_model = load_caffe(img_p)
param_provider = CaffeParamProvider(caffe_model)
model = deeplab_resnet.Res_Deeplab(21)
old_dict = model.state_dict()
new_state_dict = OrderedDict()
keys = list(model.state_dict().keys())
for var_name in keys[:]:
data = parse_pth_varnames(param_provider, var_name, layers)
new_state_dict[var_name] = torch.from_numpy(data).float()
model.load_state_dict(new_state_dict)
o = []
def hook(module, input, output):
#print module
o.append(input[0].data.numpy())
model.Scale.conv1.register_forward_hook(hook) #0, data
model.Scale.bn1.register_forward_hook(hook) #1 conv1 out
model.Scale.relu.register_forward_hook(hook) #2 batch norm out
model.Scale.maxpool.register_forward_hook(hook) #3 bn1, relu out
model.Scale.layer1._modules['0'].conv1.register_forward_hook(hook) #4, pool1 out
model.Scale.layer1._modules['1'].conv1.register_forward_hook(hook) #5, res2a out
model.Scale.layer5.conv2d_list._modules['0'].register_forward_hook(hook) #6, res5c out
model.eval()
output = model(Variable(torch.from_numpy(img_p[np.newaxis, :].transpose(0,3,1,2)).float(),volatile=True))
interp = nn.UpsamplingBilinear2d(size=(321, 321))
output_temp = interp(output[3]).cpu().data[0].numpy()
output_temp = output_temp.transpose(1,2,0)
output_temp = np.argmax(output_temp,axis = 2)
#plt.imshow(output_temp)
#plt.show()
dist_(caffe_model.blobs['data'].data,o[0])
dist_(caffe_model.blobs['conv1'].data,o[3])
dist_(caffe_model.blobs['pool1'].data,o[4])
dist_(caffe_model.blobs['res2a'].data,o[5])
dist_(caffe_model.blobs['res5c'].data,o[6])
dist_(caffe_model.blobs['fc1_voc12'].data,output[0].data.numpy())
dist_(caffe_model.blobs['fc1_voc12_res075_interp'].data,output[1].data.numpy())
dist_(caffe_model.blobs['fc1_voc12_res05'].data,output[2].data.numpy())
dist_(caffe_model.blobs['fc_fusion'].data,output[3].data.numpy())
print(('input image shape',img_p[np.newaxis, :].transpose(0,3,1,2).shape))
print('output shapes -')
for a in output:
print((a.data.numpy().shape))
torch.save(model.state_dict(),'data/MS_DeepLab_resnet_trained_VOC.pth')
def convert(img_p, layers):
caffe_model = load_caffe(img_p)
param_provider = CaffeParamProvider(caffe_model)
model = deeplab_resnet.Res_Deeplab(21)
old_dict = model.state_dict()
new_state_dict = OrderedDict()
keys = model.state_dict().keys()
for var_name in keys[:]:
data = parse_pth_varnames(param_provider, var_name, layers)
new_state_dict[var_name] = torch.from_numpy(data).float()
model.load_state_dict(new_state_dict)
o = []
def hook(module, input, output):
#print module
o.append(input[0].data.numpy())
model.Scale.conv1.register_forward_hook(hook) #0, data
model.Scale.bn1.register_forward_hook(hook) #1 conv1 out
model.Scale.relu.register_forward_hook(hook) #2 batch norm out
model.Scale.maxpool.register_forward_hook(hook) #3 bn1, relu out
model.Scale.layer1._modules['0'].conv1.register_forward_hook(
hook) #4, pool1 out
model.Scale.layer1._modules['1'].conv1.register_forward_hook(
hook) #5, res2a out
model.Scale.layer5.conv2d_list._modules['0'].register_forward_hook(
hook) #6, res5c out
model.eval()
output = model(
Variable(torch.from_numpy(img_p[np.newaxis, :].transpose(0, 3, 1,
2)).float(),
volatile=True))
dist_(caffe_model.blobs['data'].data, o[0])
dist_(caffe_model.blobs['conv1'].data, o[3])
dist_(caffe_model.blobs['pool1'].data, o[4])
dist_(caffe_model.blobs['res2a'].data, o[5])
dist_(caffe_model.blobs['res5c'].data, o[6])
dist_(caffe_model.blobs['fc1_voc12'].data, output[3].data.numpy())
print 'input image shape', img_p[np.newaxis, :].transpose(0, 3, 1, 2).shape
print 'output shapes -'
for a in output:
print a.data.numpy().shape
torch.save(model.state_dict(),
'data/MS_DeepLab_resnet_pretrained_COCO_init.pth')
def __init__(self, args):
# import most params
self.batch_size = 1
self.weight_decay = float(args['--wtDecay'])
self.base_lr = float(args['--lr'])
self.epoch_num = int(args['--epoch'])
self.label_num = int(args['--NoLabels'])
self.GPU = int(args['--GPUID'])
self.model = deeplab_resnet.Res_Deeplab(int(args['--NoLabels']))
self.model.load_state_dict(torch.load(args['--pretrain']))
self.model.float()
self.model.eval()
# if torch.cuda.device_count() > 1:
# print("Let's use", torch.cuda.device_count(), "GPUs!")
# # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
# self.model = nn.DataParallel(self.model)
self.model.cuda(self.GPU)
# data set
self.datasets = dataset.create_dataset(args['--ListPath'])
# optimization
self.criterion = nn.CrossEntropyLoss(
ignore_index=255) # use a Classification Cross-Entropy loss
self.optimizer = optim.SGD(
[{
'params': get_1x_lr_params_NOscale(self.model),
'lr': self.base_lr
}, {
'params': get_10x_lr_params(self.model),
'lr': 10 * self.base_lr
}],
lr=self.base_lr,
momentum=0.9,
weight_decay=self.weight_decay)
self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer,
5,
gamma=0.5,
last_epoch=-1)
self.optimizer.zero_grad()
# data writer
self.writer = SummaryWriter(log_dir=args['--summaryPath'])
This generator returns all the parameters for the last layer of the net,
which does the classification of pixel into classes
"""
b = []
b.append(model.Scale.layer5.parameters())
for j in range(len(b)):
for i in b[j]:
yield i
if not os.path.exists('data/snapshots'):
os.makedirs('data/snapshots')
model = deeplab_resnet.Res_Deeplab(int(args['--NoLabels']))
saved_state_dict = torch.load(
'data/MS_DeepLab_resnet_pretrained_COCO_init.pth')
if int(args['--NoLabels']) != 21:
for i in saved_state_dict:
#Scale.layer5.conv2d_list.3.weight
i_parts = i.split('.')
if i_parts[1] == 'layer5':
saved_state_dict[i] = model.state_dict()[i]
model.load_state_dict(saved_state_dict)
max_iter = int(args['--maxIter'])
batch_size = 1
weight_decay = float(args['--wtDecay'])
def main():
opt = parser.parse_args()
print(opt)
trainloader = datasets.init_data.load_data(opt)
model = deeplab_resnet.Res_Deeplab(21)
saved_state_dict = torch.load(
'/data/MS_DeepLab_resnet_pretrained_COCO_init.pth')
model.load_state_dict(saved_state_dict)
max_iter = opt.maxIter
batch_size = 1
weight_decay = opt.wtDecay
base_lr = opt.lr
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD([{
'params': get_1x_lr_params_NOscale(model),
'lr': opt.base_lr
}, {
'params': get_10x_lr_params(model),
'lr': 10 * base_lr
}],
lr=base_lr,
momentum=0.9,
weight_decay=weight_decay)
for i, data in enumerate(trainloader, 0):
images, labels = data
images = Variable(images).cuda()
out = model(images)
loss = loss_calc(out[0], label[0], gpu0)
iter_size = int(args['--iterSize'])
for i in range(len(out) - 1):
loss = loss + loss_calc(out[i + 1], label[i + 1], gpu0)
loss = loss / iter_size
loss.backward()
if iter % 1 == 0:
print 'iter = ', iter, 'of', max_iter, 'completed, loss = ', iter_size * (
loss.data.cpu().numpy())
if iter % iter_size == 0:
optimizer.step()
lr_ = lr_poly(base_lr, iter, max_iter, 0.9)
print '(poly lr policy) learning rate', lr_
optimizer = optim.SGD([{
'params': get_1x_lr_params_NOscale(model),
'lr': lr_
}, {
'params': get_10x_lr_params(model),
'lr': 10 * lr_
}],
lr=lr_,
momentum=0.9,
weight_decay=weight_decay)
optimizer.zero_grad()
if iter % 1000 == 0 and iter != 0:
print 'taking snapshot ...'
torch.save(model.state_dict(),
'data/snapshots/VOC12_scenes_' + str(iter) + '.pth')
def main():
global opt, model
opt = parser.parse_args()
print opt
##########configs########
# isSegReg = False
# isDiceLoss = True # for dice loss
# isSoftmaxLoss = True #for softmax loss
# isResidualEnhancement = False #using ensemble learning to enhance residual learning
# isViewExpansion = True #using dilation to expand receptive filed
# isAdLoss = True #for adverarail loss
# isSpatialDropOut = False
# isFocalLoss = False #we use focal loss to escale training dominated by easy examples
# isSampleImportanceFromAd = False #we set batch sample importance using the loss from adversarial training
# dropoutRate = 0.25
# lambdaAD = 0 # the coefficients before the Adversarial training
# adImportance = 0 # attention from adversarial training
# how2normalize = 3 #1. mu/(max-min); 2. mu/(percent_99 - percent_1); 3. mu/std
# lr = 1e-4 #one of the most important hyper-parameters:
# prefixModelName = 'Segmentor_wdice_wce_lrdce_viewExpansion_1111_'
# prefixPredictedFN = 'preSub_wdice_wce_lrdce_viewExpansion_1111_'
# showTrainLossEvery = 100
# showTestPerformanceEvery = 2000
# decLREvery = 25000 #decrease learning rate every xxx iterations
# saveModelEvery = 2000
# numofIters = 200000
##########configs########
if opt.isSegReg:
netG = ResSegRegNet()
elif opt.isContourLoss:
netG = ResSegContourNet(isRandomConnection=opt.isResidualEnhancement,
isSmallDilation=opt.isViewExpansion,
isSpatialDropOut=opt.isSpatialDropOut,
dropoutRate=opt.dropoutRate)
else:
netG = ResSegNet(isRandomConnection=opt.isResidualEnhancement,
isSmallDilation=opt.isViewExpansion,
isSpatialDropOut=opt.isSpatialDropOut,
dropoutRate=opt.dropoutRate)
# netG =PSPNet(num_classes=4)
netG = deeplab_resnet.Res_Deeplab(NoLabels=4)
#netG.apply(weights_init)
netG = netG.cuda()
netD = Discriminator()
netD.apply(weights_init)
netD.cuda()
params = list(netG.parameters())
print('len of params is ')
print(len(params))
print('size of params is ')
print(params[0].size())
# optimizerG =optim.SGD(netG.parameters(),lr=1e-2)
# optimizerG =optim.Adam(netG.parameters(),lr=opt.lr)
optimizerG = optim.Adam(filter(lambda p: p.requires_grad,
netG.parameters()),
lr=opt.lr)
# optimizerD =optim.SGD(netD.parameters(),lr=1e-4)
optimizerD = optim.Adam(netD.parameters(), lr=opt.lr)
criterion_MSE = nn.MSELoss()
given_weight = torch.FloatTensor([1, 4, 8, 8])
given_weight = given_weight.cuda()
# criterion_NLL2D = nn.NLLLoss2d(weight=given_weight)
criterion_CE2D = CrossEntropy2d(weight=given_weight)
criterion_BCE2D = CrossEntropy2d() #for contours
# criterion_dice = DiceLoss4Organs(organIDs=[1,2,3], organWeights=[1,1,1])
# criterion_dice = WeightedDiceLoss4Organs()
criterion_dice = myWeightedDiceLoss4Organs(organIDs=[0, 1, 2, 3],
organWeights=given_weight)
criterion_focal = myFocalLoss(4, alpha=given_weight, gamma=2)
criterion = nn.BCELoss()
criterion = criterion.cuda()
criterion_dice = criterion_dice.cuda()
criterion_MSE = criterion_MSE.cuda()
criterion_CE2D = criterion_CE2D.cuda()
criterion_BCE2D = criterion_BCE2D.cuda()
criterion_focal = criterion_focal.cuda()
softmax2d = nn.Softmax2d()
# inputs=Variable(torch.randn(1000,1,32,32)) #here should be tensor instead of variable
# targets=Variable(torch.randn(1000,10,1,1)) #here should be tensor instead of variable
# trainset=data_utils.TensorDataset(inputs, targets)
# trainloader = data_utils.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)
# inputs=torch.randn(1000,1,32,32)
# targets=torch.LongTensor(1000)
path_test = '/home/dongnie/warehouse/mrs_data'
path_test = '/shenlab/lab_stor5/dongnie/pelvic'
path_patients_h5 = '/home/dongnie/warehouse/BrainEstimation/brainH5'
path_patients_h5 = '/home/dongnie/warehouse/pelvicSeg/pelvicH5'
path_patients_h5 = '/home/dongnie/warehouse/pelvicSeg/pelvicSegRegH5'
path_patients_h5 = '/home/dongnie/warehouse/pelvicSeg/pelvicSegRegContourBatchH5'
path_patients_h5 = '/shenlab/lab_stor5/dongnie/pelvic/pelvicSeg2D64H5'
# path_patients_h5 = '/shenlab/lab_stor5/dongnie/pelvic/pelvicSegRegH5'
# path_patients_h5 = '/home/dongnie/warehouse/pelvicSeg/pelvicSegRegPartH5/' #only contains 1-15
path_patients_h5_test = '/home/dongnie/warehouse/pelvicSeg/pelvicSegRegContourH5Test'
path_patients_h5_test = '/shenlab/lab_stor5/dongnie/pelvic/pelvicSeg2D64H5Test'
# path_patients_h5_test ='/shenlab/lab_stor5/dongnie/pelvic/pelvicSegRegH5Test'
# batch_size = 10
if opt.isSegReg:
data_generator = Generator_2D_slices_variousKeys(
path_patients_h5,
opt.batchSize,
inputKey='dataMR2D',
outputKey='dataSeg2D',
regKey1='dataBladder2D',
regKey2='dataProstate2D',
regKey3='dataRectum2D')
elif opt.isContourLoss:
data_generator = Generator_2D_slicesV1(path_patients_h5,
opt.batchSize,
inputKey='dataMR2D',
segKey='dataSeg2D',
contourKey='dataContour2D')
else:
data_generator = Generator_2D_slices(path_patients_h5,
opt.batchSize,
inputKey='dataMR2D',
outputKey='dataSeg2D')
data_generator_test = Generator_2D_slices(path_patients_h5_test,
opt.batchSize,
inputKey='dataMR2D',
outputKey='dataSeg2D')
if opt.resume:
if os.path.isfile(opt.resume):
print("=> loading checkpoint '{}'".format(opt.resume))
checkpoint = torch.load(opt.resume)
# opt.start_epoch = checkpoint["epoch"] + 1
# netG.load_state_dict(checkpoint["model"].state_dict())
opt.start_epoch = 4999 + 1
netG.load_state_dict(checkpoint)
else:
print("=> no checkpoint found at '{}'".format(opt.resume))
########### We'd better use dataloader to load a lot of data,and we also should train several epoches###############
running_loss = 0.0
start = time.time()
for iter in range(opt.start_epoch, opt.numofIters + 1):
#print('iter %d'%iter)
if opt.isSegReg:
inputs, labels, regGT1, regGT2, regGT3 = data_generator.next()
elif opt.isContourLoss:
inputs, labels, contours = data_generator.next()
else:
inputs, labels = data_generator.next()
labels = np.squeeze(labels)
labels = zoomImages(labels, rate=opt.zoomRate)
labels = labels.astype(int)
if opt.isContourLoss:
contours = np.squeeze(contours)
contours = contours.astype(int)
contours = torch.from_numpy(contours)
contours = contours.cuda()
contours = Variable(contours)
inputs = torch.from_numpy(inputs)
labels = torch.from_numpy(labels)
inputs = inputs.cuda()
labels = labels.cuda()
#we should consider different data to train
#wrap them into Variable
inputs, labels = Variable(inputs), Variable(labels)
#zero the parameter gradients
#netD.zero_grad()
if opt.isAdLoss:
#forward + backward +optimizer
if opt.isSegReg:
outputG, outputReg1, outputReg2, outputReg3 = netG(inputs)
elif opt.isContourLoss:
outputG, _ = netG(inputs)
else:
outputG = netG(inputs)
outputG = softmax2d(outputG) #batach
# print 'outputG: ',outputG.size(),'labels: ',labels.size()
# print 'outputG: ', outputG.data[0].size()
outputG = outputG.data.max(1)[1]
#outputG = torch.squeeze(outputG) #[N,C,W,H]
labels = labels.unsqueeze(1) #expand the 1st dim
# print 'outputG: ',outputG.size(),'labels: ',labels.size()
outputR = labels.type(torch.FloatTensor).cuda() #output_Real
outputG = Variable(outputG.type(torch.FloatTensor).cuda())
outputD_real = netD(outputR)
# print 'size outputG: ',outputG.unsqueeze(1).size()
outputD_fake = netD(outputG.unsqueeze(1))
## (1) update D network: maximize log(D(x)) + log(1 - D(G(z)))
netD.zero_grad()
batch_size = inputs.size(0)
#print(inputs.size())
#train with real data
# real_label = torch.FloatTensor(batch_size)
# real_label.data.resize_(batch_size).fill_(1)
real_label = torch.ones(batch_size, 1)
real_label = real_label.cuda()
#print(real_label.size())
real_label = Variable(real_label)
#print(outputD_real.size())
loss_real = criterion(outputD_real, real_label)
loss_real.backward()
#train with fake data
fake_label = torch.zeros(batch_size, 1)
# fake_label = torch.FloatTensor(batch_size)
# fake_label.data.resize_(batch_size).fill_(0)
fake_label = fake_label.cuda()
fake_label = Variable(fake_label)
loss_fake = criterion(outputD_fake, fake_label)
loss_fake.backward()
lossD = loss_real + loss_fake
optimizerD.step()
## (2) update G network: minimize the L1/L2 loss, maximize the D(G(x))
#we want to fool the discriminator, thus we pretend the label here to be real. Actually, we can explain from the
#angel of equation (note the max and min difference for generator and discriminator)
if opt.isAdLoss:
if opt.isSegReg:
outputG, outputReg1, outputReg2, outputReg3 = netG(inputs)
elif opt.isContourLoss:
outputG, _ = netG(inputs)
else:
outputG = netG(inputs)
outputG = outputG.data.max(1)[1]
outputG = Variable(outputG.type(torch.FloatTensor).cuda())
# print 'outputG shape, ',outputG.size()
outputD = netD(outputG.unsqueeze(1))
averProb = outputD.data.cpu().mean()
# print 'prob: ',averProb
# adImportance = computeAttentionWeight(averProb)
adImportance = computeSampleAttentionWeight(averProb)
lossG_D = opt.lambdaAD * criterion(
outputD, real_label
) #note, for generator, the label for outputG is real
lossG_D.backward(retain_graph=True)
if opt.isSegReg:
outputG, outputReg1, outputReg2, outputReg3 = netG(inputs)
elif opt.isContourLoss:
outputG, outputContour = netG(inputs)
else:
outputG = netG(
inputs) #here I am not sure whether we should use twice or not
netG.zero_grad()
outputG = outputG[0]
# print 'outputG size: ',outputG.size(),' label size: ',labels.size()
if opt.isFocalLoss:
lossG_focal = criterion_focal(outputG, torch.squeeze(labels))
lossG_focal.backward(retain_graph=True) #compute gradients
if opt.isSoftmaxLoss:
if opt.isSampleImportanceFromAd:
lossG_G = (1 + adImportance) * criterion_CE2D(
outputG, torch.squeeze(labels))
else:
lossG_G = criterion_CE2D(outputG, torch.squeeze(labels))
lossG_G.backward(retain_graph=True) #compute gradients
if opt.isContourLoss:
lossG_contour = criterion_BCE2D(outputContour, contours)
lossG_contour.backward(retain_graph=True)
# criterion_dice(outputG,torch.squeeze(labels))
# print 'hahaN'
if opt.isSegReg:
lossG_Reg1 = criterion_MSE(outputReg1, regGT1)
lossG_Reg2 = criterion_MSE(outputReg2, regGT2)
lossG_Reg3 = criterion_MSE(outputReg3, regGT3)
lossG_Reg = lossG_Reg1 + lossG_Reg2 + lossG_Reg3
lossG_Reg.backward()
if opt.isDiceLoss:
# print 'isDiceLoss line278'
# criterion_dice = myWeightedDiceLoss4Organs(organIDs=[0,1,2,3], organWeights=[1,4,8,6])
if opt.isSampleImportanceFromAd:
loss_dice = (1 + adImportance) * criterion_dice(
outputG, torch.squeeze(labels))
else:
loss_dice = criterion_dice(outputG, torch.squeeze(labels))
# loss_dice = myDiceLoss4Organs(outputG,torch.squeeze(labels)) #succeed
# loss_dice.backward(retain_graph=True) #compute gradients for dice loss
loss_dice.backward() #compute gradients for dice loss
#lossG_D.backward()
#for other losses, we can define the loss function following the pytorch tutorial
optimizerG.step() #update network parameters
# print 'gradients of parameters****************************'
# [x.grad.data for x in netG.parameters()]
# print x.grad.data[0]
# print '****************************'
if opt.isDiceLoss and opt.isSoftmaxLoss and opt.isAdLoss and opt.isSegReg and opt.isFocalLoss:
lossG = opt.lambdaAD * lossG_D + lossG_G + loss_dice.data[
0] + lossG_Reg + lossG_focal
if opt.isDiceLoss and opt.isFocalLoss and opt.isAdLoss and opt.isSegReg:
lossG = opt.lambdaAD * lossG_D + lossG_focal + loss_dice.data[
0] + lossG_Reg
if opt.isDiceLoss and opt.isSoftmaxLoss and opt.isAdLoss and opt.isSegReg:
lossG = opt.lambdaAD * lossG_D + lossG_G + loss_dice.data[
0] + lossG_Reg
elif opt.isSoftmaxLoss and opt.isAdLoss and opt.isSegReg:
lossG = opt.lambdaAD * lossG_D + lossG_G + lossG_Reg
elif opt.isDiceLoss and opt.isAdLoss and opt.isSegReg:
lossG = opt.lambdaAD * lossG_D + loss_dice.data[0] + lossG_Reg
elif opt.isDiceLoss and opt.isSoftmaxLoss and opt.isAdLoss:
lossG = opt.lambdaAD * lossG_D + lossG_G + loss_dice.data[0]
elif opt.isDiceLoss and opt.isFocalLoss and opt.isAdLoss:
lossG = opt.lambdaAD * lossG_D + lossG_focal + loss_dice.data[0]
elif opt.isSoftmaxLoss and opt.isAdLoss:
lossG = opt.lambdaAD * lossG_D + lossG_G
elif opt.isFocalLoss and opt.isAdLoss:
lossG = opt.lambdaAD * lossG_D + lossG_focal
elif opt.isDiceLoss and opt.isAdLoss:
lossG = opt.lambdaAD * lossG_D + loss_dice.data[0]
elif opt.isSoftmaxLoss:
lossG = lossG_G
#print('loss for generator is %f'%lossG.data[0])
#print statistics
running_loss = running_loss + lossG.data[0]
# print 'running_loss is ',running_loss,' type: ',type(running_loss)
# print type(outputD_fake.cpu().data[0].numpy())
if iter % opt.showTrainLossEvery == 0: #print every 2000 mini-batches
print '************************************************'
print 'time now is: ' + time.asctime(time.localtime(time.time()))
if opt.isAdLoss:
print 'the outputD_real for iter {}'.format(
iter), ' is ', outputD_real.cpu().data[0].numpy()[0]
print 'the outputD_fake for iter {}'.format(
iter), ' is ', outputD_fake.cpu().data[0].numpy()[0]
print 'loss for discriminator at iter ', iter, ' is %f' % lossD.data[
0]
# print 'running loss is ',running_loss
print 'average running loss for generator between iter [%d, %d] is: %.3f' % (
iter - 100 + 1, iter, running_loss / 100)
print 'total loss for generator at iter ', iter, ' is %f' % lossG.data[
0]
if opt.isDiceLoss and opt.isSoftmaxLoss and opt.isAdLoss and opt.isSegReg:
print 'lossG_D, lossG_G and loss_dice loss_Reg are %.2f, %.2f and %.2f respectively.' % (
lossG_D.data[0], lossG_G.data[0], loss_dice.data[0],
lossG_Reg.data[0])
elif opt.isDiceLoss and opt.isSoftmaxLoss and opt.isAdLoss:
print 'lossG_D, lossG_G and loss_dice are %.2f, %.2f and %.2f respectively.' % (
lossG_D.data[0], lossG_G.data[0], loss_dice.data[0])
elif opt.isDiceLoss and opt.isFocalLoss and opt.isAdLoss:
print 'lossG_D, lossG_focal and loss_dice are %.2f, %.2f and %.2f respectively.' % (
lossG_D.data[0], lossG_focal.data[0], loss_dice.data[0])
elif opt.isSoftmaxLoss and opt.isAdLoss:
print 'lossG_D and lossG_G are %.2f and %.2f respectively.' % (
lossG_D.data[0], lossG_G.data[0])
elif opt.isFocalLoss and opt.isAdLoss:
print 'lossG_D and lossG_focal are %.2f and %.2f respectively.' % (
lossG_D.data[0], lossG_focal.data[0])
elif opt.isDiceLoss and opt.isAdLoss:
print 'lossG_D and loss_dice are %.2f and %.2f respectively.' % (
lossG_D.data[0], loss_dice.data[0])
elif opt.isSoftmaxLoss:
print ' lossG_G are %.2f respectively.' % (lossG_G.data[0])
if opt.isContourLoss:
print 'lossG_contour is {}'.format(lossG_contour.data[0])
print 'cost time for iter [%d, %d] is %.2f' % (
iter - 100 + 1, iter, time.time() - start)
print '************************************************'
running_loss = 0.0
start = time.time()
if iter % opt.saveModelEvery == 0: #save the model
torch.save(netG.state_dict(), opt.prefixModelName + '%d.pt' % iter)
print 'save model: ' + opt.prefixModelName + '%d.pt' % iter
if iter % opt.decLREvery == 0 and iter > 0:
opt.lr = opt.lr * 0.1
adjust_learning_rate(optimizerG, opt.lr)
print 'now the learning rate is {}'.format(opt.lr)
if iter % opt.showTestPerformanceEvery == 0: #test one subject
# to test on the validation dataset in the format of h5
inputs, labels = data_generator_test.next()
labels = np.squeeze(labels)
labels = zoomImages(labels, rate=opt.zoomRate)
labels = labels.astype(int)
inputs = torch.from_numpy(inputs)
labels = torch.from_numpy(labels)
inputs = inputs.cuda()
labels = labels.cuda()
inputs, labels = Variable(inputs), Variable(labels)
if opt.isSegReg:
outputG, outputReg1, outputReg2, outputReg3 = netG(inputs)
elif opt.isContourLoss:
outputG, _ = netG(inputs)
else:
outputG = netG(
inputs
) #here I am not sure whether we should use twice or not
outputG = outputG[0]
lossG_G = criterion_CE2D(outputG, torch.squeeze(labels))
loss_dice = criterion_dice(outputG, torch.squeeze(labels))
print '.......come to validation stage: iter {}'.format(
iter), '........'
print 'lossG_G and loss_dice are %.2f and %.2f respectively.' % (
lossG_G.data[0], loss_dice.data[0])
####release all the unoccupied memory####
torch.cuda.empty_cache()
mr_test_itk = sitk.ReadImage(
os.path.join(path_test, 'img50_nocrop.nii.gz'))
ct_test_itk = sitk.ReadImage(
os.path.join(path_test, 'img50_label_nie_nocrop.nii.gz'))
mrnp = sitk.GetArrayFromImage(mr_test_itk)
mu = np.mean(mrnp)
ctnp = sitk.GetArrayFromImage(ct_test_itk)
#for training data in pelvicSeg
if opt.how2normalize == 1:
maxV, minV = np.percentile(mrnp, [99, 1])
print 'maxV,', maxV, ' minV, ', minV
mrnp = (mrnp - mu) / (maxV - minV)
print 'unique value: ', np.unique(ctnp)
#for training data in pelvicSeg
if opt.how2normalize == 2:
maxV, minV = np.percentile(mrnp, [99, 1])
print 'maxV,', maxV, ' minV, ', minV
mrnp = (mrnp - mu) / (maxV - minV)
print 'unique value: ', np.unique(ctnp)
#for training data in pelvicSegRegH5
if opt.how2normalize == 3:
std = np.std(mrnp)
mrnp = (mrnp - mu) / std
print 'maxV,', np.ndarray.max(mrnp), ' minV, ', np.ndarray.min(
mrnp)
# full image version with average over the overlapping regions
# ct_estimated = testOneSubject(mrnp,ctnp,[3,168,112],[1,168,112],[1,8,8],netG,'Segmentor_model_%d.pt'%iter)
if opt.how2normalize == 4:
maxV, minV = np.percentile(mrnp, [99.2, 1])
print 'maxV is: ', np.ndarray.max(mrnp)
mrnp[np.where(mrnp > maxV)] = maxV
print 'maxV is: ', np.ndarray.max(mrnp)
mu = np.mean(mrnp)
std = np.std(mrnp)
mrnp = (mrnp - mu) / std
print 'maxV,', np.ndarray.max(mrnp), ' minV, ', np.ndarray.min(
mrnp)
matFA = mrnp
matGT = ctnp
matGT = zoomImages(matGT, rate=opt.zoomRate)
matOut, _ = testOneSubject(matFA, matGT, 4, [3, 64, 64], [1, 9, 9],
[1, 9, 9], netG,
opt.prefixModelName + '%d.pt' % iter)
ct_estimated = np.zeros(
[ctnp.shape[0], ctnp.shape[1], ctnp.shape[2]])
print 'matOut shape: ', matOut.shape
ct_estimated = matOut
tmp_ctnp = ctnp
ctnp = zoomImages(tmp_ctnp, opt.zoomRate)
# ct_estimated[:,y1:y2,x1:x2] = matOut
ct_estimated = np.rint(ct_estimated)
ct_estimated = denoiseImg(ct_estimated,
kernel=np.ones((20, 20, 20)))
diceBladder = dice(ct_estimated, ctnp, 1)
diceProstate = dice(ct_estimated, ctnp, 2)
diceRectumm = dice(ct_estimated, ctnp, 3)
print 'pred: ', ct_estimated.dtype, ' shape: ', ct_estimated.shape
print 'gt: ', ctnp.dtype, ' shape: ', ct_estimated.shape
print 'dice1 = ', diceBladder, ' dice2= ', diceProstate, ' dice3= ', diceRectumm
volout = sitk.GetImageFromArray(ct_estimated)
sitk.WriteImage(
volout, opt.prefixPredictedFN + '{}'.format(iter) + '.nii.gz')
# netG.save_state_dict('Segmentor_model_%d.pt'%iter)
# netD.save_state_dic('Discriminator_model_%d.pt'%iter)
print('Finished Training')
import sys
import numpy as np
import torch
from torch.autograd import Variable
import torch.nn.functional as F
import deeplab_resnet
import torch.nn as nn
import cv2
from config import Config
max_label = Config.class_num
gpu0 = 0
model = deeplab_resnet.Res_Deeplab(max_label)
model.eval()
counter = 0
model.cuda(gpu0)
saved_state_dict = torch.load(
os.path.join(Config.model_path, Config.model_name + "_20000.pth"))
model.load_state_dict(saved_state_dict)
classes = np.array(('background', 'robot_hand', 'inhand_object'))
colormap = [(0, 0, 0), (0.5, 0, 0),
(0, 0.5, 0)] + [(0, 0, 0)] * (256 - len(classes))
colormap = [colormap]
colormap = np.array(colormap) * 255
colormap = colormap.astype(np.uint8)
def main():
args = docopt(docstr, version='v0.1')
print(args)
cudnn.enabled = True
gpu0 = int(args['--gpu0'])
base_lr = float(args['--lr'])
max_iter = int(args['--maxIter'])
iter_size = int(args['--iterSize'])
weight_decay = float(args['--wtDecay'])
if not os.path.exists('data/' + args['--snapshots']):
os.makedirs('data/' + args['--snapshots'])
model = deeplab_resnet.Res_Deeplab(int(args['--NoLabels']), args['--dgf'],
4, 1e-2)
if args['--ft']:
saved_state_dict = torch.load(args['--ft_model_path'])
else:
saved_state_dict = torch.load(
'data/MS_DeepLab_resnet_pretrained_COCO_init.pth')
model_dict = model.state_dict()
model_dict.update(saved_state_dict)
model.load_state_dict(model_dict)
model.float().eval().cuda(gpu0)
optimizer = optim.SGD([{
'params': get_1x_lr_params_NOscale(model),
'lr': base_lr
}, {
'params': get_10x_lr_params(args, model),
'lr': 10 * base_lr
}],
lr=base_lr,
momentum=0.9,
weight_decay=weight_decay)
optimizer.zero_grad()
img_list = read_file(args['--LISTpath'])
data_list = []
# make list for 10 epocs, though we will only use the first max_iter*batch_size entries of this list
for i in range(10):
np.random.shuffle(img_list)
data_list.extend(img_list)
data_gen = chunker(data_list, 1)
for iter in range(max_iter + 1):
inputs, label = get_data_from_chunk_v2(args, next(data_gen))
inputs = [Variable(input).cuda(gpu0) for input in inputs]
loss = loss_calc(model(*inputs), label, gpu0) / iter_size
loss.backward()
if iter % 1 == 0:
print('iter = ', iter, 'of', max_iter, 'completed, loss = ',
iter_size * (loss.data.cpu().numpy()))
if iter % iter_size == 0:
optimizer.step()
lr_ = lr_poly(base_lr, iter, max_iter, 0.9)
print('(poly lr policy) learning rate', lr_)
optimizer = optim.SGD([{
'params': get_1x_lr_params_NOscale(model),
'lr': lr_
}, {
'params': get_10x_lr_params(args, model),
'lr': 10 * lr_
}],
lr=lr_,
momentum=0.9,
weight_decay=weight_decay)
optimizer.zero_grad()
if iter % 1000 == 0 and iter != 0:
print('taking snapshot ...')
torch.save(
model.state_dict(), 'data/' + args['--snapshots'] +
'/VOC12_scenes_' + str(iter) + '.pth')
"""
This generator returns all the parameters for the last layer of the net,
which does the classification of pixel into classes
"""
b = []
b.append(model.Scale.layer5.parameters())
for j in range(len(b)):
for i in b[j]:
yield i
if not os.path.exists(Config.snapshot_path):
os.makedirs(Config.snapshot_path)
model = deeplab_resnet.Res_Deeplab(args.NoLabels)
saved_state_dict = torch.load(os.path.join(Config.model_path, "MS_DeepLab_resnet_pretrained_COCO_init.pth"))
if args.NoLabels != 21:
for i in saved_state_dict:
#Scale.layer5.conv2d_list.3.weight
i_parts = i.split('.')
if i_parts[1]=='layer5':
saved_state_dict[i] = model.state_dict()[i]
model.load_state_dict(saved_state_dict)
max_iter = args.maxIter
batch_size = 1
