def main(): opt = TestOptions() args = opt.initialize() if not os.path.exists(args.save): os.makedirs(args.save) model = CreateModel(args) model.eval() model.cuda() targetloader = CreateTrgDataLoader(args) for index, batch in enumerate(targetloader): if index % 100 == 0: print '%d processd' % index image, _, name = batch output = model(Variable(image).cuda()) output = nn.functional.upsample(output, (1024, 2048), mode='bilinear', align_corners=True).cpu().data[0].numpy() output = output.transpose(1,2,0) output_nomask = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) output_col = colorize_mask(output_nomask) output_nomask = Image.fromarray(output_nomask) name = name[0].split('/')[-1] output_nomask.save('%s/%s' % (args.save, name)) output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0])) compute_mIoU(args.gt_dir, args.save, args.devkit_dir, args.restore_from)
def main(): testing_entropy = 0 args = parse_args() if not os.path.exists(args.save): os.makedirs(args.save) args.init_weights = root_base+'/snapshots/' + args.model_name model = CreateModel(args) model.eval() model.cuda() targetloader = CreateTrgDataLoader(args) for index, batch in tqdm.tqdm(enumerate(targetloader)): image, _, name,_ = batch output = model(Variable(image).cuda()) output = nn.functional.softmax(output, dim=1) testing_entropy += self_entropy(roll_axis(output.cpu().data[0].numpy())) output = nn.functional.upsample(output, (1024, 2048), mode='bilinear', align_corners=True).cpu().data[0].numpy() #output = np.multiply(output,priors) output = output.transpose(1,2,0) output_nomask = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) output_col = colorize_mask(output_nomask) output_nomask = Image.fromarray(output_nomask) name = name[0].split('/')[-1] output_nomask.save('%s/%s' % (args.save, name)) output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0])) mIou_ = compute_mIoU(args.gt_dir, args.save, args.devkit_dir, '')
def main(): opt = TestOptions() args = opt.initialize() if not os.path.exists(args.save): os.makedirs(args.save) maxiou = 0 root = args.restore_from for i in range(1, 51): #args.restore_from = './snapshots/gta2city_vgg/gta5_{0:d}'.format(i*2000) args.restore_from = root + '_{0:d}'.format(i * 2000) model = CreateModel(args) print(args.restore_from) model.eval() model.cuda() targetloader = CreateTrgDataLoader(args) for index, batch in enumerate(targetloader): if index % 100 == 0: print '%d processd' % index image, _, name = batch output1, output2 = model(Variable(image).cuda()) output = output1 + output2 output = nn.functional.softmax(output, dim=1) output = nn.functional.upsample( output, (1024, 2048), mode='bilinear', align_corners=True).cpu().data[0].numpy() output = output.transpose(1, 2, 0) output_nomask = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) output_col = colorize_mask(output_nomask) output_nomask = Image.fromarray(output_nomask) name = name[0].split('/')[-1] output_nomask.save('%s/%s' % (args.save, name)) output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0])) iou = compute_mIoU(args.gt_dir, args.save, args.devkit_dir, args.restore_from) #print(save_path) if iou > maxiou: maxiou = iou best_model = args.restore_from print('The best model is {:s}, the best IOU is {:f}'.format( best_model, maxiou))
def main(): # torch.manual_seed(1234) # torch.cuda.manual_seed(1234) opt = TrainOptions() args = opt.initialize() _t = {'iter time' : Timer()} model_name = args.source + '_to_' + args.target if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) os.makedirs(os.path.join(args.snapshot_dir, 'logs')) opt.print_options(args) sourceloader, targetloader = CreateSrcDataLoader(args), CreateTrgDataLoader(args) targetloader_iter, sourceloader_iter = iter(targetloader), iter(sourceloader) model, optimizer = CreateModel(args) model_D, optimizer_D = CreateDiscriminator(args) start_iter = 0 if args.restore_from is not None: start_iter = int(args.restore_from.rsplit('/', 1)[1].rsplit('_')[1]) train_writer = tensorboardX.SummaryWriter(os.path.join(args.snapshot_dir, "logs", model_name)) bce_loss = torch.nn.BCEWithLogitsLoss() cent_loss=ConditionalEntropyLoss() cudnn.enabled = True cudnn.benchmark = True model.train() model.cuda() model_D.train() model_D.cuda() loss = ['loss_seg_src', 'loss_seg_trg', 'loss_D_trg_fake', 'loss_D_src_real', 'loss_D_trg_real'] _t['iter time'].tic() pbar = tqdm(range(start_iter,args.num_steps_stop)) #for i in range(start_iter, args.num_steps): for i in pbar: model.adjust_learning_rate(args, optimizer, i) model_D.adjust_learning_rate(args, optimizer_D, i) optimizer.zero_grad() optimizer_D.zero_grad() for param in model_D.parameters(): param.requires_grad = False src_img, src_lbl, _, _ = sourceloader_iter.next() src_img, src_lbl = Variable(src_img).cuda(), Variable(src_lbl.long()).cuda() src_seg_score = model(src_img) loss_seg_src = CrossEntropy2d(src_seg_score, src_lbl) #loss_seg_src = model.loss loss_seg_src.backward() if args.data_label_folder_target is not None: trg_img, trg_lbl, _, _ = targetloader_iter.next() trg_img, trg_lbl = Variable(trg_img).cuda(), Variable(trg_lbl.long()).cuda() trg_seg_score = model(trg_img) loss_seg_trg = model.loss else: trg_img, _, name = targetloader_iter.next() trg_img = Variable(trg_img).cuda() trg_seg_score = model(trg_img) #ipdb.set_trace() loss_seg_trg= cent_loss(trg_seg_score) #loss_seg_trg= entropy_loss(F.softmax(trg_seg_score)) #loss_seg_trg = 0 outD_trg = model_D(F.softmax(trg_seg_score)) loss_D_trg_fake = bce_loss(outD_trg, Variable(torch.FloatTensor(outD_trg.data.size()).fill_(0)).cuda()) #loss_D_trg_fake = model_D.loss loss_trg = args.lambda_adv_target * (loss_D_trg_fake + loss_seg_trg) loss_trg.backward() for param in model_D.parameters(): param.requires_grad = True src_seg_score, trg_seg_score = src_seg_score.detach(), trg_seg_score.detach() outD_src = model_D(F.softmax(src_seg_score)) loss_D_src_real = bce_loss(outD_src, Variable(torch.FloatTensor(outD_src.data.size()).fill_(0)).cuda())/ 2 #loss_D_src_real = model_D.loss / 2 loss_D_src_real.backward() outD_trg = model_D(F.softmax(trg_seg_score)) loss_D_trg_real = bce_loss(outD_trg, Variable(torch.FloatTensor(outD_trg.data.size()).fill_(1)).cuda())/ 2 #loss_D_trg_real = model_D.loss / 2 loss_D_trg_real.backward() d_loss=loss_D_src_real.data+ loss_D_trg_real.data optimizer.step() optimizer_D.step() for m in loss: train_writer.add_scalar(m, eval(m), i+1) if (i+1) % args.save_pred_every == 0: print 'taking snapshot ...' torch.save(model.state_dict(), os.path.join(args.snapshot_dir, '%s_' %(args.source) +str(i+1)+'.pth' )) torch.save(model_D.state_dict(), os.path.join(args.snapshot_dir, '%s_' %(args.source) +str(i+1)+'_D.pth' )) if (i+1) % args.print_freq == 0: _t['iter time'].toc(average=False) print '[it %d][src seg loss %.4f][adv loss %.4f][d loss %.4f][lr %.4f][%.2fs]' % \ (i + 1, loss_seg_src.data, loss_D_trg_fake.data,d_loss,optimizer.param_groups[0]['lr']*10000, _t['iter time'].diff) if i + 1 > args.num_steps_stop: print 'finish training' break _t['iter time'].tic()
def main(): # torch.manual_seed(1234) # torch.cuda.manual_seed(1234) opt = TrainOptions() args = opt.initialize() _t = {'iter time': Timer()} model_name = args.source + '_to_' + args.target if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) os.makedirs(os.path.join(args.snapshot_dir, 'logs')) opt.print_options(args) sourceloader, targetloader = CreateSrcDataLoader( args), CreateTrgDataLoader(args) targetloader_iter, sourceloader_iter = iter(targetloader), iter( sourceloader) model, optimizer = CreateModel(args) model_D, optimizer_D = CreateDiscriminator(args) start_iter = 0 if args.restore_from is not None: start_iter = int(args.restore_from.rsplit('/', 1)[1].rsplit('_')[1]) train_writer = tensorboardX.SummaryWriter( os.path.join(args.snapshot_dir, "logs", model_name)) bce_loss = torch.nn.BCEWithLogitsLoss() l1_loss = torch.nn.L1Loss() cos_loss = torch.nn.CosineSimilarity(dim=0, eps=1e-06) cudnn.enabled = True cudnn.benchmark = True model.train() model.cuda() model_D.train() model_D.cuda() loss = [ 'loss_seg_src', 'loss_seg_trg', 'loss_D_trg_fake', 'loss_D_src_real', 'loss_D_trg_real' ] _t['iter time'].tic() pbar = tqdm(range(start_iter, args.num_steps_stop)) #for i in range(start_iter, args.num_steps): for i in pbar: model.adjust_learning_rate(args, optimizer, i) model_D.adjust_learning_rate(args, optimizer_D, i) optimizer.zero_grad() optimizer_D.zero_grad() for param in model_D.parameters(): param.requires_grad = False src_img, src_lbl, _, _ = sourceloader_iter.next() src_img, src_lbl = Variable(src_img).cuda(), Variable( src_lbl.long()).cuda() src_seg_score, src_seg_score2 = model(src_img) loss_seg_src1 = CrossEntropy2d(src_seg_score, src_lbl) loss_seg_src2 = CrossEntropy2d(src_seg_score2, src_lbl) loss_seg_src = loss_seg_src1 + loss_seg_src2 loss_seg_src.backward() if args.data_label_folder_target is not None: trg_img, trg_lbl, _, _ = targetloader_iter.next() trg_img, trg_lbl = Variable(trg_img).cuda(), Variable( trg_lbl.long()).cuda() trg_seg_score = model(trg_img) loss_seg_trg = model.loss else: trg_img, _, name = targetloader_iter.next() trg_img = Variable(trg_img).cuda() trg_seg_score, trg_seg_score2 = model(trg_img) loss_seg_trg = 0 outD_trg = model_D(F.softmax(trg_seg_score)) outD_trg2 = model_D(F.softmax(trg_seg_score2)) loss_D_trg_fake1 = bce_loss( outD_trg, Variable(torch.FloatTensor(outD_trg.data.size()).fill_(0)).cuda()) loss_D_trg_fake2 = bce_loss( outD_trg2, Variable(torch.FloatTensor(outD_trg2.data.size()).fill_(0)).cuda()) loss_D_trg_fake = loss_D_trg_fake1 + loss_D_trg_fake2 loss_agree = l1_loss(F.softmax(trg_seg_score), F.softmax(trg_seg_score2)) loss_trg = args.lambda_adv_target * loss_D_trg_fake + loss_seg_trg + loss_agree loss_trg.backward() #Weight Discrepancy Loss W5 = None W6 = None if args.model == 'DeepLab2': for (w5, w6) in zip(model.layer5.parameters(), model.layer6.parameters()): if W5 is None and W6 is None: W5 = w5.view(-1) W6 = w6.view(-1) else: W5 = torch.cat((W5, w5.view(-1)), 0) W6 = torch.cat((W6, w6.view(-1)), 0) #ipdb.set_trace() #loss_weight = (torch.matmul(W5, W6) / (torch.norm(W5) * torch.norm(W6)) + 1) # +1 is for a positive loss # loss_weight = loss_weight * damping * 2 loss_weight = args.weight_div * (cos_loss(W5, W6) + 1) loss_weight.backward() for param in model_D.parameters(): param.requires_grad = True src_seg_score, trg_seg_score = src_seg_score.detach( ), trg_seg_score.detach() src_seg_score2, trg_seg_score2 = src_seg_score2.detach( ), trg_seg_score2.detach() outD_src = model_D(F.softmax(src_seg_score)) loss_D_src_real1 = bce_loss( outD_src, Variable(torch.FloatTensor( outD_src.data.size()).fill_(0)).cuda()) / 2 outD_src2 = model_D(F.softmax(src_seg_score2)) loss_D_src_real2 = bce_loss( outD_src2, Variable(torch.FloatTensor( outD_src2.data.size()).fill_(0)).cuda()) / 2 loss_D_src_real = loss_D_src_real1 + loss_D_src_real2 loss_D_src_real.backward() outD_trg = model_D(F.softmax(trg_seg_score)) loss_D_trg_real1 = bce_loss( outD_trg, Variable(torch.FloatTensor( outD_trg.data.size()).fill_(1)).cuda()) / 2 outD_trg2 = model_D(F.softmax(trg_seg_score2)) loss_D_trg_real2 = bce_loss( outD_trg2, Variable(torch.FloatTensor( outD_trg2.data.size()).fill_(1)).cuda()) / 2 loss_D_trg_real = loss_D_trg_real1 + loss_D_trg_real2 loss_D_trg_real.backward() d_loss = loss_D_src_real.data + loss_D_trg_real.data optimizer.step() optimizer_D.step() for m in loss: train_writer.add_scalar(m, eval(m), i + 1) if (i + 1) % args.save_pred_every == 0: print 'taking snapshot ...' torch.save( model.state_dict(), os.path.join(args.snapshot_dir, '%s_' % (args.source) + str(i + 1) + '.pth')) torch.save( model_D.state_dict(), os.path.join(args.snapshot_dir, '%s_' % (args.source) + str(i + 1) + '_D.pth')) if (i + 1) % args.print_freq == 0: _t['iter time'].toc(average=False) print '[it %d][src seg loss %.4f][adv loss %.4f][d loss %.4f][agree loss %.4f][div loss %.4f][lr %.4f][%.2fs]' % \ (i + 1, loss_seg_src.data, loss_D_trg_fake.data,d_loss,loss_agree.data,loss_weight.data, optimizer.param_groups[0]['lr']*10000, _t['iter time'].diff) if i + 1 > args.num_steps_stop: print 'finish training' break _t['iter time'].tic()
def main(): opt = TrainOptions() args = opt.initialize() _t = {'iter time': Timer()} model_name = args.source + '_to_' + args.target if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) os.makedirs(os.path.join(args.snapshot_dir, 'logs')) opt.print_options(args) sourceloader, targetloader = CreateSrcDataLoader( args), CreateTrgDataLoader(args) targetloader_iter, sourceloader_iter = iter(targetloader), iter( sourceloader) model, optimizer = CreateModel(args) model_D1, optimizer_D1 = CreateDiscriminator(args, 1) model_D2, optimizer_D2 = CreateDiscriminator(args, 2) start_iter = 0 if args.restore_from is not None: start_iter = int(args.restore_from.rsplit('/', 1)[1].rsplit('_')[1]) #train_writer = tensorboardX.SummaryWriter(os.path.join(args.snapshot_dir, "logs", model_name)) bce_loss = torch.nn.BCEWithLogitsLoss() interp_target = nn.Upsample(size=(1024, 1024), mode='bilinear', align_corners=True) interp_source = nn.Upsample(size=(1024, 1024), mode='bilinear', align_corners=True) cudnn.enabled = True cudnn.benchmark = True model.train() model.cuda() model_D1.train() model_D1.cuda() model_D2.train() model_D2.cuda() weight_loss = WeightedBCEWithLogitsLoss() loss = [ 'loss_seg_src', 'loss_seg_trg', 'loss_D_trg_fake', 'loss_D_src_real', 'loss_D_trg_real' ] _t['iter time'].tic() load_selected_samples = args.load_selected_samples if load_selected_samples is None: total_num = args.total_number else: clean_ids = [i_id.strip() for i_id in open(load_selected_samples)] total_num = len(clean_ids) remember_rate = args.remember_rate loss_list, name_list, dice_list = [], [], [] print('total_num:', total_num) predict_sum_disc = 0 predict_sum_cup = 0 noise_sum_disc = 0 noise_sum_cup = 0 threshold = 0.4 for i in range(start_iter, start_iter + total_num): model.adjust_learning_rate(args, optimizer, i) model_D1.adjust_learning_rate(args, optimizer_D1, i) model_D2.adjust_learning_rate(args, optimizer_D2, i) optimizer.zero_grad() optimizer_D1.zero_grad() optimizer_D2.zero_grad() ##train G for param in model_D1.parameters(): param.requires_grad = False for param in model_D2.parameters(): param.requires_grad = False #import pdb;pdb.set_trace() try: src_img, src_lbl, weight_map, name = sourceloader_iter.next() except StopIteration: sourceloader_iter = iter(sourceloader) src_img, src_lbl, weight_map, name = sourceloader_iter.next() src_img, src_lbl, weight_map = Variable(src_img).cuda(), Variable( src_lbl.long()).cuda(), Variable(weight_map.long()).cuda() src_seg_score1, src_seg_score2, src_seg_score3, src_seg_score4 = model( src_img, lbl=src_lbl, weight=None) loss_seg_src = model.loss loss_seg_src.backward() loss_list.append(loss_seg_src) name_list.append(name) print(i, name) #import pdb;pdb.set_trace() output = nn.functional.softmax(src_seg_score2, dim=1) output = nn.functional.upsample( output, (2056, 2124), mode='bilinear', align_corners=True).cpu().data[0].numpy() output = output.transpose(1, 2, 0) # (1634,1634,3) output_mask = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) # (1644,1634) unique:[0,1,2] predict_disc_dice, predict_cup_dice = calculate_dice( args.data_dir, output_mask, name) print('===>predict disc_dice:' + str(round(predict_disc_dice, 3)) + '\t' + 'cup_dice:' + str(round(predict_cup_dice, 3))) #import pdb;pdb.set_trace() label = torch.unsqueeze(src_lbl, 0) label = nn.functional.upsample( label.float(), size=(2056, 2124), mode='bilinear', align_corners=True).cpu().data[0].numpy().squeeze() noise_disc_dice, noise_cup_dice = calculate_dice( args.data_dir, label, name) print('===>noise-included disc_dice:' + str(round(noise_disc_dice, 3)) + '\t' + 'cup_dice:' + str(round(noise_cup_dice, 3))) predict_sum_disc += predict_disc_dice predict_sum_cup += predict_cup_dice noise_sum_disc += noise_disc_dice noise_sum_cup += noise_cup_dice #import pdb;pdb.set_trace() if load_selected_samples is not None: if (2 - predict_disc_dice - predict_cup_dice) > threshold: src_seg_score1, src_seg_score2, src_seg_score3, src_seg_score4 = model( src_img, lbl=src_lbl, weight=weight_map) weightloss = 0.05 * model.loss print('weightloss:', weightloss) weightloss.backward() if args.data_label_folder_target is not None: trg_img, trg_lbl, _, _ = targetloader_iter.next() trg_img, trg_lbl = Variable(trg_img).cuda(), Variable( trg_lbl.long()).cuda() trg_seg_score1, trg_seg_score2, trg_seg_score3, trg_seg_score4 = model( trg_img, lbl=trg_lbl) loss_seg_trg = model.loss else: trg_img, _, _ = targetloader_iter.next() trg_img = Variable(trg_img).cuda() trg_seg_score1, trg_seg_score2, trg_seg_score3, trg_seg_score4 = model( trg_img) loss_seg_trg = 0 outD1_trg = model_D1(F.softmax(trg_seg_score1), 0) outD2_trg = model_D2(F.softmax(trg_seg_score2), 0) #import pdb;pdb.set_trace() outD1_trg = interp_target(outD1_trg) outD2_trg = interp_target(outD2_trg) if i > 9001: #import pdb;pdb.set_trace() weight_map1 = prob_2_entropy(F.softmax(trg_seg_score1)) weight_map2 = prob_2_entropy(F.softmax(trg_seg_score2)) loss_D1_trg_fake = weight_loss( outD1_trg, Variable(torch.FloatTensor( outD1_trg.data.size()).fill_(0)).cuda(), weight_map1, 0.3, 1) loss_D2_trg_fake = weight_loss( outD2_trg, Variable(torch.FloatTensor( outD2_trg.data.size()).fill_(0)).cuda(), weight_map2, 0.3, 1) else: loss_D1_trg_fake = model_D1.loss loss_D2_trg_fake = model_D2.loss #loss_D_trg_fake = model_D1.loss*0.2 + model_D2.loss loss_D_trg_fake = loss_D1_trg_fake * 0.2 + loss_D2_trg_fake loss_trg = args.lambda_adv_target * loss_D_trg_fake + loss_seg_trg loss_trg.backward() ###train D for param in model_D1.parameters(): param.requires_grad = True for param in model_D2.parameters(): param.requires_grad = True src_seg_score1, src_seg_score2, src_seg_score3, src_seg_score4, trg_seg_score1, trg_seg_score2, trg_seg_score3, trg_seg_score4 = src_seg_score1.detach( ), src_seg_score2.detach(), src_seg_score3.detach( ), src_seg_score4.detach(), trg_seg_score1.detach( ), trg_seg_score2.detach(), trg_seg_score3.detach( ), trg_seg_score4.detach() outD1_src = model_D1(F.softmax(src_seg_score1), 0) outD2_src = model_D2(F.softmax(src_seg_score2), 0) loss_D1_src_real = model_D1.loss / 2 loss_D1_src_real.backward() loss_D2_src_real = model_D2.loss / 2 loss_D2_src_real.backward() loss_D_src_real = loss_D1_src_real + loss_D2_src_real outD1_trg = model_D1(F.softmax(trg_seg_score1), 1) outD2_trg = model_D2(F.softmax(trg_seg_score2), 1) outD1_trg = interp_target(outD1_trg) outD2_trg = interp_target(outD2_trg) if i > 9001: weight_map1 = prob_2_entropy(F.softmax(trg_seg_score1)) weight_map2 = prob_2_entropy(F.softmax(trg_seg_score2)) loss_D1_trg_real = weight_loss( outD1_trg, Variable(torch.FloatTensor( outD1_trg.data.size()).fill_(1)).cuda(), weight_map1, 0.3, 1) / 2 loss_D2_trg_real = weight_loss( outD2_trg, Variable(torch.FloatTensor( outD2_trg.data.size()).fill_(1)).cuda(), weight_map2, 0.3, 1) / 2 else: loss_D1_trg_real = model_D1.loss / 2 loss_D2_trg_real = model_D2.loss / 2 loss_D1_trg_real.backward() loss_D2_trg_real.backward() loss_D_trg_real = loss_D1_trg_real + loss_D2_trg_real optimizer.step() optimizer_D1.step() optimizer_D2.step() # for m in loss: # train_writer.add_scalar(m, eval(m), i+1) if (i + 1 == start_iter + total_num) and load_selected_samples is not None: print('taking snapshot ', args.snapshot_dir, args.source + '_' + str(total_num)) torch.save( model.state_dict(), os.path.join(args.snapshot_dir, '%s_' % (args.source) + str(total_num) + '.pth')) torch.save( model_D1.state_dict(), os.path.join( args.snapshot_dir, '%s_' % (args.source) + str(total_num) + '_D1.pth')) torch.save( model_D2.state_dict(), os.path.join( args.snapshot_dir, '%s_' % (args.source) + str(total_num) + '_D2.pth')) if (i + 1) % args.print_freq == 0: _t['iter time'].toc(average=False) print ('[it %d][src seg loss %.4f][trg loss %.4f][trg seg loss %.4f][lr %.4f][%.2fs]' % \ (i + 1, loss_seg_src.data,loss_trg.data, loss_seg_trg.data,optimizer.param_groups[0]['lr']*10000, _t['iter time'].diff)) if i + 1 > args.num_steps_stop: print('finish training') break _t['iter time'].tic() if i + 1 == start_iter + total_num and load_selected_samples is None: ind_sorted = np.argsort(loss_list) loss_sorted = np.array(loss_list)[ind_sorted] num_remember = int(remember_rate * len(loss_sorted)) clean_ind_update = ind_sorted[:num_remember] clean_name_update = np.array(name_list)[clean_ind_update] #import pdb;pdb.set_trace() #dice = np.array(dice_list)[clean_ind_update] noise_ind_sorted = np.argsort(-np.array(loss_list)) noise_loss_sorted = np.array(loss_list)[noise_ind_sorted] noise_num_remember = int(remember_rate * len(noise_loss_sorted)) noise_ind_update = noise_ind_sorted[:noise_num_remember] noise_name_update = np.array(name_list)[noise_ind_update] with open(os.path.join(args.save_selected_samples), "w") as f: for i in range(len(clean_name_update)): #f.write(str(clean_name_update[i][0]) +'\t'+ str(dice[:,0][i]) + '\t'+ str(dice[:,1][i]) +'\n') f.write(str(clean_name_update[i][0]) + '\n') with open(os.path.join(args.noise_selected_samples), "w") as g: for j in range(len(noise_name_update)): g.write(str(noise_name_update[j][0]) + '\n') print(args.save_selected_samples, 'Sample selection finished!') break print('\n predict disc_coef = {0:.4f}, cup_coef = {1:.4f}'.format( predict_sum_disc / total_num, predict_sum_cup / total_num)) print('\n noise-included disc_coef = {0:.4f}, cup_coef = {1:.4f}'.format( noise_sum_disc / total_num, noise_sum_cup / total_num))
def main(): opt = TrainOptions() # loading train options(arg parser) args = opt.initialize() # get arguments os.environ["CUDA_VISIBLE_DEVICES"] = args.GPU _t = {'iter time': Timer()} model_name = args.source + '_to_' + args.target if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) os.makedirs(os.path.join(args.snapshot_dir, 'logs')) opt.print_options(args) # print set options sourceloader, targetloader = CreateSrcDataLoader( args), CreateTrgDataLoader(args) sourceloader_iter, targetloader_iter = iter(sourceloader), iter( targetloader) model, optimizer = CreateModel(args) start_iter = 0 if args.restore_from is not None: start_iter = int(args.restore_from.rsplit('/', 1)[1].rsplit('_')[1]) cudnn.enabled = True cudnn.benchmark = True model.train() model.cuda() # losses to log loss = ['loss_seg_src', 'loss_seg_trg'] loss_train = 0.0 loss_val = 0.0 loss_train_list = [] loss_val_list = [] mean_img = torch.zeros(1, 1) class_weights = Variable(CS_weights).cuda() _t['iter time'].tic() for i in range(start_iter, args.num_steps): model.adjust_learning_rate(args, optimizer, i) # adjust learning rate optimizer.zero_grad() # zero grad src_img, src_lbl, _, _ = sourceloader_iter.next() # new batch source trg_img, trg_lbl, _, _ = targetloader_iter.next() # new batch target scr_img_copy = src_img.clone() if mean_img.shape[-1] < 2: B, C, H, W = src_img.shape mean_img = IMG_MEAN.repeat(B, 1, H, W) #-------------------------------------------------------------------# # 1. source to target, target to target src_in_trg = FDA_source_to_target(src_img, trg_img, L=args.LB) # src_lbl trg_in_trg = trg_img # 2. subtract mean src_img = src_in_trg.clone() - mean_img # src, src_lbl trg_img = trg_in_trg.clone() - mean_img # trg, trg_lbl #-------------------------------------------------------------------# # evaluate and update params ##### src_img, src_lbl = Variable(src_img).cuda(), Variable( src_lbl.long()).cuda() # to gpu src_seg_score = model(src_img, lbl=src_lbl, weight=class_weights, ita=args.ita) # forward pass loss_seg_src = model.loss_seg # get loss loss_ent_src = model.loss_ent # get target loss, only entropy for backpro trg_img, trg_lbl = Variable(trg_img).cuda(), Variable( trg_lbl.long()).cuda() # to gpu trg_seg_score = model(trg_img, lbl=trg_lbl, weight=class_weights, ita=args.ita) # forward pass loss_seg_trg = model.loss_seg # get loss loss_ent_trg = model.loss_ent triger_ent = 0.0 if i > args.switch2entropy: triger_ent = 1.0 loss_all = loss_seg_src + triger_ent * args.entW * loss_ent_trg # loss of seg on src, and ent on s and t loss_all.backward() optimizer.step() loss_train += loss_seg_src.detach().cpu().numpy() loss_val += loss_seg_trg.detach().cpu().numpy() if (i + 1) % args.save_pred_every == 0: print('taking snapshot ...') torch.save( model.state_dict(), os.path.join(args.snapshot_dir, '%s_' % (args.source) + str(i + 1) + '.pth')) if (i + 1) % args.print_freq == 0: _t['iter time'].toc(average=False) print('[it %d][src seg loss %.4f][trg seg loss %.4f][lr %.4f][%.2fs]' % \ (i + 1, loss_seg_src.data, loss_seg_trg.data, optimizer.param_groups[0]['lr']*10000, _t['iter time'].diff) ) sio.savemat(args.tempdata, { 'src_img': src_img.cpu().numpy(), 'trg_img': trg_img.cpu().numpy() }) loss_train /= args.print_freq loss_val /= args.print_freq loss_train_list.append(loss_train) loss_val_list.append(loss_val) sio.savemat(args.matname, { 'loss_train': loss_train_list, 'loss_val': loss_val_list }) loss_train = 0.0 loss_val = 0.0 if i + 1 > args.num_steps_stop: print('finish training') break _t['iter time'].tic()
def main(): opt = TestOptions() args = opt.initialize() os.environ["CUDA_VISIBLE_DEVICES"] = args.GPU if not os.path.exists(args.save): os.makedirs(args.save) # 3 models are used because MBT method is used. args.restore_from = args.restore_opt1 model1 = CreateModel(args) model1.eval() model1.cuda() args.restore_from = args.restore_opt2 model2 = CreateModel(args) model2.eval() model2.cuda() args.restore_from = args.restore_opt3 model3 = CreateModel(args) model3.eval() model3.cuda() targetloader = CreateTrgDataLoader(args) # change the mean for different dataset other than CS IMG_MEAN = np.array((104.00698793, 116.66876762, 122.67891434), dtype=np.float32) IMG_MEAN = torch.reshape(torch.from_numpy(IMG_MEAN), (1, 3, 1, 1)) mean_img = torch.zeros(1, 1) # ------------------------------------------------- # # compute scores and save them with torch.no_grad(): for index, batch in enumerate(targetloader): if index % 100 == 0: print('%d processd' % index) image, _, name = batch # 1. get image # create mean image if mean_img.shape[-1] < 2: B, C, H, W = image.shape # 2. get mean image mean_img = IMG_MEAN.repeat(B, 1, H, W) image = image.clone() - mean_img # 3, image - mean_img image = Variable(image).cuda() # forward output1 = model1(image) output1 = nn.functional.softmax(output1, dim=1) output2 = model2(image) output2 = nn.functional.softmax(output2, dim=1) output3 = model3(image) output3 = nn.functional.softmax(output3, dim=1) a, b = 0.3333, 0.3333 output = a * output1 + b * output2 + (1.0 - a - b) * output3 output = nn.functional.interpolate( output, (1024, 2048), mode='bilinear', align_corners=True).cpu().data[0].numpy() #output = nn.functional.upsample( output, (1024, 2048), mode='bilinear', align_corners=True).cpu().data[0].numpy() output = output.transpose(1, 2, 0) output_nomask = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) output_col = colorize_mask(output_nomask) output_nomask = Image.fromarray(output_nomask) name = name[0].split('/')[-1] output_nomask.save('%s/%s' % (args.save, name)) output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0])) # scores computed and saved # ------------------------------------------------- # compute_mIoU(args.gt_dir, args.save, args.devkit_dir, args.restore_from)
def main(): opt = TrainOptions() args = opt.initialize() _t = {'iter time': Timer()} model_name = args.source + '_to_' + args.target if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) os.makedirs(os.path.join(args.snapshot_dir, 'logs')) opt.print_options(args) sourceloader, targetloader = CreateSrcDataLoader( args), CreateTrgDataLoader(args) sourceloader_iter, targetloader_iter = iter(sourceloader), iter( targetloader) model, optimizer = CreateModel(args) model_D, optimizer_D = CreateDiscriminator(args) start_iter = 0 if args.restore_from is not None: start_iter = int(args.restore_from.rsplit('/', 1)[1].rsplit('_')[1]) train_writer = tensorboardX.SummaryWriter( os.path.join(args.snapshot_dir, "logs", model_name)) bce_loss = torch.nn.BCEWithLogitsLoss() cudnn.enabled = True cudnn.benchmark = True model.train() model.cuda() model_D.train() model_D.cuda() loss = [ 'loss_seg_src', 'loss_seg_trg', 'loss_D_trg_fake', 'loss_D_src_real', 'loss_D_trg_real' ] _t['iter time'].tic() for i in range(start_iter, args.num_steps): model.adjust_learning_rate(args, optimizer, i) model_D.adjust_learning_rate(args, optimizer_D, i) optimizer.zero_grad() optimizer_D.zero_grad() """ train Segmentation model """ for param in model_D.parameters(): param.requires_grad = False ''' train on [S':transferred source images] for loss_seg_src ''' src_img, src_lbl, _, _ = sourceloader_iter.next() src_img, src_lbl = Variable(src_img).cuda(), Variable( src_lbl.long()).cuda() src_seg_score = model(src_img, lbl=src_lbl) loss_seg_src = model.loss loss_seg_src.backward() ''' train on [T:target images] for loss_seg_trg ''' if args.data_label_folder_target is not None: trg_img, trg_lbl, _, _ = targetloader_iter.next() trg_img, trg_lbl = Variable(trg_img).cuda(), Variable( trg_lbl.long()).cuda() trg_seg_score = model(trg_img, lbl=trg_lbl) loss_seg_trg = model.loss else: trg_img, _, name = targetloader_iter.next() trg_img = Variable(trg_img).cuda() trg_seg_score = model(trg_img) loss_seg_trg = 0 outD_trg = model_D(F.softmax(trg_seg_score), 0) loss_D_trg_fake = model_D.loss loss_trg = args.lambda_adv_target * loss_D_trg_fake + loss_seg_trg loss_trg.backward() """ train Discriminator """ for param in model_D.parameters(): param.requires_grad = True src_seg_score, trg_seg_score = src_seg_score.detach( ), trg_seg_score.detach() outD_src = model_D(F.softmax(src_seg_score), 0) loss_D_src_real = model_D.loss / 2 loss_D_src_real.backward() outD_trg = model_D(F.softmax(trg_seg_score), 1) loss_D_trg_real = model_D.loss / 2 loss_D_trg_real.backward() """ update segmentation model & discriminator model """ optimizer.step() optimizer_D.step() for m in loss: train_writer.add_scalar(m, eval(m), i + 1) if (i + 1) % args.save_pred_every == 0: print('taking snapshot ...') torch.save( model.state_dict(), os.path.join(args.snapshot_dir, '%s_' % (args.source) + str(i + 1) + '.pth')) if (i + 1) % args.print_freq == 0: _t['iter time'].toc(average=False) print('[it %d][src seg loss %.4f][lr %.4f][%.2fs]' % (i + 1, loss_seg_src.data, optimizer.param_groups[0]['lr'] * 10000, _t['iter time'].diff)) if i + 1 > args.num_steps_stop: print('finish training') break _t['iter time'].tic()
def train_adp(self,round_): self.args.data_label_folder_target = root_base +'/dataset/generated_data/' self.args.shuffel_ = True self.args.data_dir_target = root_base +'/dataset/generated_data/' self.args.data_list_target = self.args.data_gen_list self.args.batch_size = 1 _t = {'iter time' : Timer()} self.args.num_steps = int(self.cnt_img * self.args.epoch_per_round * self.args.no_of_patches_per_image) sourceloader, targetloader = CreateSrcDataLoader(self.args), CreateTrgDataLoader(self.args) targetloader_iter, sourceloader_iter = iter(targetloader), iter(sourceloader) start_iter = 0 train_writer = tensorboardX.SummaryWriter(os.path.join(self.args.snapshot_dir, "logs", self.model_name)) cudnn.enabled = True cudnn.benchmark = True self.model.train() self.model.cuda() loss = ['loss_seg_src', 'loss_seg_trg'] _t['iter time'].tic() for i in range(start_iter, self.args.num_steps): self.model.adjust_learning_rate(self.args, self.optimizer, i) self.optimizer.zero_grad() src_img, src_lbl, _, _,_ = sourceloader_iter.next() src_img, src_lbl = Variable(src_img).cuda(), Variable(src_lbl.long()).cuda() src_seg_score = self.model(src_img, lbl=src_lbl) loss_seg_src = self.model.loss loss_src = torch.mean(loss_seg_src) ############################## loss_src.backward() trg_img, trg_lbl, _, _ = targetloader_iter.next() trg_img, trg_lbl = Variable(trg_img).cuda(), Variable(trg_lbl.long()).cuda() trg_seg_score = self.model(trg_img, lbl=trg_lbl) ############################ loss_seg_trg = self.model.loss ##########Focal loss############ loss_trg_2 = torch.mean(loss_seg_trg) if self.args.focal_loss: pt = torch.exp(-loss_seg_trg) loss_trg = loss_seg_trg * (1-pt)**self.args.gamma trg_fcl = torch.mean(loss_trg) else: trg_fcl = 0 loss_trg = self.args.beta *trg_fcl + loss_trg_2 loss_trg.backward() src_seg_score, trg_seg_score = src_seg_score.detach(), trg_seg_score.detach() self.optimizer.step() if (i+1) % self.args.saving_step == 0 : print ('taking snapshot ...') if(args.focal_loss): torch.save(self.model.state_dict(), os.path.join(self.args.snapshot_dir, '%s' %(self.args.source+"_to_" +self.args.target)+"_w_focal_loss_"+args.model +'.pth' )) else: torch.save(self.model.state_dict(), os.path.join(self.args.snapshot_dir, '%s' %(self.args.source+"_to_" +self.args.target)+"_wo_focal_loss" +args.model +'.pth' )) if (i+1) % 100 == 0: _t['iter time'].toc(average=False) print ('[it %d][src seg loss %.4f][lr %.4f][%.2fs]' % \ (i + 1, loss_src.data, self.optimizer.param_groups[0]['lr']*10000, _t['iter time'].diff)) _t['iter time'].tic()
def Generate(self): self.args.data_label_folder_target = None stop_point = self.cnt_img self.args.data_dir_target = root_base +'/dataset/cityscapes' self.args.data_list_target = self.args.sorted_list self.args.batch_size = 1 self.model.eval() self.model.cuda() targetloader = CreateTrgDataLoader(self.args) count = 0 f = open(self.args.data_gen_list,"w+") print("***********Generating Random Scale Invariant Example***********") print("Total it : ",stop_point+1) for index, batch in tqdm.tqdm(enumerate(targetloader)): if(index > stop_point): break image, _, name, rl_img = batch fn = name[0] output = self.model(Variable(image).cuda()) output = nn.functional.softmax(output, dim=1) rl_img = rl_img.cpu().data[0].numpy()/255.0 output=output.cpu().data[0].numpy() Y_pred = channel_revert(output) ############################################### map_ = np.zeros((IMG_H,IMG_W)) SE = self_entropy(Y_pred) SE = ent_normalization(SE,Y_pred) for il in range(IMG_H): for ij in range(IMG_W): if(SE[il][ij] >= self.entropy_th_class[0][np.argmax(Y_pred[0][il][ij],axis = -1)]): map_[il][ij] = 0 else: map_[il][ij] = 1 for patch in range(self.args.no_of_patches_per_image): if(self.args.patch_size[1] == 1024): X_Patch_resize = rl_img Y_Patch_resize = Y_pred map_patch_re = map_ else: point_X_s = self.args.patch_size[1] point_Y_s = self.args.patch_size[0] point_X = random.randint(0, IMG_W-point_X_s) point_Y = random.randint(0, IMG_H - point_Y_s) x_PATCH = rl_img[ point_Y:point_Y+point_Y_s, point_X:point_X+point_X_s].reshape(point_Y_s,point_X_s,3) y_PATCH = Y_pred[0, point_Y:point_Y+point_Y_s, point_X:point_X+point_X_s].reshape(1,point_Y_s,point_X_s,19) map_patch = map_[point_Y:point_Y+point_Y_s, point_X:point_X+point_X_s].reshape(point_Y_s,point_X_s) ################resize ############### x_PATCH = np.rollaxis(x_PATCH, axis = -1).reshape(1,3,point_Y_s,point_X_s) x_PATCH = torch.from_numpy(x_PATCH).float().to(device) y_PATCH = np.rollaxis(y_PATCH[0], axis = -1).reshape(1,19,point_Y_s,point_X_s) y_PATCH = torch.from_numpy(y_PATCH).float().to(device) map_patch = map_patch.reshape(1,1,point_Y_s,point_X_s) map_patch = torch.from_numpy(map_patch).float().to(device) X_Patch_resize = nn.functional.upsample(x_PATCH, (IMG_H, IMG_W), mode='bilinear', align_corners=True).cpu().data[0].numpy() X_Patch_resize = np.rollaxis(X_Patch_resize, axis = -1) X_Patch_resize = np.rollaxis(X_Patch_resize, axis = -1) Y_Patch_resize = nn.functional.upsample(y_PATCH, (IMG_H, IMG_W), mode='bilinear', align_corners=True).cpu().data[0].numpy() Y_Patch_resize = np.rollaxis(Y_Patch_resize, axis = -1) Y_Patch_resize = np.rollaxis(Y_Patch_resize, axis = -1).reshape(1,IMG_H,IMG_W,19) map_patch_re = nn.functional.upsample(map_patch, (IMG_H, IMG_W), mode='bilinear', align_corners=True).cpu().data[0].numpy() map_patch_re = map_patch_re[0] f.write(str(count+patch)+".png"+'\n') self.save_data(X_Patch_resize, Y_Patch_resize.reshape(IMG_H,IMG_W,19), self.args.generated_data,(count+patch),map_patch_re) count+=self.args.no_of_patches_per_image + 1 print("*****************END****************")
def class_wise_sort(self): self.args.data_label_folder_target = None self.args.data_dir_target = root_base +'/dataset/cityscapes' self.args.data_list_target = root_base +'/dataset/cityscapes_list/train.txt' self.args.num_steps = self.args.total_no_of_target self.args.batch_size = 1 self.model.eval() self.model.cuda() targetloader = CreateTrgDataLoader(self.args) data_save_class = {classes[0]:[],classes[1]:[],classes[2]:[],classes[3]:[],classes[4]:[],classes[5]:[],classes[6]:[],classes[7]:[],classes[8]:[],classes[9]:[],classes[10]:[],classes[11]:[],classes[12]:[],classes[13]:[],classes[14]:[],classes[15]:[],classes[16]:[],classes[17]:[],classes[18]:[],"f_name_"+classes[0]:[],"f_name_"+classes[1]:[],"f_name_"+classes[2]:[],"f_name_"+classes[3]:[],"f_name_"+classes[4]:[],"f_name_"+classes[5]:[],"f_name_"+classes[6]:[],"f_name_"+classes[7]:[],"f_name_"+classes[8]:[],"f_name_"+classes[9]:[],"f_name_"+classes[10]:[],"f_name_"+classes[11]:[],"f_name_"+classes[12]:[],"f_name_"+classes[13]:[],"f_name_"+classes[14]:[],"f_name_"+classes[15]:[],"f_name_"+classes[16]:[],"f_name_"+classes[17]:[],"f_name_"+classes[18]:[]} data_save_class_sorted = {classes[0]:[],classes[1]:[],classes[2]:[],classes[3]:[],classes[4]:[],classes[5]:[],classes[6]:[],classes[7]:[],classes[8]:[],classes[9]:[],classes[10]:[],classes[11]:[],classes[12]:[],classes[13]:[],classes[14]:[],classes[15]:[],classes[16]:[],classes[17]:[],classes[18]:[],"f_name_"+classes[0]:[],"f_name_"+classes[1]:[],"f_name_"+classes[2]:[],"f_name_"+classes[3]:[],"f_name_"+classes[4]:[],"f_name_"+classes[5]:[],"f_name_"+classes[6]:[],"f_name_"+classes[7]:[],"f_name_"+classes[8]:[],"f_name_"+classes[9]:[],"f_name_"+classes[10]:[],"f_name_"+classes[11]:[],"f_name_"+classes[12]:[],"f_name_"+classes[13]:[],"f_name_"+classes[14]:[],"f_name_"+classes[15]:[],"f_name_"+classes[16]:[],"f_name_"+classes[17]:[],"f_name_"+classes[18]:[]} for index, batch in tqdm.tqdm(enumerate(targetloader)): image, _, name ,__= batch fn = name[0] output = self.model(Variable(image).cuda()) output = nn.functional.softmax(output, dim=1) output=output.cpu().data[0].numpy() y_pred = channel_revert(output) MP_max = np.max(y_pred,axis = -1)[0] LP_argmax = np.argmax(y_pred,axis = -1)[0] for c in range(len(classes) ): M_c = MP_max[LP_argmax == c] if(M_c.any()): data_save_class['f_name_'+classes[c]].append(fn) data_save_class[classes[c]].append(np.mean(M_c)) f = open(self.args.sorted_list,"w+") self.cnt_img=0 over_all_list = [] for c in range(len(classes) ): if(self.args.source == 'synthia'): if(c !=9 and c !=14 and c !=16): max_prob_sorted_ascending, f_name_arrays = zip(*sorted(zip(data_save_class[classes[c]], data_save_class['f_name_'+classes[c]]))) f_name_arrays = f_name_arrays[::-1] max_prob_sorted_ascending = max_prob_sorted_ascending[::-1] data_save_class_sorted[classes[c]] = max_prob_sorted_ascending data_save_class_sorted['f_name_'+classes[c]] = f_name_arrays select_imgs = int(len(f_name_arrays) * (self.args.p/19)) selct = f_name_arrays[0:select_imgs] ######Dynamic threshold of each class############## f_n = f_name_arrays[select_imgs] img= load_single_image(f_n) img = img.reshape(1,3,IMG_H,IMG_W) img = torch.from_numpy(img).float().to(device) out_ = self.model(Variable(img).cuda()) out_ = nn.functional.softmax(out_, dim=1) out_=out_.cpu().data[0].numpy() y_pred = channel_revert(out_) Y_argmax = np.argmax(y_pred,axis = -1)[0] SE_main = self_entropy(y_pred) SE_main = ent_normalization(SE_main,y_pred) class_se = SE_main[Y_argmax == c] mean_class_se = np.mean(class_se) self.entropy_th_class[0][c] = mean_class_se ############################################## over_all_list += selct else: max_prob_sorted_ascending, f_name_arrays = zip(*sorted(zip(data_save_class[classes[c]], data_save_class['f_name_'+classes[c]]))) f_name_arrays = f_name_arrays[::-1] max_prob_sorted_ascending = max_prob_sorted_ascending[::-1] data_save_class_sorted[classes[c]] = max_prob_sorted_ascending data_save_class_sorted['f_name_'+classes[c]] = f_name_arrays select_imgs = int(len(f_name_arrays) * (self.args.p/19)) selct = f_name_arrays[0:select_imgs] ######Dynamic threshold of each class############## f_n = f_name_arrays[select_imgs] img= load_single_image(f_n) img = img.reshape(1,3,IMG_H,IMG_W) img = torch.from_numpy(img).float().to(device) out_ = self.model(Variable(img).cuda()) out_ = nn.functional.softmax(out_, dim=1) out_=out_.cpu().data[0].numpy() y_pred = channel_revert(out_) Y_argmax = np.argmax(y_pred,axis = -1)[0] SE_main = self_entropy(y_pred) SE_main = ent_normalization(SE_main,y_pred) class_se = SE_main[Y_argmax == c] mean_class_se = np.mean(class_se) self.entropy_th_class[0][c] = mean_class_se ############################################## over_all_list += selct over_all_list = list(dict.fromkeys(over_all_list)) for loop in over_all_list: f.write(loop+'\n') self.cnt_img +=1
def main(): opt = TestOptions() args = opt.initialize() if not os.path.exists(args.save): os.makedirs(args.save) model = CreateModel(args) model.eval() model.cuda() targetloader = CreateTrgDataLoader(args) id_to_trainid = {2: 0, 1: 128, 0: 255} for index, batch in enumerate(targetloader): if index % 10 == 0: print('%d processd' % index) image, _, name = batch _, output, _, _ = model(Variable(image).cuda()) #[1,3,129,129] #import pdb;pdb.set_trace() output = nn.functional.softmax(output, dim=1) output = nn.functional.upsample( output, (1634, 1634), mode='bilinear', align_corners=True).cpu().data[0].numpy() output = output.transpose(1, 2, 0) #(1634,1634,3) ''' output_crop = output[14:526, 626:1138, 0:2] np.save("/extracephonline/medai_data2/zhengdzhang/eyes/qikan/cai/output_crop_1.npy", output_crop) #crop_img = Image.fromarray(output_crop) #crop_img.save("/extracephonline/medai_data2/zhengdzhang/eyes/qikan/cai/crop_img.png") ''' output_nomask = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) #(1644,1634) unique:[0,1,2] cup_mask = get_bool(output_nomask, 0) disc_mask = get_bool(output_nomask, 0) + get_bool(output_nomask, 1) disc_mask = disc_mask.astype(np.uint8) cup_mask = cup_mask.astype(np.uint8) for i in range(5): disc_mask = scipy.signal.medfilt2d(disc_mask, 19) cup_mask = scipy.signal.medfilt2d(cup_mask, 19) disc_mask = morphology.binary_erosion(disc_mask, morphology.diamond(7)).astype( np.uint8) # return 0,1 cup_mask = morphology.binary_erosion(cup_mask, morphology.diamond(7)).astype( np.uint8) # return 0,1 disc_mask = get_largest_fillhole(disc_mask) cup_mask = get_largest_fillhole(cup_mask) disc_mask = morphology.binary_dilation(disc_mask, morphology.diamond(7)).astype( np.uint8) # return 0,1 cup_mask = morphology.binary_dilation(cup_mask, morphology.diamond(7)).astype( np.uint8) # return 0,1 disc_mask = get_largest_fillhole(disc_mask).astype( np.uint8) # return 0,1 cup_mask = get_largest_fillhole(cup_mask).astype(np.uint8) output_nomask = disc_mask + cup_mask output_col = np.ones(output_nomask.shape, dtype=np.float32) for k, v in id_to_trainid.items(): output_col[output_nomask == k] = v output_col = Image.fromarray(output_col.astype(np.uint8)) output_nomask = Image.fromarray(output_nomask) name = name[0].split('.')[0] + '.png' output_nomask.save('%s/%s' % (args.save, name)) output_col.save('%s/color_%s' % (args.save, name)) disc_dice, cup_dice = calculate_dice(args.gt_dir, args.save, args.devkit_dir) print('===> disc_dice:' + str(round(disc_dice, 3)) + '\t' + 'cup_dice:' + str(round(cup_dice, 3)))
def main(): opt = TrainOptions() args = opt.initialize() _t = {'iter time': Timer()} model_name = args.source + '_to_' + args.target if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) os.makedirs(os.path.join(args.snapshot_dir, 'logs')) opt.print_options(args) sourceloader, targetloader = CreateSrcDataLoader( args), CreateTrgDataLoader(args) targetloader_iter, sourceloader_iter = iter(targetloader), iter( sourceloader) model, optimizer = CreateModel(args) model_D1, optimizer_D1 = CreateDiscriminator(args, 1) model_D2, optimizer_D2 = CreateDiscriminator(args, 2) start_iter = 0 if args.restore_from is not None: start_iter = int(args.restore_from.rsplit('/', 1)[1].rsplit('_')[1]) train_writer = tensorboardX.SummaryWriter( os.path.join(args.snapshot_dir, "logs", model_name)) bce_loss = torch.nn.BCEWithLogitsLoss() interp_target = nn.Upsample(size=(1024, 1024), mode='bilinear', align_corners=True) interp_source = nn.Upsample(size=(1024, 1024), mode='bilinear', align_corners=True) cudnn.enabled = True cudnn.benchmark = True model.train() model.cuda() model_D1.train() model_D1.cuda() model_D2.train() model_D2.cuda() weight_loss = WeightedBCEWithLogitsLoss() weight_map_loss = WeightMapLoss() loss = [ 'loss_seg_src', 'loss_seg_trg', 'loss_D_trg_fake', 'loss_D_src_real', 'loss_D_trg_real' ] _t['iter time'].tic() for i in range(start_iter, args.num_steps): print(i) model.adjust_learning_rate(args, optimizer, i) model_D1.adjust_learning_rate(args, optimizer_D1, i) model_D2.adjust_learning_rate(args, optimizer_D2, i) optimizer.zero_grad() optimizer_D1.zero_grad() optimizer_D2.zero_grad() ##train G for param in model_D1.parameters(): param.requires_grad = False for param in model_D2.parameters(): param.requires_grad = False try: src_img, src_lbl, weight_map, _ = sourceloader_iter.next() except StopIteration: sourceloader_iter = iter(sourceloader) src_img, src_lbl, weight_map, _ = sourceloader_iter.next() src_img, src_lbl, weight_map = Variable(src_img).cuda(), Variable( src_lbl.long()).cuda(), Variable(weight_map.long()).cuda() src_seg_score1, src_seg_score2, src_seg_score3, src_seg_score4 = model( src_img, lbl=src_lbl, weight=weight_map) #import pdb;pdb.set_trace() #WeightLoss1 = weight_map_loss(src_seg_score1, src_lbl, weight_map) #WeightLoss2 = weight_map_loss(src_seg_score2, src_lbl, weight_map) loss_seg_src = model.loss #print('WeightLoss2, WeightLoss1:', WeightLoss2.data, WeightLoss1.data) loss_seg_src.backward() if args.data_label_folder_target is not None: trg_img, trg_lbl, _, name = targetloader_iter.next() trg_img, trg_lbl = Variable(trg_img).cuda(), Variable( trg_lbl.long()).cuda() trg_seg_score1, trg_seg_score2, trg_seg_score3, trg_seg_score4 = model( trg_img, lbl=trg_lbl) loss_seg_trg = model.loss else: trg_img, _, name = targetloader_iter.next() trg_img = Variable(trg_img).cuda() trg_seg_score1, trg_seg_score2, trg_seg_score3, trg_seg_score4 = model( trg_img) loss_seg_trg = 0 outD1_trg = model_D1(F.softmax(trg_seg_score1), 0) outD2_trg = model_D2(F.softmax(trg_seg_score2), 0) #import pdb;pdb.set_trace() outD1_trg = interp_target(outD1_trg) #[1, 1, 1024, 1024] outD2_trg = interp_target(outD2_trg) ''' if i > 9001: #import pdb;pdb.set_trace() weight_map1 = prob_2_entropy(F.softmax(trg_seg_score1)) #[1, 1, 1024, 1024] weight_map2 = prob_2_entropy(F.softmax(trg_seg_score2)) #[1, 1, 1024, 1024] loss_D1_trg_fake = weight_loss(outD1_trg, Variable(torch.FloatTensor(outD1_trg.data.size()).fill_(0)).cuda(), weight_map1, 0.3, 1) loss_D2_trg_fake = weight_loss(outD2_trg, Variable(torch.FloatTensor(outD2_trg.data.size()).fill_(0)).cuda(), weight_map2, 0.3, 1) else: loss_D1_trg_fake = model_D1.loss loss_D2_trg_fake = model_D2.loss loss_D_trg_fake = loss_D1_trg_fake*0.2 + loss_D2_trg_fake ''' loss_D_trg_fake = model_D1.loss * 0.2 + model_D2.loss loss_trg = args.lambda_adv_target * loss_D_trg_fake + loss_seg_trg loss_trg.backward() ###train D for param in model_D1.parameters(): param.requires_grad = True for param in model_D2.parameters(): param.requires_grad = True src_seg_score1, src_seg_score2, src_seg_score3, src_seg_score4, trg_seg_score1, trg_seg_score2, trg_seg_score3, trg_seg_score4 = src_seg_score1.detach( ), src_seg_score2.detach(), src_seg_score3.detach( ), src_seg_score4.detach(), trg_seg_score1.detach( ), trg_seg_score2.detach(), trg_seg_score3.detach( ), trg_seg_score4.detach() outD1_src = model_D1(F.softmax(src_seg_score1), 0) outD2_src = model_D2(F.softmax(src_seg_score2), 0) loss_D1_src_real = model_D1.loss / 2 loss_D1_src_real.backward() loss_D2_src_real = model_D2.loss / 2 loss_D2_src_real.backward() loss_D_src_real = loss_D1_src_real + loss_D2_src_real outD1_trg = model_D1(F.softmax(trg_seg_score1), 1) outD2_trg = model_D2(F.softmax(trg_seg_score2), 1) outD1_trg = interp_target(outD1_trg) outD2_trg = interp_target(outD2_trg) if i > 9001: weight_map1 = prob_2_entropy(F.softmax(trg_seg_score1)) weight_map2 = prob_2_entropy(F.softmax(trg_seg_score2)) loss_D1_trg_real = weight_loss( outD1_trg, Variable(torch.FloatTensor( outD1_trg.data.size()).fill_(1)).cuda(), weight_map1, 0.3, 1) / 2 loss_D2_trg_real = weight_loss( outD2_trg, Variable(torch.FloatTensor( outD2_trg.data.size()).fill_(1)).cuda(), weight_map2, 0.3, 1) / 2 else: loss_D1_trg_real = model_D1.loss / 2 loss_D2_trg_real = model_D2.loss / 2 loss_D1_trg_real.backward() loss_D2_trg_real.backward() loss_D_trg_real = loss_D1_trg_real + loss_D2_trg_real optimizer.step() optimizer_D1.step() optimizer_D2.step() for m in loss: train_writer.add_scalar(m, eval(m), i + 1) if (i + 1) % args.save_pred_every == 0: print('taking snapshot ...') torch.save( model.state_dict(), os.path.join(args.snapshot_dir, '%s_' % (args.source) + str(i + 1) + '.pth')) torch.save( model_D1.state_dict(), os.path.join(args.snapshot_dir, '%s_' % (args.source) + str(i + 1) + '_D1.pth')) torch.save( model_D2.state_dict(), os.path.join(args.snapshot_dir, '%s_' % (args.source) + str(i + 1) + '_D2.pth')) if (i + 1) % args.print_freq == 0: _t['iter time'].toc(average=False) print ('[it %d][src seg loss %.4f][trg seg loss %.4f][lr %.4f][%.2fs]' % \ (i + 1, loss_seg_src.data,loss_seg_trg.data, optimizer.param_groups[0]['lr']*10000, _t['iter time'].diff)) if i + 1 > args.num_steps_stop: print('finish training') break _t['iter time'].tic()
def main(): opt = TrainOptions() args = opt.initialize() os.environ["CUDA_VISIBLE_DEVICES"] = args.GPU _t = {'iter time': Timer()} model_name = args.source + '_to_' + args.target if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) os.makedirs(os.path.join(args.snapshot_dir, 'logs')) opt.print_options(args) sourceloader, targetloader = CreateSrcDataLoader( args), CreateTrgDataLoader(args) sourceloader_iter, targetloader_iter = iter(sourceloader), iter( targetloader) pseudotrgloader = CreatePseudoTrgLoader(args) pseudoloader_iter = iter(pseudotrgloader) model, optimizer = CreateModel(args) start_iter = 0 if args.restore_from is not None: start_iter = int(args.restore_from.rsplit('/', 1)[1].rsplit('_')[1]) if args.restore_optim_from is not None: optimizer.load_state_dict(torch.load(args.restore_optim_from)) for state in optimizer.state.values(): for k, v in state.items(): if isinstance(v, torch.Tensor): state[k] = v.cuda() cudnn.enabled = True cudnn.benchmark = True model.train() model.cuda() wandb.watch(model, log='gradient', log_freq=1) # losses to log loss = ['loss_seg_src', 'loss_seg_psu'] loss_train = 0.0 loss_val = 0.0 loss_pseudo = 0.0 loss_train_list = [] loss_val_list = [] loss_pseudo_list = [] mean_img = torch.zeros(1, 1) class_weights = Variable(CS_weights).cuda() _t['iter time'].tic() for i in range(start_iter, args.num_steps): model.adjust_learning_rate(args, optimizer, i) # adjust learning rate optimizer.zero_grad() # zero grad src_img, src_lbl, _, _ = sourceloader_iter.next() # new batch source trg_img, trg_lbl, _, _ = targetloader_iter.next() # new batch target psu_img, psu_lbl, _, _ = pseudoloader_iter.next() scr_img_copy = src_img.clone() if mean_img.shape[-1] < 2: B, C, H, W = src_img.shape mean_img = IMG_MEAN.repeat(B, 1, H, W) #-------------------------------------------------------------------# # 1. source to target, target to target src_in_trg = FDA_source_to_target(src_img, trg_img, L=args.LB) # src_lbl trg_in_trg = trg_img # 2. subtract mean src_img = src_in_trg.clone() - mean_img # src_1, trg_1, src_lbl trg_img = trg_in_trg.clone() - mean_img # trg_1, trg_0, trg_lbl psu_img = psu_img.clone() - mean_img #-------------------------------------------------------------------# # evaluate and update params ##### src_img, src_lbl = Variable(src_img).cuda(), Variable( src_lbl.long()).cuda() # to gpu src_seg_score = model(src_img, lbl=src_lbl, weight=class_weights, ita=args.ita) # forward pass loss_seg_src = model.loss_seg # get loss loss_ent_src = model.loss_ent # use pseudo label as supervision psu_img, psu_lbl = Variable(psu_img).cuda(), Variable( psu_lbl.long()).cuda() psu_seg_score = model(psu_img, lbl=psu_lbl, weight=class_weights, ita=args.ita) loss_seg_psu = model.loss_seg loss_ent_psu = model.loss_ent loss_all = loss_seg_src + (loss_seg_psu + args.entW * loss_ent_psu ) # loss of seg on src, and ent on s and t loss_all.backward() optimizer.step() loss_train += loss_seg_src.detach().cpu().numpy() loss_val += loss_seg_psu.detach().cpu().numpy() if (i + 1) % args.save_pred_every == 0: print('taking snapshot ...') torch.save( model.state_dict(), os.path.join(args.snapshot_dir, '%s_' % (args.source) + str(i + 1) + '.pth')) torch.save( optimizer.state_dict(), os.path.join(args.snapshot_dir_optim, '%s_' % (args.source) + '.pth')) wandb.log({ "src seg loss": loss_seg_src.data, "psu seg loss": loss_seg_psu.data, "learnign rate": optimizer.param_groups[0]['lr'] * 10000 }) if (i + 1) % args.print_freq == 0: _t['iter time'].toc(average=False) print('[it %d][src seg loss %.4f][psu seg loss %.4f][lr %.4f][%.2fs]' % \ (i + 1, loss_seg_src.data, loss_seg_psu.data, optimizer.param_groups[0]['lr']*10000, _t['iter time'].diff) ) sio.savemat(args.tempdata, { 'src_img': src_img.cpu().numpy(), 'trg_img': trg_img.cpu().numpy() }) loss_train /= args.print_freq loss_val /= args.print_freq loss_train_list.append(loss_train) loss_val_list.append(loss_val) sio.savemat(args.matname, { 'loss_train': loss_train_list, 'loss_val': loss_val_list }) loss_train = 0.0 loss_val = 0.0 if i + 1 > args.num_steps_stop: print('finish training') break _t['iter time'].tic()
def main(): args = arg_parser.Parse() os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus logger = Logger(args.log_dir) logger.PrintAndLogArgs(args) saver = ImageAndLossSaver(args.tb_logs_dir, logger.log_folder, args.checkpoints_dir, args.save_pics_every) source_loader, target_train_loader, target_eval_loader = CreateSrcDataLoader( args), CreateTrgDataLoader(args, 'train'), CreateTrgDataLoader(args, 'val') epoch_size = np.maximum(len(target_train_loader.dataset), len(source_loader.dataset)) steps_per_epoch = int(np.floor(epoch_size / args.batch_size)) source_loader.dataset.SetEpochSize(epoch_size) target_train_loader.dataset.SetEpochSize(epoch_size) generator = model.DeepLPFNet() generator = nn.DataParallel(generator.cuda()) generator_criterion = model.GeneratorLoss() generator_optimizer = optim.Adam(generator.parameters(), lr=args.generator_lr, betas=(0.9, 0.999), eps=1e-08) discriminator = model.Discriminator() discriminator = nn.DataParallel(discriminator.cuda()) discriminator_criterion = model.DiscriminatorLoss() discriminator_optimizer = optim.Adam(discriminator.parameters(), lr=args.discriminator_lr, betas=(0.9, 0.999), eps=1e-08) semseg_net, semseg_optimizer = CreateModel(args) semseg_net = nn.DataParallel(semseg_net.cuda()) logger.info('######### Network created #########') logger.info('Architecture of Generator:\n' + str(generator)) logger.info('Architecture of Discriminator:\n' + str(discriminator)) logger.info('Architecture of Backbone net:\n' + str(semseg_net)) for epoch in range(args.num_epochs): generator.train() discriminator.train() semseg_net.train() saver.Reset() discriminate_src = True source_loader_iter, target_train_loader_iter, target_eval_loader_iter = iter( source_loader), iter(target_train_loader), iter(target_eval_loader) logger.info('#################[Epoch %d]#################' % (epoch + 1)) for batch_num in range(steps_per_epoch): start_time = time.time() training_discriminator = (batch_num >= args.generator_boost) and ( batch_num - args.generator_boost) % ( args.discriminator_iters + args.generator_iters) < args.discriminator_iters src_img, src_lbl, src_shapes, src_names = source_loader_iter.next( ) # new batch source trg_eval_img, trg_eval_lbl, trg_shapes, trg_names = target_train_loader_iter.next( ) # new batch target generator_optimizer.zero_grad() discriminator_optimizer.zero_grad() semseg_optimizer.zero_grad() src_input_batch = Variable(src_img, requires_grad=False).cuda() src_label_batch = Variable(src_lbl, requires_grad=False).cuda() trg_input_batch = Variable(trg_eval_img, requires_grad=False).cuda() # trg_label_batch = Variable(trg_lbl, requires_grad=False).cuda() src_in_trg = generator(src_input_batch, trg_input_batch) # G(S,T) if training_discriminator: #train discriminator if discriminate_src == True: discriminator_src_in_trg = discriminator( src_in_trg) # D(G(S,T)) discriminator_trg = None # D(T) else: discriminator_src_in_trg = None # D(G(S,T)) discriminator_trg = discriminator(trg_input_batch) # D(T) discriminate_src = not discriminate_src loss = discriminator_criterion(discriminator_src_in_trg, discriminator_trg) else: #train generator and semseg net discriminator_trg = discriminator(trg_input_batch) # D(T) predicted, loss_seg, loss_ent = semseg_net( src_in_trg, lbl=src_label_batch) # F(G(S.T)) src_in_trg_labels = torch.argmax(predicted, dim=1) loss = generator_criterion(loss_seg, loss_ent, args.entW, discriminator_trg) saver.WriteLossHistory(training_discriminator, loss.item()) loss.backward() if training_discriminator: # train discriminator discriminator_optimizer.step() else: # train generator and semseg net generator_optimizer.step() semseg_optimizer.step() saver.running_time += time.time() - start_time if (not training_discriminator) and saver.SaveImagesIteration: saver.SaveTrainImages(epoch, src_img[0, :, :, :], src_in_trg[0, :, :, :], src_lbl[0, :, :], src_in_trg_labels[0, :, :]) if (batch_num + 1) % args.print_every == 0: logger.PrintAndLogData(saver, epoch, batch_num, args.print_every) if (batch_num + 1) % args.save_checkpoint == 0: saver.SaveModelsCheckpoint(semseg_net, discriminator, generator, epoch, batch_num) #Validation: semseg_net.eval() rand_samp_inds = np.random.randint(0, len(target_eval_loader.dataset), 5) rand_batchs = np.floor(rand_samp_inds / args.batch_size).astype(np.int) cm = torch.zeros((NUM_CLASSES, NUM_CLASSES)).cuda() for val_batch_num, (trg_eval_img, trg_eval_lbl, _, _) in enumerate(target_eval_loader): with torch.no_grad(): trg_input_batch = Variable(trg_eval_img, requires_grad=False).cuda() trg_label_batch = Variable(trg_eval_lbl, requires_grad=False).cuda() pred_softs_batch = semseg_net(trg_input_batch) pred_batch = torch.argmax(pred_softs_batch, dim=1) cm += compute_cm_batch_torch(pred_batch, trg_label_batch, IGNORE_LABEL, NUM_CLASSES) print('Validation: saw', val_batch_num * args.batch_size, 'examples') if (val_batch_num + 1) in rand_batchs: rand_offset = np.random.randint(0, args.batch_size) saver.SaveValidationImages( epoch, trg_input_batch[rand_offset, :, :, :], trg_label_batch[rand_offset, :, :], pred_batch[rand_offset, :, :]) iou, miou = compute_iou_torch(cm) saver.SaveEpochAccuracy(iou, miou, epoch) logger.info( 'Average accuracy of Epoch #%d on target domain: mIoU = %2f' % (epoch + 1, miou)) logger.info( '-----------------------------------Epoch #%d Finished-----------------------------------' % (epoch + 1)) del cm, trg_input_batch, trg_label_batch, pred_softs_batch, pred_batch saver.tb.close() logger.info('Finished training.')