def test(): device = torch.device(args.devices if torch.cuda.is_available() else "cpu") #test_dataset = Training_Dataset(args.test_dir, (args.image_size,args.image_size),(args.noise, args.noise_param)) # test_dataset = HongZhang_Dataset("/data_1/data/Noise2Noise/shenqingbiao/0202", "/data_1/data/Noise2Noise/hongzhang") test_dataset = HongZhang_TestDataset("/data_1/data/红章图片/test/hongzhang", (256, 256)) test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False) # choose the model if args.model == "unet": model = UNet(in_channels=args.image_channels, out_channels=args.image_channels) elif args.model == "srresnet": model = SRResnet(args.image_channels, args.image_channels) elif args.model == "eesp": model = EESPNet_Seg(args.image_channels, 2) else: model = UNet(in_channels=args.image_channels, out_channels=args.image_channels) print('loading model') # model.load_state_dict(torch.load(model_path)) # model.eval() # model.to(device) if args.resume_model: resume_model(model, args.resume_model) model.eval() model.to(device) # result_dir = args.denoised_dir # if not os.path.exists(result_dir): # os.mkdir(result_dir) for batch_idx, image in enumerate(test_loader): #PIL_ShowTensor(torch.squeeze(source)) #PIL_ShowTensor2(torch.squeeze(source),torch.squeeze(noise)) image = image.to(device) denoised_img = model(image).detach().cpu() CV2_showTensors(image.cpu(),denoised_img,timeout=5000)
def main(): logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s') device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Using device {device}') parser = argparse.ArgumentParser() parser.add_argument("--config", type=str, required=True, help="Path to (.yml) config file.") parser.add_argument( "--load-checkpoint", type=str, default="", help="Path to load saved checkpoint from.", ) configargs = parser.parse_args() # Read config file. cfg = None with open(configargs.config, "r") as f: cfg_dict = yaml.load(f, Loader=yaml.FullLoader) print(cfg_dict) # set up network in/out channels details # n_channels=3 for RGB images # n_classes is the number of probabilities you want to get per pixel # - For 1 class and background, use n_classes=1 # - For 2 classes, use n_classes=1 # - For N > 2 classes, use n_classes=N net = UNet(n_channels=1, n_classes=1, bilinear=True) #print(net) logging.info( f'Network:\n' f'\t{net.n_channels} input channels\n' f'\t{net.n_classes} output channels (classes)\n' f'\t{"Bilinear" if net.bilinear else "Transposed conv"} upscaling') net.to(device=device) # set up custome_loss_fn custome_loss_fn = None if cfg_dict.get('loss_function', None) == 'dice_loss': logging.info( f"\n Use custom loss function-{cfg_dict.get('loss_function', None)}" ) custome_loss_fn = DiceLoss # start training TrainNet(Net=net, device=device, root_imgs_dir=cfg_dict.get('base_dir', None), imgs_dir_name=cfg_dict.get("image_dir_suffix", None), mask_dir_name=cfg_dict.get("mask_dir_suffix", None), dir_checkpoint=cfg_dict.get("checkpoint_dir", None), epochs=cfg_dict.get("epochs", 5), batch_size=cfg_dict.get("batch_size", 1), lr=cfg_dict.get("learning_rate", 0.0001), val_percent=cfg_dict.get("validation", 0.2), save_checkpoints=True, img_scale=cfg_dict.get("scale", 1), custome_loss_fn=custome_loss_fn)
def main(args): dataset_kwargs = { 'transforms': {}, 'max_length': None, 'sensor_resolution': None, 'preload_events': False, 'num_bins': 16, 'voxel_method': { 'method': 'random_k_events', 'k': 60000, 't': 0.5, 'sliding_window_w': 500, 'sliding_window_t': 0.1 } } unet_kwargs = { 'base_num_channels': 32, # written as '64' in EVFlowNet tf code 'num_encoders': 4, 'num_residual_blocks': 2, # transition 'num_output_channels': 2, # (x, y) displacement 'skip_type': 'concat', 'norm': None, 'use_upsample_conv': True, 'kernel_size': 3, 'channel_multiplier': 2, 'num_bins': 16 } torch.autograd.set_detect_anomaly(True) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') ev_loader = EventDataLoader(args.h5_file_path, batch_size=1, num_workers=6, shuffle=True, pin_memory=True, dataset_kwargs=dataset_kwargs) H, W = ev_loader.H, ev_loader.W model = UNet(unet_kwargs) model = model.to(device) model.train() crop = CropParameters(W, H, 4) print("=== Let's use", torch.cuda.device_count(), "GPUs!") if torch.cuda.device_count() > 1: model = nn.DataParallel(model) optimizer = torch.optim.Adam(model.parameters(), lr=1e-5, betas=(0.9, 0.999)) # optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, betas=(0.9, 0.999), weight_decay=0.01) # raise # tmp_voxel = crop.pad(torch.randn(1, 9, H, W).to(device)) # F, P = profile(model, inputs=(tmp_voxel, )) for idx in range(10): # for i, item in enumerate(tqdm(ev_loader)): for i, item in enumerate(ev_loader): events = item['events'] voxel = item['voxel'].to(device) voxel = crop.pad(voxel) model.zero_grad() optimizer.zero_grad() flow = model(voxel) * 10 flow = torch.clamp(flow, min=-40, max=40) loss = compute_loss(events, flow) loss.backward() # cvshow_voxel_grid(voxel.squeeze()[0:2].cpu().numpy()) # raise optimizer.step() if i % 10 == 0: print( idx, i, '\t', "{0:.2f}".format(loss.data.item()), "{0:.2f}".format(torch.max(flow[0, 0]).item()), "{0:.2f}".format(torch.min(flow[0, 0]).item()), "{0:.2f}".format(torch.max(flow[0, 1]).item()), "{0:.2f}".format(torch.min(flow[0, 1]).item()), ) xs, ys, ts, ps = events print_voxel = voxel[0].sum(axis=0).cpu().numpy() print_flow = flow[0].clone().detach().cpu().numpy() print_co = warp_events_with_flow_torch( (xs[0][ps[0] == 1], ys[0][ps[0] == 1], ts[0][ps[0] == 1], ps[0][ps[0] == 1]), flow[0].clone().detach(), sensor_size=(H, W)) print_co = crop.pad(print_co) print_co = print_co.cpu().numpy() cvshow_all(idx=idx * 10000 + i, voxel=print_voxel, flow=flow[0].clone().detach().cpu().numpy(), frame=None, compensated=print_co)
logging.info('Created checkpoint directory') except OSError: pass torch.save(net.state_dict(), dir_checkpoint + f'CP_epoch{epoch + 1}.pth') logging.info(f'Checkpoint {epoch + 1} saved !') writer.close() if __name__ == '__main__': device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') net = UNet(n_channels=3, n_classes=1, bilinear=True) net.to(device=device) # faster convolutions, but more memory # cudnn.benchmark = True try: train_net(net=net, epochs=5, batch_size=1, lr=0.001, device=device, val_percent=10.0 / 100) except KeyboardInterrupt: torch.save(net.state_dict(), 'INTERRUPTED.pth') logging.info('Saved interrupt') try: sys.exit(0)
def predict(input_images: list = None, target_images: list = None, config_file: str = None, save_file_suffix=None): ''' Compute output masked and its contours graphs given the "list" of input images filenames. Args: input_images (list[str]): list of input images filenames, if None then input filenames are given by argument list instead target_images (list[str]): list of target images mask filenames, if None then target filenames are given by argument list instead config_file (list[str]): path to the configuation file that specify evaluation detail, if None then config file path are given by argument list instead Returns: out_files (list[str]): list of output maksed filenames countors_outs_files (list[str]): list of countours of output maksed filenames dc_val_records (list[float]): list of dice coefficient of each target masked image and output(predicted) masked image ''' logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s') device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Using device {device}') # specify configuration file if config_file is None: parser = argparse.ArgumentParser() parser.add_argument("--config", type=str, help="Path to (.yml) config file.") parser.add_argument('--input_images', '-i', metavar='INPUT', nargs='+', help='filenames of input images') parser.add_argument('--target_images', '-t', metavar='INPUT', nargs='+', help='filenames of target mask images') configargs = parser.parse_args() config_file = configargs.config # Read config file. cfg = None with open(config_file, "r") as f: cfg_dict = yaml.load(f, Loader=yaml.FullLoader) print(cfg_dict) # set up network in/out channels details # n_channels=3 for RGB images # n_classes is the number of probabilities you want to get per pixel # - For 1 class and background, use n_classes=1 # - For 2 classes, use n_classes=1 # - For N > 2 classes, use n_classes=N net = UNet(n_channels=1, n_classes=1, bilinear=True) #print(net) net.to(device=device) net.load_state_dict( torch.load(cfg_dict.get('model_weights', None), map_location=device)) # In the case of ignoring parameters, input filenames are given by argument list instead if input_images is None: input_images = configargs.input_images if target_images is None: target_images = configargs.target_images logging.info("Model loaded !") out_files = get_output_filenames(in_files=input_images, output_dir=cfg_dict.get( 'output_dir', None), suffix=save_file_suffix) countors_outs_files = [] dc_val_records = [] # start evaluating for i, (filename, target_filename) in enumerate(zip(input_images, target_images)): logging.info( f"\nPredicting image {filename}, Target image {target_filename}") img = Image.open(filename) target = Image.open(target_filename) mask, dc_val = predict_img(net=net, full_img=img, target_img=target, scale_factor=cfg_dict.get('scale', 1), out_threshold=cfg_dict.get( 'mask_threshold', 0.5), device=device) if cfg_dict.get('save', True): out_filename = out_files[i] result, contours = mask_to_image(mask, fn=contours_fn) result.save(out_files[i]) out_contour = out_files[i].replace(".jpg", "-contour.jpg") contours.save(out_contour) countors_outs_files.append(out_contour) # Record DC value for evaluation dc_val_records.append(dc_val) logging.info( f"\nMask saved to {out_files[i]}, Countour saved to {out_contour}" ) return out_files, countors_outs_files, dc_val_records
def train(): # prepare the dataloader device = torch.device(args.devices if torch.cuda.is_available() else "cpu") #dataset = Training_Dataset(args.image_dir, (args.image_size, args.image_size), (args.noise,args.noise_param)) # dataset = HongZhang_Dataset("/data_1/data/Noise2Noise/shenqingbiao/0202", "/data_1/data/Noise2Noise/hongzhang") # dataset = HongZhang_Dataset2("/data_1/data/红章图片", (256, 256)) dataset = HongZhang_Dataset3("/data_1/data/红章图片/6_12", (256, 256)) dataset_length = len(dataset) train_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4) # choose the model if args.model == "unet": model = UNet(in_channels=args.image_channels, out_channels=args.image_channels) elif args.model == "srresnet": model = SRResnet(args.image_channels, args.image_channels) elif args.model == "eesp": model = EESPNet_Seg(args.image_channels, 2) else: model = UNet(in_channels=args.image_channels, out_channels=args.image_channels) model = model.to(device) # choose the loss type if args.loss == "l2": criterion = nn.MSELoss() elif args.loss == "l1": criterion = nn.L1Loss() elif args.loss == "ssim": criterion = SSIM() # resume the mode if needed if args.resume_model: resume_model(model, args.resume_model) optim = Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=True) #scheduler = lr_scheduler.StepLR(optim, step_size=args.scheduler_step, gamma=0.5) scheduler = lr_scheduler.MultiStepLR(optim, milestones=[20, 40], gamma=0.1) model.train() print(model) # start to train print("Starting Training Loop...") since = time.time() for epoch in range(args.epochs): print('Epoch {}/{}'.format(epoch, args.epochs - 1)) print('-' * 10) running_loss = 0.0 scheduler.step() for batch_idx, (target, source) in enumerate(train_loader): source = source.to(device) target = target.to(device) denoised_source = model(source) if args.loss == "ssim": loss = 1 - criterion(denoised_source, Variable(target)) else: loss = criterion(denoised_source, Variable(target)) optim.zero_grad() loss.backward() optim.step() running_loss += loss.item() * source.size(0) if batch_idx % args.steps_show == 0: print('{}/{} Current loss {}'.format(batch_idx, len(train_loader), loss.item())) epoch_loss = running_loss / dataset_length print('{} Loss: {:.4f}'.format('current ' + str(epoch), epoch_loss)) if (epoch + 1) % args.save_per_epoch == 0: save_model(model, epoch + 1) time_elapsed = time.time() - since print('Training complete in {:.0f}m {:.0f}s'.format( time_elapsed // 60, time_elapsed % 60))
if net.n_classes > 1: logging.info( 'Validation cross entropy: {}'.format(val_score)) # writer.add_scalar('Loss/test', val_score, global_step) else: logging.info( 'Validation Dice Coeff: {}'.format(val_score)) # writer.add_scalar('Dice/test', val_score, global_step) if save_cp: try: os.mkdir(dir_checkpoint) logging.info('Created checkpoint directory') except OSError: pass if ((epoch + 1) % 20 == 0): torch.save(net.state_dict(), dir_checkpoint + f'CP_epoch{epoch + 1}.pth') logging.info(f'Checkpoint {epoch + 1} saved !') if __name__ == '__main__': net = UNet(n_channels=3, n_classes=1) device = "cuda:0" net = net.to(device) # net.load_state_dict(torch.load(r"checkpoints/CP_epoch10.pth", map_location=device)) # print(f'Model loaded') # wandb.watch(net) train_net(net, device=device)
def train(train_sources, eval_source): path = sys.argv[1] dr = DataReader(path, train_sources) dr.read() print(len(dr.train.x)) batch_size = 8 device = torch.device('cpu') if torch.cuda.is_available(): device = torch.device('cuda') dataset_s_train = MultiDomainDataset(dr.train.x, dr.train.y, dr.train.vendor, device, DomainAugmentation()) dataset_s_dev = MultiDomainDataset(dr.dev.x, dr.dev.y, dr.dev.vendor, device) dataset_s_test = MultiDomainDataset(dr.test.x, dr.test.y, dr.test.vendor, device) loader_s_train = DataLoader(dataset_s_train, batch_size, shuffle=True) dr_eval = DataReader(path, [eval_source]) dr_eval.read() dataset_eval_dev = MultiDomainDataset(dr_eval.dev.x, dr_eval.dev.y, dr_eval.dev.vendor, device) dataset_eval_test = MultiDomainDataset(dr_eval.test.x, dr_eval.test.y, dr_eval.test.vendor, device) dataset_da_train = MultiDomainDataset(dr.train.x+dr_eval.train.x, dr.train.y+dr_eval.train.y, dr.train.vendor+dr_eval.train.vendor, device, DomainAugmentation()) loader_da_train = DataLoader(dataset_da_train, batch_size, shuffle=True) segmentator = UNet() discriminator = Discriminator(n_domains=len(train_sources)) discriminator.to(device) segmentator.to(device) sigmoid = nn.Sigmoid() selector = Selector() s_criterion = nn.BCELoss() d_criterion = nn.CrossEntropyLoss() s_optimizer = optim.AdamW(segmentator.parameters(), lr=0.0001, weight_decay=0.01) d_optimizer = optim.AdamW(discriminator.parameters(), lr=0.001, weight_decay=0.01) a_optimizer = optim.AdamW(segmentator.encoder.parameters(), lr=0.001, weight_decay=0.01) lmbd = 1/150 s_train_losses = [] s_dev_losses = [] d_train_losses = [] eval_domain_losses = [] train_dices = [] dev_dices = [] eval_dices = [] epochs = 3 da_loader_iter = iter(loader_da_train) for epoch in tqdm(range(epochs)): s_train_loss = 0.0 d_train_loss = 0.0 for index, sample in enumerate(loader_s_train): img = sample['image'] target_mask = sample['target'] da_sample = next(da_loader_iter, None) if epoch == 100: s_optimizer.defaults['lr'] = 0.001 d_optimizer.defaults['lr'] = 0.0001 if da_sample is None: da_loader_iter = iter(loader_da_train) da_sample = next(da_loader_iter, None) if epoch < 50 or epoch >= 100: # Training step of segmentator predicted_activations, inner_repr = segmentator(img) predicted_mask = sigmoid(predicted_activations) s_loss = s_criterion(predicted_mask, target_mask) s_optimizer.zero_grad() s_loss.backward() s_optimizer.step() s_train_loss += s_loss.cpu().detach().numpy() if epoch >= 50: # Training step of discriminator predicted_activations, inner_repr = segmentator(da_sample['image']) predicted_activations = predicted_activations.clone().detach() inner_repr = inner_repr.clone().detach() predicted_vendor = discriminator(predicted_activations, inner_repr) d_loss = d_criterion(predicted_vendor, da_sample['vendor']) d_optimizer.zero_grad() d_loss.backward() d_optimizer.step() d_train_loss += d_loss.cpu().detach().numpy() if epoch >= 100: # adversarial training step predicted_mask, inner_repr = segmentator(da_sample['image']) predicted_vendor = discriminator(predicted_mask, inner_repr) a_loss = -1 * lmbd * d_criterion(predicted_vendor, da_sample['vendor']) a_optimizer.zero_grad() a_loss.backward() a_optimizer.step() lmbd += 1/150 inference_model = nn.Sequential(segmentator, selector, sigmoid) inference_model.to(device) inference_model.eval() d_train_losses.append(d_train_loss / len(loader_s_train)) s_train_losses.append(s_train_loss / len(loader_s_train)) s_dev_losses.append(calculate_loss(dataset_s_dev, inference_model, s_criterion, batch_size)) eval_domain_losses.append(calculate_loss(dataset_eval_dev, inference_model, s_criterion, batch_size)) train_dices.append(calculate_dice(inference_model, dataset_s_train)) dev_dices.append(calculate_dice(inference_model, dataset_s_dev)) eval_dices.append(calculate_dice(inference_model, dataset_eval_dev)) segmentator.train() date_time = datetime.now().strftime("%m%d%Y_%H%M%S") model_path = os.path.join(pathlib.Path(__file__).parent.absolute(), "model", "weights", "segmentator"+str(date_time)+".pth") torch.save(segmentator.state_dict(), model_path) util.plot_data([(s_train_losses, 'train_losses'), (s_dev_losses, 'dev_losses'), (d_train_losses, 'discriminator_losses'), (eval_domain_losses, 'eval_domain_losses')], 'losses.png') util.plot_dice([(train_dices, 'train_dice'), (dev_dices, 'dev_dice'), (eval_dices, 'eval_dice')], 'dices.png') inference_model = nn.Sequential(segmentator, selector, sigmoid) inference_model.to(device) inference_model.eval() print('Dice on annotated: ', calculate_dice(inference_model, dataset_s_test)) print('Dice on unannotated: ', calculate_dice(inference_model, dataset_eval_test))
def train(args): result_path = 'result/%s/'%args.model if not os.path.exists(result_path): os.makedirs(result_path) os.makedirs('%simage'%result_path) os.makedirs('%scheckpoint'%result_path) train_set = MyDataset('train', args.label_type, 512) train_loader = DataLoader( train_set, batch_size=args.batchsize, shuffle=True, num_workers=args.num_workers) # device = 'cuda:0' device = 'cuda:6' if torch.cuda.device_count()>1 else 'cuda:0' out_channels = 1 if args.label_type=='msk' else 2 print(out_channels) if args.model=='unet': print('using unet as model!') model = UNet(out_channels=out_channels) elif args.model=='deeplab': print('using deeplab as model!') model = torch.hub.load('pytorch/vision:v0.9.0', 'deeplabv3_resnet101', pretrained=False) else: print('no model!') model = model.to(device) # model = nn.DataParallel(model) img_show = train_set.__getitem__(0)['x'] img_show = torch.tensor(img_show).to(device).float()[None, :] model.train() criterion = nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) loss_list = [] loss_best = 10 for epo in tqdm(range(1, args.epochs+1), ascii=True): epo_loss = [] for idx, item in enumerate(train_loader): x = item['x'].to(device, dtype=torch.float) y = item['y'].to(device, dtype=torch.float) optimizer.zero_grad() if args.model=='unet': pred = model(x) elif args.model=='deeplab': pred = model(x)['out'][:,0][:,None] # print(y.shape, pred.shape) loss = criterion(pred, y) # print(loss.item()) epo_loss.append(loss.data.item()) loss.backward() optimizer.step() epo_loss_mean = np.array(epo_loss).mean() # print(epo_loss_mean) loss_list.append(epo_loss_mean) plot_loss(loss_list, '%simage/loss.png'%result_path) with torch.no_grad(): if args.model=='unet': pred = model(img_show.clone()) elif args.model=='deeplab': pred = model(img_show.clone())['out'][:,0][:,None] # y = model(img_show) # print(img_show.shape) if args.label_type=='msk': x = img_show[0].cpu().detach().numpy().transpose((1,2,0)) y = pred[0, 0].cpu().detach().numpy() elif args.label_type=='flow': x = img_show[0].cpu().detach().numpy().transpose((1,2,0)) y = pred[0].cpu().detach().numpy().transpose((1,2,0)) plt.subplot(121) plt.imshow(x*255) plt.subplot(122) plt.imshow(y[:,:,0]) plt.savefig('%simage/%d.png'%(result_path, epo)) plt.clf() #loss if epo % 3 ==0: torch.save(model, '%scheckpoint/%d.pt'%(result_path, epo)) if epo_loss_mean < loss_best: loss_best = epo_loss_mean torch.save(model, '%scheckpoint/best.pt'%(result_path)) np.save('%sloss.npy'%result_path, np.array(loss_list))