def main(args): print('===================================\n', ) print("Root directory: {}".format(args.name)) args.exp_dir = os.path.join(os.path.join(RES_DIR, args.name), "evaluation_trinity_test") if not os.path.isdir(args.exp_dir): os.makedirs(args.exp_dir) print("EXP PATH: {}".format(args.exp_dir)) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args.device = device assert args.checkpoint_seg is not None, "Need trained .pth!" model_seg = load_models( mode="segmentation", device=device, args=args, ) assert args.checkpoint_disc is not None, "Need trained .pth!" model_disc = load_models( mode="single_discriminator", device=device, args=args, ) transforms_shape_target = dual_transforms.Compose([ dual_transforms.CenterCrop((400, 400)), dual_transforms.Scale(args.image_size[0]), ]) photo_transformer = PhotometricTransform(photometric_transform_config) traindata = OpenEDSDataset_withLabels( root=os.path.join(args.target_root, "train"), image_size=args.image_size, data_to_train="", shape_transforms=transforms_shape_target, photo_transforms=photo_transformer, train_bool=False, ) train_loader = torch.utils.data.DataLoader( traindata, batch_size=1, shuffle=False, num_workers=0, pin_memory=True, ) trainloader_iter = enumerate(train_loader) args.total_iterations = traindata.__len__() // args.batch_size model_seg.eval() model_disc.eval() confidence_map = [] confidence_mean_score = [] confidence_cnt = [] imagename = [] imageindex = [] dst_folder = os.path.join(args.exp_dir, "visualization") if not os.path.isdir(dst_folder): os.mkdir(dst_folder) image_confidence = np.empty((traindata.__len__(), 224, 224, 2)) # pred_all = np.empty((traindata.__len__(),4,224,224)) # with open("dataloaders/eye/trinity_top_200.pkl", "rb") as f: # top_200_lists = pickle.load(f) with open("dataloaders/eye/top_1_adv.pkl", "rb") as f: top_1_lists = pickle.load(f) with open("dataloaders/eye/top_2_adv.pkl", "rb") as f: top_2_lists = pickle.load(f) for i_iter in range(args.total_iterations): if i_iter % 1000 == 0: print("Processing {} ..........".format(i_iter)) # if not (traindata.train_data_list[i_iter] in top_200_lists): # continue # if traindata.train_data_list[i_iter] in top_1_lists: # continue if traindata.train_data_list[i_iter] in top_2_lists: continue imageindex.append(i_iter) imagename.append(traindata.train_data_list[i_iter]) _, batch = next(trainloader_iter) images, labels = batch images = Variable(images).to(args.device) labels = Variable(labels.long()).to(args.device) pred = model_seg(images) pred_softmax = F.softmax(pred, dim=1) D_out = model_disc(pred_softmax) D_out = torch.sigmoid(D_out) D_out = D_out[0, 0, :, :].detach().cpu().numpy() pred = np.argmax(pred.detach().cpu().numpy(), axis=1)[0, :, :] # fig = plt.figure() # ax = fig.add_subplot(231) # ax.imshow(images[0,0,:,:].detach().cpu().numpy(), cmap="gray") # ax.set_xticks([]) # ax.set_yticks([]) # # ax = fig.add_subplot(232) # ax.imshow(labels[0, :, :].detach().cpu().numpy()) # ax.set_xticks([]) # ax.set_yticks([]) # # ax = fig.add_subplot(233) # ax.imshow(pred, cmap="gray") # ax.set_xticks([]) # ax.set_yticks([]) # # ax = fig.add_subplot(234) # ax.imshow(D_out, cmap="gray") # ax.set_xticks([]) # ax.set_yticks([]) D_out_mean = D_out.mean() D_out_mean_map = (D_out > D_out_mean) * 1 # labels = labels[0,:,:].detach().cpu().numpy() # semi_ignore_mask = (D_out < D_out_mean) # # pseudo_gt = labels.copy() # # pseudo_gt[semi_ignore_mask] = 4 # # pseudo_gt = pseudo_gt.astype(np.uint8) # filename = traindata.train_data_list[i_iter].replace("/images/", "/masks/") # filename = filename.replace("/train/", "/train_pseudo/") # filename = filename.replace(".png", ".npy") # np.save(filename, semi_ignore_mask) # # print(D_out_mean_map.shape) # ax = fig.add_subplot(235) # ax.imshow(D_out_mean_map) # ax.set_xticks([]) # ax.set_yticks([]) # plt.tight_layout() # filename = ntpath.basename(traindata.train_data_list[i_iter]) # filename = os.path.join(dst_folder, filename) # plt.savefig(filename) # im_filename = traindata.train_data_list[i_iter].replace("/train/", "/train_pseudo/") # os.system("cp %s %s" % (traindata.train_data_list[i_iter], im_filename)) confidence_mean_score.append(D_out_mean) confidence_cnt.append(D_out_mean_map.sum()) ### generate confidence map ### # confidence_map.append(D_out_mean) # imagename.append(ntpath.basename(traindata.train_data_list[i_iter])) # image_confidence[i_iter,:,:,0] = images[0,0,:,:].detach().cpu().numpy() # image_confidence[i_iter,:,:,1] = D_out # pred_all[i_iter, ...] = pred_softmax[0,...].detach().cpu().numpy() ### generate confidence map ### # with open("%s/confidence_map_top1_adv.pkl" % (args.exp_dir), "wb") as f: # pickle.dump([imageindex, imagename, confidence_mean_score, confidence_cnt], f) with open("%s/confidence_map_top2_adv.pkl" % (args.exp_dir), "wb") as f: pickle.dump( [imageindex, imagename, confidence_mean_score, confidence_cnt], f)
def main(args): print('===================================\n', ) print("Root directory: {}".format(args.name)) args.exp_dir = os.path.join(RES_DIR, args.name) if not os.path.isdir(args.exp_dir): os.makedirs(args.exp_dir) print("EXP PATH: {}".format(args.exp_dir)) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args.device = device logger = make_logger(filename="TrainVal.log", args=args) if args.tensorboard: writer = SummaryWriter(args.exp_dir) else: writer = None print("===================================") print("====== Loading Training Data ======") print("===================================") shape_transformer = dual_transforms.Compose([ dual_transforms.CenterCrop((400, 400)), dual_transforms.Scale(args.image_size[0]), ]) photo_transformer = PhotometricTransform(photometric_transform_config) source_data = UnityDataset( root=args.source_root, image_size=args.image_size, shape_transforms=shape_transformer, photo_transforms=photo_transformer, train_bool=False, ) args.tot_source = source_data.__len__() args.total_iterations = args.num_epochs * args.tot_source // args.batch_size args.iters_to_eval = args.epoch_to_eval * args.tot_source // args.batch_size train_loader = torch.utils.data.DataLoader( source_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) model_seg = load_models( mode="segmentation", device=device, args=args, ) optimizer_seg = optim.Adam( model_seg.parameters(), lr=args.lr_seg, betas=(args.beta1, 0.999), ) optimizer_seg.zero_grad() seg_loss_target = torch.nn.CrossEntropyLoss().to(device) trainloader_iter = enumerate(train_loader) loss_seg_min = float("inf") miou_max = float("-inf") for i_iter in range(args.total_iterations): loss_seg_value = 0 model_seg.train() optimizer_seg.zero_grad() adjust_learning_rate( optimizer=optimizer_seg, learning_rate=args.lr_seg, i_iter=i_iter, max_steps=args.total_iterations, power=0.9, ) try: _, batch = next(trainloader_iter) except StopIteration: trainloader_iter = enumerate(train_loader) _, batch = next(trainloader_iter) images, labels = batch images = Variable(images).to(args.device) labels = Variable(labels.long()).to(args.device) pred = model_seg(images) loss_seg = seg_loss_target(pred, labels) current_loss_seg = loss_seg.item() loss_seg_value += current_loss_seg loss_seg.backward() optimizer_seg.step() pred_img = pred.argmax(dim=1, keepdim=True) flat_pred = pred_img.detach().cpu().numpy().flatten() flat_gt = labels.detach().cpu().numpy().flatten() miou, _ = compute_mean_iou(flat_pred=flat_pred, flat_label=flat_gt) logger.info('iter = {0:8d}/{1:8d} ' 'loss_seg = {2:.3f} ' 'mIoU = {3:.3f} '.format( i_iter, args.total_iterations, loss_seg_value, miou, )) if args.tensorboard and (writer != None): writer.add_scalar('Train/Cross_Entropy', current_loss_seg, i_iter) writer.add_scalar('Train/mIoU', miou, i_iter) if i_iter % args.iters_to_eval == 0: filename = os.path.join( args.exp_dir, "Target_img_trainiter_{}.png".format(i_iter)) target_img = images.float() gen_target_img = torchvision.utils.make_grid(target_img, padding=2, normalize=True) torchvision.utils.save_image(gen_target_img, filename) filename = os.path.join( args.exp_dir, "Unity_pred_trainiter_{}.png".format(i_iter)) pred_img = pred_img.float() gen_img = torchvision.utils.make_grid(pred_img, padding=2, normalize=True) torchvision.utils.save_image(gen_img, filename) is_better_ss = current_loss_seg < loss_seg_min if is_better_ss: loss_seg_min = current_loss_seg torch.save(model_seg.state_dict(), os.path.join(args.exp_dir, "model_train_best.pth")) if miou > miou_max: miou_max = miou torch.save(model_seg.state_dict(), os.path.join(args.exp_dir, "model_train_best_miou.pth")) logger.info("==========================================") logger.info("Training DONE!") if args.tensorboard and (writer != None): writer.close()
def main(args): print('===================================\n', ) print("Root directory: {}".format(args.name)) args.exp_dir = os.path.join(os.path.join(RES_DIR, args.name), "evaluation_trinity_test") if not os.path.isdir(args.exp_dir): os.makedirs(args.exp_dir) print("EXP PATH: {}".format(args.exp_dir)) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args.device = device uncertainty_scores = np.load("uncertainty_scores.npy") with (open( "results/unity_to_trinity_UDA_semi_10/evaluation_trinity_test/confidence_map_UDA.pkl", "rb")) as f: adv_scores = pickle.load(f) index = np.argsort(adv_scores[1])[::-1] ascore = np.empty((2, 8916)) ascore[0] = index ascore[1] = adv_scores[1][index] assert args.checkpoint_seg is not None, "Need trained .pth!" model_seg = load_models( mode="segmentation", device=device, args=args, ) # assert args.checkpoint_disc is not None, "Need trained .pth!" # model_disc = load_models( # mode="single_discriminator", # device=device, # args=args, # ) class FeatureExtractor(torch.nn.Module): def __init__(self, submodule, extracted_layers): super(FeatureExtractor, self).__init__() self.submodule = submodule self.extracted_layers = extracted_layers def forward(self, x): for name, module in self.submodule._modules.items(): x = module(x) print(name) if name in self.extracted_layers: return x['x5'] exact_list = ["pretrained_net"] featExactor = FeatureExtractor(model_seg, exact_list) # a = torch.randn(1, 3, 224, 224) # a = Variable(a).to(args.device) # x = myexactor(a) # print(x) transforms_shape_target = dual_transforms.Compose([ dual_transforms.CenterCrop((400, 400)), dual_transforms.Scale(args.image_size[0]), ]) photo_transformer = PhotometricTransform(photometric_transform_config) traindata = OpenEDSDataset_withLabels( root=os.path.join(args.target_root, "train"), image_size=args.image_size, data_to_train="", shape_transforms=transforms_shape_target, photo_transforms=photo_transformer, train_bool=False, ) train_loader = torch.utils.data.DataLoader( traindata, batch_size=1, shuffle=False, num_workers=0, pin_memory=True, ) trainloader_iter = enumerate(train_loader) args.total_iterations = traindata.__len__() // args.batch_size model_seg.eval() feature_maps = np.empty((traindata.__len__(), 25088)) count = 0 for i_iter in range(args.total_iterations): if i_iter % 1000 == 0: print("Processing {} ..........".format(i_iter)) # if uncertainty_scores[1][i_iter] == uncertainty_scores[1].max(): # print("uncertain", i_iter, traindata.train_data_list[i_iter], uncertainty_scores[1][i_iter], adv_scores[1][i_iter]) # count += 1 # if adv_scores[1][i_iter] == adv_scores[1].max(): # print("adv", i_iter, traindata.train_data_list[i_iter], uncertainty_scores[1][i_iter], adv_scores[1][i_iter]) # count += 1 # # if count==2: # break # u_idx = np.where(uncertainty_scores[0]==i_iter)[0] a_idx = np.where(ascore[0] == i_iter)[0] _, batch = next(trainloader_iter) images, labels = batch images = Variable(images).to(args.device) # labels = Variable(labels.long()).to(args.device) feat = featExactor(images) # feature_maps[i_iter] = feat.view(1,-1)[0].detach().cpu().numpy() feature_maps[a_idx] = feat.view(1, -1)[0].detach().cpu().numpy() # print("Saving feature maps .......................") # np.save("feature_maps_UDA_original_order.npy", feature_maps) # print("Saving feature maps .......................") # np.save("feature_maps_UDA_ascore.npy", feature_maps) from sklearn.metrics.pairwise import cosine_similarity dist = cosine_similarity(feature_maps) # # A = np.matmul(feature_maps.transpose(), feature_maps) # D = A.diagonal() # distance_map = np.power(D, 0.5) * A * np.power(D, -0.5) # print("Saving distance maps .......................") np.save("distance_maps_UDA_ascore.npy", dist)
def main(args): print('===================================\n', ) print("Root directory: {}".format(args.name)) args.exp_dir = os.path.join(RES_DIR, args.name) if not os.path.isdir(args.exp_dir): os.makedirs(args.exp_dir) print("EXP PATH: {}".format(args.exp_dir)) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args.device = device logger = make_logger(filename="TrainVal.log", args=args) if args.tensorboard: writer = SummaryWriter(args.exp_dir) else: writer = None print("===================================") print("====== Loading Training Data ======") print("===================================") shape_transformer = dual_transforms.Compose([ dual_transforms.CenterCrop((400, 400)), dual_transforms.Scale(args.image_size[0]), ]) photo_transformer = PhotometricTransform(photometric_transform_config) joint_data = JointDataset( root_source=args.source_root, image_size=args.image_size, data_to_train="dataloaders/eye/trinity_train_200.pkl", shape_transforms=shape_transformer, photo_transforms=photo_transformer, train_bool=False, ) args.tot_data = joint_data.__len__() args.total_iterations = args.num_epochs * args.tot_data // args.batch_size args.iters_to_eval = args.epoch_to_eval * args.tot_data // args.batch_size print("===================================") print("========= Loading Val Data ========") print("===================================") val_target_data = OpenEDSDataset_withLabels( root=os.path.join(args.target_root, "validation"), image_size=args.image_size, data_to_train="", shape_transforms=shape_transformer, photo_transforms=None, train_bool=False, ) train_loader = torch.utils.data.DataLoader( joint_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) val_loader = torch.utils.data.DataLoader( val_target_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) model_seg = load_models( mode="segmentation", device=device, args=args, ) optimizer_seg = optim.Adam( model_seg.parameters(), lr=args.lr_seg, betas=(args.beta1, 0.999), ) optimizer_seg.zero_grad() # class_weight_target = 1.0 / train_target_data.get_class_probability().to(device) seg_loss_target = torch.nn.CrossEntropyLoss().to(device) trainloader_iter = enumerate(train_loader) val_loss, val_miou = [], [] val_loss_f = float("inf") val_miou_f = float("-inf") loss_seg_min = float("inf") for i_iter in range(args.total_iterations): loss_seg_value = 0 model_seg.train() optimizer_seg.zero_grad() adjust_learning_rate( optimizer=optimizer_seg, learning_rate=args.lr_seg, i_iter=i_iter, max_steps=args.total_iterations, power=0.9, ) try: _, batch = next(trainloader_iter) except StopIteration: trainloader_iter = enumerate(train_loader) _, batch = next(trainloader_iter) images, labels = batch images = Variable(images).to(args.device) labels = Variable(labels.long()).to(args.device) pred = model_seg(images) loss_seg = seg_loss_target(pred, labels) current_loss_seg = loss_seg.item() loss_seg_value += current_loss_seg loss_seg.backward() optimizer_seg.step() logger.info('iter = {0:8d}/{1:8d} ' 'loss_seg = {2:.3f} '.format( i_iter, args.total_iterations, loss_seg_value, )) current_epoch = i_iter * args.batch_size // args.tot_data if i_iter % args.iters_to_eval == 0: val_loss_f, val_miou_f = validate_baseline( i_iter=i_iter, val_loader=val_loader, model=model_seg, epoch=current_epoch, logger=logger, writer=writer, val_loss=val_loss_f, val_iou=val_miou_f, args=args, ) val_loss.append(val_loss_f) val_loss_f = np.min(np.array(val_loss)) val_miou.append(val_miou_f) val_miou_f = np.max(np.array(val_miou)) if args.tensorboard and (writer != None): writer.add_scalar('Val/Cross_Entropy_Target', val_loss_f, i_iter) writer.add_scalar('Val/mIoU_Target', val_miou_f, i_iter) is_better_ss = current_loss_seg < loss_seg_min if is_better_ss: loss_seg_min = current_loss_seg torch.save(model_seg.state_dict(), os.path.join(args.exp_dir, "model_train_best.pth")) logger.info("==========================================") logger.info("Training DONE!") if args.tensorboard and (writer != None): writer.close() with open("%s/train_performance.pkl" % args.exp_dir, "wb") as f: pickle.dump([val_loss, val_miou], f) logger.info("==========================================") logger.info("Evaluating on test data ...") testdata = OpenEDSDataset_withLabels( root=os.path.join(args.target_root, "test"), image_size=args.image_size, data_to_train="", shape_transforms=shape_transformer, photo_transforms=None, train_bool=False, ) test_loader = torch.utils.data.DataLoader( testdata, batch_size=1, shuffle=False, num_workers=0, pin_memory=True, ) pm = run_testing( dataset=testdata, test_loader=test_loader, model=model_seg, args=args, ) logger.info('Global Mean Accuracy: {:.3f}'.format(np.array(pm.GA).mean())) logger.info('Mean IOU: {:.3f}'.format(np.array(pm.IOU).mean())) logger.info('Mean Recall: {:.3f}'.format(np.array(pm.Recall).mean())) logger.info('Mean Precision: {:.3f}'.format(np.array(pm.Precision).mean())) logger.info('Mean F1: {:.3f}'.format(np.array(pm.F1).mean())) IOU_ALL = np.array(pm.Iou_all) logger.info( "Back: {:.4f}, Sclera: {:.4f}, Iris: {:.4f}, Pupil: {:.4f}".format( IOU_ALL[:, 0].mean(), IOU_ALL[:, 1].mean(), IOU_ALL[:, 2].mean(), IOU_ALL[:, 3].mean(), ))
def main(args): print('===================================\n', ) print("Root directory: {}".format(args.name)) args.exp_dir = os.path.join(RES_DIR, args.name) if not os.path.isdir(args.exp_dir): os.makedirs(args.exp_dir) print("EXP PATH: {}".format(args.exp_dir)) assert args.trinity_data_train_with_labels != "", "Indicate trained data in Trinity!" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") args.device = device logger = make_logger(filename="TrainVal.log", args=args) if args.tensorboard: writer = SummaryWriter(args.exp_dir) else: writer = None print("===================================") print("====== Loading Training Data ======") print("===================================") transforms_shape_source = dual_transforms.Compose([ dual_transforms.CenterCrop((400, 400)), dual_transforms.Scale(args.image_size[0]), ]) transforms_shape_target = dual_transforms.Compose([ dual_transforms.CenterCrop((400, 400)), dual_transforms.Scale(args.image_size[0]), ]) photo_transformer = PhotometricTransform(photometric_transform_config) source_data = UnityDataset( root=args.source_root, image_size=args.image_size, shape_transforms=transforms_shape_source, photo_transforms=photo_transformer, train_bool=False, ) target_data = OpenEDSDataset_withLabels( root=os.path.join(args.target_root, "train"), image_size=args.image_size, data_to_train=args.trinity_data_train_with_labels, shape_transforms=transforms_shape_target, photo_transforms=photo_transformer, train_bool=False, ) args.tot_source = source_data.__len__() args.total_iterations = args.num_epochs * source_data.__len__( ) // args.batch_size args.iters_to_eval = args.epoch_to_eval * source_data.__len__( ) // args.batch_size args.iter_source_to_eval = args.epoch_to_eval_source * source_data.__len__( ) // args.batch_size args.iter_semi_start = args.epoch_semi_start * source_data.__len__( ) // args.batch_size print("===================================") print("========= Loading Val Data ========") print("===================================") val_target_data = OpenEDSDataset_withLabels( root=os.path.join(args.target_root, "validation"), image_size=args.image_size, data_to_train="", shape_transforms=transforms_shape_target, photo_transforms=None, train_bool=False, ) # class_weight_source = 1.0 / source_data.get_class_probability().to(device) source_loader = torch.utils.data.DataLoader( source_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) target_loader = torch.utils.data.DataLoader( target_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) val_loader = torch.utils.data.DataLoader( val_target_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, ) model_seg = load_models( mode="segmentation", device=device, args=args, ) model_disc = load_models( mode="single_discriminator", device=device, args=args, ) optimizer_seg = optim.Adam( model_seg.parameters(), lr=args.lr_seg, betas=(args.beta1, 0.999), ) optimizer_seg.zero_grad() optimizer_disc = torch.optim.Adam( model_disc.parameters(), lr=args.lr_disc, betas=(args.beta1, 0.999), ) optimizer_disc.zero_grad() seg_loss_source = torch.nn.CrossEntropyLoss().to(device) gan_loss = torch.nn.BCEWithLogitsLoss().to(device) semi_loss = torch.nn.CrossEntropyLoss(ignore_index=-1) history_true_mask = ImagePool(args.pool_size) history_fake_mask = ImagePool(args.pool_size) trainloader_iter = enumerate(source_loader) targetloader_iter = enumerate(target_loader) val_loss, val_miou = run_training_SSDA( trainloader_source=source_loader, trainloader_target=target_loader, trainloader_iter=trainloader_iter, targetloader_iter=targetloader_iter, val_loader=val_loader, model_seg=model_seg, model_disc=model_disc, gan_loss=gan_loss, seg_loss=seg_loss_source, semi_loss_criterion=semi_loss, optimizer_seg=optimizer_seg, optimizer_disc=optimizer_disc, history_pool_true=history_true_mask, history_pool_fake=history_fake_mask, logger=logger, writer=writer, args=args, ) with open("%s/train_performance.pkl" % args.exp_dir, "wb") as f: pickle.dump([val_loss, val_miou], f) logger.info("==========================================") logger.info("Evaluating on test data ...") testdata = OpenEDSDataset_withLabels( root=os.path.join(args.target_root, "test"), image_size=args.image_size, data_to_train="", shape_transforms=transforms_shape_target, photo_transforms=None, train_bool=False, ) test_loader = torch.utils.data.DataLoader( testdata, batch_size=1, shuffle=False, num_workers=0, pin_memory=True, ) pm = run_testing( dataset=testdata, test_loader=test_loader, model=model_seg, args=args, ) logger.info('Global Mean Accuracy: {:.3f}'.format(np.array(pm.GA).mean())) logger.info('Mean IOU: {:.3f}'.format(np.array(pm.IOU).mean())) logger.info('Mean Recall: {:.3f}'.format(np.array(pm.Recall).mean())) logger.info('Mean Precision: {:.3f}'.format(np.array(pm.Precision).mean())) logger.info('Mean F1: {:.3f}'.format(np.array(pm.F1).mean())) IOU_ALL = np.array(pm.Iou_all) logger.info( "Back: {:.4f}, Sclera: {:.4f}, Iris: {:.4f}, Pupil: {:.4f}".format( IOU_ALL[:, 0].mean(), IOU_ALL[:, 1].mean(), IOU_ALL[:, 2].mean(), IOU_ALL[:, 3].mean(), ))