# device = "cuda" if torch.cuda.is_available() and use_gpu else "cpu"
device = torch.device('cuda')
batch_size = 1
batch_size_val = 1
num_workers = 0
lr = 0.001
max_epoch = 100
root_dir = '/home/jizong/WorkSpace/LogBarrierCNNs/ACDC-2D-All'
model_dir = 'model'
size_min = 5
size_max = 20
cuda_device = "0"
os.environ["CUDA_VISIBLE_DEVICES"] = cuda_device
color_transform = Colorize()
transform = transforms.Compose([
transforms.ToTensor()
])
mask_transform = transforms.Compose([
transforms.ToTensor()
])
train_set = medicalDataLoader.MedicalImageDataset('train', root_dir, transform=transform, mask_transform=mask_transform,
augment=True, equalize=False)
def main():
## Here we have to split the fully annotated dataset and unannotated dataset
split_ratio = 0.2
random_index = np.random.permutation(len(train_set))
def plot_seg(im):
im = (Colorize()(im).squeeze(0).numpy()).astype(np.uint8)
plt.imshow(im)
plt.show()
def __init__(self, opt):
super(MCDModel, self).__init__(opt)
print('-------------- Networks initializing -------------')
# specify the training losses you want to print out. The program will call base_model.get_current_losses
self.lossNames = [
'loss{}'.format(i)
for i in ['Source', 'TargetClassifier', 'TargetFeature']
]
# specify the training miou you want to print out. The program will call base_model.get_current_losses
self.miouNames = ['miou{}'.format(i) for i in ['Source', 'Target']]
self.miouSource = IouEval(opt.nClass)
self.miouTarget = IouEval(opt.nClass)
# specify the images you want to save/display. The program will call base_model.get_current_visuals
# only image doesn't have prefix
imageNamesSource = ['source', 'sourcePred', 'sourceGnd']
imageNamesTarget = ['target', 'targetPred', 'targetGnd']
self.imageNames = imageNamesSource + imageNamesTarget
# specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks
# naming is by the input domain
self.modelNames = [
'net{}'.format(i)
for i in ['Features', 'Classifier1', 'Classifier2']
]
# load/define networks
# The naming conversion is different from those used in the paper
# Code (paper): G_RGB (G), G_D (F), D_RGB (D_Y), D_D (D_X)
self.netFeatures = self.initNet(
DRNSegBase(model_name=opt.net,
n_class=opt.nClass,
input_ch=opt.inputCh))
self.netClassifier1 = self.initNet(
DRNSegPixelClassifier(n_class=opt.nClass))
self.netClassifier2 = self.initNet(
DRNSegPixelClassifier(n_class=opt.nClass))
self.set_requires_grad(
[self.netFeatures, self.netClassifier1, self.netClassifier2], True)
# define loss functions
self.criterionCycle = torch.nn.L1Loss()
self.criterionSeg = CrossEntropyLoss2d(opt)
self.criterionDis = Distance(opt)
# initialize optimizers
self.optimizerF = getOptimizer(itertools.chain(
self.netFeatures.parameters()),
lr=opt.lr,
momentum=opt.momentum,
opt=opt.opt,
weight_decay=opt.weight_decay)
self.optimizerC = getOptimizer(itertools.chain(
self.netClassifier1.parameters(),
self.netClassifier2.parameters()),
lr=opt.lr,
momentum=opt.momentum,
opt=opt.opt,
weight_decay=opt.weight_decay)
self.optimizers = []
self.optimizers.append(self.optimizerF)
self.optimizers.append(self.optimizerC)
self.colorize = Colorize()
print('--------------------------------------------------')
def __init__(self, opt):
super(CycleMcdModel, self).__init__(opt)
print('-------------- Networks initializing -------------')
self.mode = None
# specify the training losses you want to print out. The program will call base_model.get_current_losses
self.lossNames = [
'loss{}'.format(i) for i in [
'GenA', 'DisA', 'CycleA', 'IdtA', 'DisB', 'GenB', 'CycleB',
'IdtB', 'Supervised', 'UnsupervisedClassifier',
'UnsupervisedFeature'
]
]
self.lossGenA, self.lossDisA, self.lossCycleA, self.lossIdtA = 0, 0, 0, 0
self.lossGenB, self.lossDisB, self.lossCycleB, self.lossIdtB = 0, 0, 0, 0
self.lossSupervised, self.lossUnsupervisedClassifier, self.lossUnsupervisedFeature = 0, 0, 0
# define loss functions
self.criterionGAN = networks.GANLoss(use_lsgan=opt.lsgan).to(
opt.device) # lsgan = True use MSE loss, False use BCE loss
self.criterionCycle = torch.nn.L1Loss()
self.criterionIdt = torch.nn.L1Loss()
self.criterionSeg = CrossEntropyLoss2d(opt) # 2d for each pixels
self.criterionDis = Distance(opt)
# specify the training miou you want to print out. The program will call base_model.get_current_mious
self.miouNames = [
'miou{}'.format(i) for i in
['SupervisedA', 'UnsupervisedA', 'SupervisedB', 'UnsupervisedB']
]
self.miouSupervisedA = IouEval(opt.nClass)
self.miouUnsupervisedA = IouEval(opt.nClass)
self.miouSupervisedB = IouEval(opt.nClass)
self.miouUnsupervisedB = IouEval(opt.nClass)
# specify the images you want to save/display. The program will call base_model.get_current_visuals
# only image doesn't have prefix
imageNamesA = [
'realA', 'fakeA', 'recA', 'idtA', 'supervisedA', 'predSupervisedA',
'gndSupervisedA', 'unsupervisedA', 'predUnsupervisedA',
'gndUnsupervisedA'
]
imageNamesB = [
'realB', 'fakeB', 'recB', 'idtB', 'supervisedB', 'predSupervisedB',
'gndSupervisedB', 'unsupervisedB', 'predUnsupervisedB',
'gndUnsupervisedB'
]
self.imageNames = imageNamesA + imageNamesB
self.realA, self.fakeA, self.recA, self.idtA = None, None, None, None
self.supervisedA, self.predSupervisedA, self.gndSupervisedA = None, None, None
self.unsupervisedA, self.predUnsupervisedA, self.gndUnsupervisedA = None, None, None
self.realB, self.fakeB, self.recB, self.idtB = None, None, None, None
self.supervisedB, self.predSupervisedB, self.gndSupervisedB = None, None, None
self.unsupervisedB, self.predUnsupervisedB, self.gndUnsupervisedB = None, None, None
# specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks
# naming is by the input domain
# Cycle gan model: 'GenA', 'DisA', 'GenB', 'DisB'
# Mcd model : 'Features', 'Classifier1', 'Classifier2'
self.modelNames = [
'net{}'.format(i) for i in [
'GenA', 'DisA', 'GenB', 'DisB', 'Features', 'Classifier1',
'Classifier2'
]
]
# load/define networks
# The naming conversion is different from those used in the paper
# Code (paper): G_RGB (G), G_D (F), D_RGB (D_Y), D_D (D_X)
self.netGenA = networks.define_G(opt.inputCh, opt.inputCh, opt.ngf,
opt.which_model_netG, opt.norm,
opt.dropout, opt.init_type,
opt.init_gain, opt.gpuIds)
self.netDisA = networks.define_D(opt.inputCh, opt.inputCh,
opt.which_model_netD, opt.n_layers_D,
opt.norm, not opt.lsgan,
opt.init_type, opt.init_gain,
opt.gpuIds)
self.netGenB = networks.define_G(opt.inputCh, opt.inputCh, opt.ngf,
opt.which_model_netG, opt.norm,
opt.dropout, opt.init_type,
opt.init_gain, opt.gpuIds)
self.netDisB = networks.define_D(opt.inputCh, opt.inputCh,
opt.which_model_netD, opt.n_layers_D,
opt.norm, not opt.lsgan,
opt.init_type, opt.init_gain,
opt.gpuIds)
self.netFeatures = self.initNet(
DRNSegBase(model_name=opt.segNet,
n_class=opt.nClass,
input_ch=opt.inputCh))
self.netClassifier1 = self.initNet(
DRNSegPixelClassifier(n_class=opt.nClass))
self.netClassifier2 = self.initNet(
DRNSegPixelClassifier(n_class=opt.nClass))
self.set_requires_grad([
self.netGenA, self.netGenB, self.netDisA, self.netDisB,
self.netFeatures, self.netClassifier1, self.netClassifier2
], True)
# define image pool
self.fakeAPool = ImagePool(opt.pool_size)
self.fakeBPool = ImagePool(opt.pool_size)
# initialize optimizers
self.optimizerG = getOptimizer(itertools.chain(
self.netGenA.parameters(), self.netGenB.parameters()),
opt=opt.cycleOpt,
lr=opt.lr,
beta1=opt.beta1,
momentum=opt.momentum,
weight_decay=opt.weight_decay)
self.optimizerD = getOptimizer(itertools.chain(
self.netDisA.parameters(), self.netDisB.parameters()),
opt=opt.cycleOpt,
lr=opt.lr,
beta1=opt.beta1,
momentum=opt.momentum,
weight_decay=opt.weight_decay)
self.optimizerF = getOptimizer(itertools.chain(
self.netFeatures.parameters()),
opt=opt.mcdOpt,
lr=opt.lr,
beta1=opt.beta1,
momentum=opt.momentum,
weight_decay=opt.weight_decay)
self.optimizerC = getOptimizer(itertools.chain(
self.netClassifier1.parameters(),
self.netClassifier2.parameters()),
opt=opt.mcdOpt,
lr=opt.lr,
beta1=opt.beta1,
momentum=opt.momentum,
weight_decay=opt.weight_decay)
self.optimizers = []
self.optimizers.append(self.optimizerG)
self.optimizers.append(self.optimizerD)
self.optimizers.append(self.optimizerF)
self.optimizers.append(self.optimizerC)
self.colorize = Colorize()
print('--------------------------------------------------')