def __init__(self) : super(GANCascadedUNet,self).__init__() self.layer_1 = UNet(1,1, segmentation = False) self.layer_2 = UNet(2,4, segmentation = False) self.layer_3 = UNet(6,1, segmentation = False) self.layer_4 = UNet(2,4, segmentation = False) self.layer_5 = UNet(6,1, segmentation = False)
def start(): model = UNet(num_channels=1, num_classes=2) criterion = DICELossMultiClass() dir_path = './Data/images/' sizes = [] file_names = [] for f in os.listdir(dir_path): im = Image.open(os.path.join(dir_path, f)) print(im.size) sizes.append(im.size) file_names.append(f) print(sizes) print(file_names) test_dataset = TestDataset('./Data/', im_size=[256, 256], transform=tr.ToTensor()) test_loader = DataLoader(test_dataset, batch_size=4, shuffle=False, num_workers=1) print("Test Data : ", len(test_loader.dataset)) model.load_state_dict( torch.load('unet-model-16-100-0.001', map_location='cpu')) test_only(model, test_loader, criterion, sizes, file_names)
def load_net_cifar(model_loc): """ Make a model Network must be saved in the form model_name-depth, where this is a unique identifier """ model_file = Path(model_loc).name.rsplit('_')[0] model_name = model_file.split('-')[0] print('Loading model_file', model_file) if (model_name == 'lenet'): model = LeNet(10, 32) elif (model_name == 'vggnet'): model = VGG(int(model_file.split('-')[1]), 10, 32) elif (model_name == 'resnet'): model = ResNet(int(model_file.split('-')[1]), 10, 32) elif (model_name == 'wide'): model = Wide_ResNet( model_file.split('-')[2][0:2], model_file.split('-')[2][2:4], 0, 10, 32) elif (model_name == 'unet'): # this is a terrible way to do it but should be saving class, den parameters to load them individually loss = DenoiseLoss(n=1, hard_mining=0, norm=False) classifier = ResNet(50, num_classes=10, IM_SIZE=32) if (Path(model_loc).name.rsplit('_')[1] == 'model'): denoiser = UNet(3, 3, stochastic=False) elif (Path(model_loc).name.rsplit('_')[1] == 'stochastic'): denoiser = UNet(3, 3, stochastic=True) model = DenoiseNet(classifer=classifier, denoiser=denoiser, loss=loss) else: print( 'Error : Network should be either [LeNet / VGGNet / ResNet / Wide_ResNet / unet' ) sys.exit(0) model.load_state_dict(torch.load(model_loc)['state_dict']) return model
def main(opt): writer = SummaryWriter() log_dir = writer.get_logdir() os.makedirs(os.path.join(log_dir, "images"), exist_ok=True) os.makedirs(os.path.join(log_dir, "test"), exist_ok=True) device = torch.device( "cuda") if torch.cuda.is_available() else torch.device("cpu") # Initialize generator and discriminator generator = UNet(opt.sample_num, opt.channels, opt.batch_size, opt.alpha) discriminator = Discriminator(opt.batch_size, opt.alpha) generator.to(device=device) discriminator.to(device=device) # Optimizers optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr_g, betas=(opt.b1, opt.b2)) optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr_d, betas=(opt.b1, opt.b2)) if opt.mode == 'train': generator = train(writer, log_dir, device, generator, discriminator, optimizer_G, optimizer_D, opt) test(opt, log_dir, generator=generator) if opt.mode == 'test': test(opt, log_dir) test_moving(opt, log_dir)
def main(): # Define the net net = UNet(in_channels=Dataset.in_channels, out_channels=Dataset.out_channels) #.cuda() criterion = DiceLoss() optimizer = optim.Adam(net.parameters(), lr=1e-4) # Parpered the dataset dataloader = torch.utils.data.DataLoader( Dataset('/home/mooziisp/GitRepos/unet/data/membrane/train/image'), batch_size=1, shuffle=True, num_workers=2, pin_memory=True) # TODO validation # TODO acc and loss record # TODO intergated with tensorboard for epoch in range(100): for i, (images, targets) in enumerate(dataloader): #images, targets = images.cuda(), targets.cuda() preds = net(images) loss = criterion(preds, targets) optimizer.zero_grad() loss.backward() optimizer.step() print(f'{epoch*30+i}it, loss: {loss.item():.3f}') return net
def __init__(self, proxy_map, opts,startup_check=False): super(SpecificWorker, self).__init__(proxy_map) self.Period = 50 if startup_check: self.startup_check() else: self.timer.timeout.connect(self.compute) self.timer.start(self.Period) self.opts = opts # Copy user options self.wait_frames = 5 # Wait for these many frames to get stored before calling gaitset (see gaitSet/pretreatment.py) # Dictonary to store mapping between tracking id and segmentation mask self.id2mask = {} # Segmentation module segmentation_model = UNet(backbone="resnet18", num_classes=2) segmentation_model.load_state_dict(torch.load("./src/PretrainedModels/UNet_ResNet18.pth", map_location="cpu")['state_dict'], strict=False) self.segmentation_inference = ImageListInference(segmentation_model,opts) # Gait feature extraction self.gait_model = SetNet(256) self.gait_model.load_state_dict(torch.load("./src/PretrainedModels/GaitSet_CASIA-B_73_False_256_0.2_128_full_30-80000-encoder.ptm",map_location="cpu"),strict=False) if self.opts.gpu >= 0: self.gait_model.cuda() self.gait_model.eval() # Variables to store data for computation self.lock = False # Lock from taking any more input until features are calculated self.input_image_list = [] # Store images in this list before performing segmentation self.input_tracking_id = [] # Store respective tracking id here
def __init__(self, batch_size, image_paths, write_path, state_dict): self.batch_size = batch_size self.image_paths = glob.glob(image_paths) self.batches = Anonymizer.get_number_of_batches( self.image_paths, self.batch_size) self.write_path = write_path self.model = UNet() self.state_dict = state_dict
def __init__(self, side_length, batch_size, seed, image_paths, state_dict): self.side_length = side_length self.batch_size = batch_size self.seed = seed self.image_paths = glob.glob(image_paths) self.batches = Tester.get_number_of_batches(self.image_paths, self.batch_size) self.model = UNet() self.loader = Loader(self.side_length) self.state_dict = state_dict
def get_model(model_path, model_type): """ :param model_path: :param model_type: 'UNet', 'UNet11', 'UNet16', 'AlbuNet34' :return: """ num_classes = 1 if model_type == 'UNet11': model = UNet11(num_classes=num_classes) elif model_type == 'UNet16': model = UNet16(num_classes=num_classes) elif model_type == 'AlbuNet34': model = AlbuNet34(num_classes=num_classes) elif model_type == 'UNet': model = UNet(num_classes=num_classes) else: model = UNet(num_classes=num_classes) state = torch.load(str(model_path)) state = { key.replace('module.', ''): value for key, value in state['model'].items() } model.load_state_dict(state) if torch.cuda.is_available(): return model.cuda() model.eval() return model
def __init__(self, hparams): super(TimeTransfer, self).__init__() self.hparams = hparams # networks self.unet = UNet(3, hparams.hidden_dim) self.data_dir = os.path.expanduser(hparams.data_dir) self.split_indices = prefix_sum(hparams.data_split) self.device = torch.device('cuda' if hparams.gpus > 0 else 'cpu') self.criteria = PerceptualLoss().to(self.device)
def __init__(self, hparams, train_size: int, class_weight: Optional[Tensor] = None): # model, criterion, and prediction self.model = UNet(ch_in=2, ch_out=1, **hparams.model) self.sigmoid = torch.nn.Sigmoid() self.criterion = torch.nn.BCEWithLogitsLoss(reduction='none') self.class_weight = class_weight # for prediction self.frame2time = hparams.hop_size / hparams.sample_rate self.T_6s = round(6 / self.frame2time) - 1 self.T_12s = round(12 / self.frame2time) - 1 self.metrics = ('precision', 'recall', 'F1') # optimizer and scheduler self.optimizer = AdamW( self.model.parameters(), lr=hparams.learning_rate, weight_decay=hparams.weight_decay, ) self.scheduler = CosineLRWithRestarts(self.optimizer, batch_size=hparams.batch_size, epoch_size=train_size, **hparams.scheduler) self.scheduler.step() self.f1_last_restart = -1 # device device_for_summary = self._init_device(hparams.device, hparams.out_device) # summary self.writer = SummaryWriter(logdir=hparams.logdir) path_summary = Path(self.writer.logdir, 'summary.txt') if not path_summary.exists(): print_to_file(path_summary, summary, (self.model, (2, 128, 16 * hparams.model['stride'][1]**4)), dict(device=device_for_summary)) # save hyperparameters path_hparam = Path(self.writer.logdir, 'hparams.txt') if not path_hparam.exists(): with path_hparam.open('w') as f: for var in vars(hparams): value = getattr(hparams, var) print(f'{var}: {value}', file=f)
def test(opt, log_dir, generator=None): device = torch.device( "cuda") if torch.cuda.is_available() else torch.device("cpu") if generator == None: generator = UNet(opt.sample_num, opt.channels, opt.batch_size, opt.alpha) checkpoint = torch.load(opt.load_model, map_location=device) generator.load_state_dict(checkpoint['g_state_dict']) del checkpoint torch.cuda.empty_cache() generator.to(device) generator.eval() dataloader = torch.utils.data.DataLoader(MyDataset_test(opt), opt.batch_size, shuffle=True, num_workers=0) for i, (imgs, filename) in enumerate(dataloader): with torch.no_grad(): test_img = generator(imgs.to(device)) filename = filename[0].split('/')[-1] filename = "test/" + filename + '.png' test_img = convert_im(test_img, os.path.join(log_dir, filename), nrow=5, normalize=True, save_im=True)
def test_moving(opt, log_dir, generator=None): device = torch.device( "cuda") if torch.cuda.is_available() else torch.device("cpu") if generator == None: generator = UNet(opt.sample_num, opt.channels, opt.batch_size, opt.alpha) checkpoint = torch.load(opt.load_model, map_location=device) generator.load_state_dict(checkpoint['g_state_dict']) del checkpoint torch.cuda.empty_cache() generator.to(device) generator.eval() dataloader = torch.utils.data.DataLoader( MyDataset_test_moving(opt), opt.batch_size, shuffle=True, num_workers=opt.num_workers_dataloader) for i, (imgs, filename) in enumerate(dataloader): with torch.no_grad(): filename = filename[0].split('/')[-1] for k in range(len(imgs)): test_img = generator(imgs[k].to(device)) folder_path = os.path.join(log_dir, "test/%s" % filename) os.makedirs(folder_path, exist_ok=True) filename_ = filename + '_' + str(k) + '.png' test_img = convert_im(test_img, os.path.join(folder_path, filename_), nrow=5, normalize=True, save_im=True)
def train(device, model_path, dataset_path): """ Trains the network according on the dataset_path """ network = UNet(1, 3).to(device) optimizer = torch.optim.Adam(network.parameters()) criteria = torch.nn.MSELoss() dataset = GrayColorDataset(dataset_path, transform=train_transform) loader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True, num_workers=cpu_count()) if os.path.exists(model_path): network.load_state_dict(torch.load(model_path)) for _ in tqdm.trange(10, desc="Epoch"): network.train() for gray, color in tqdm.tqdm(loader, desc="Training", leave=False): gray, color = gray.to(device), color.to(device) optimizer.zero_grad() pred_color = network(gray) loss = criteria(pred_color, color) loss.backward() optimizer.step() torch.save(network.state_dict(), model_path)
def denoise_3d(mat, model=None, filepath='/home/jupyter/notebooks/checkpoints/3d_denoise.pt', return_detached=True, batch_size=5000, device=torch.device('cuda')): if model is None: model = UNet(in_channels=1, num_classes=1, out_channels=[4, 8, 16], num_conv=2, n_dim=3, kernel_size=[3, 3, 3], same_shape=True).to(device) model.load_state_dict(torch.load(filepath)) with torch.no_grad(): num_batches = (mat.size(0) + batch_size - 1)//batch_size mat = torch.cat([model(mat[batch_size*i:batch_size*(i+1)]) for i in range(num_batches)], dim=0) if return_detached: mat = mat.detach() for k in [k for k in locals().keys() if k!='mat']: del locals()[k] torch.cuda.empty_cache() return mat
def get_model(m_config, d_config, args): if m_config.unet: kwargs = { 'n_channels': 3, 'dropout': m_config.dropout, 'deeper': not (d_config.dataset == 'CIFAR10' or d_config.dataset == 'CELEBA'), 'ngf': m_config.ngf } return UNet(**kwargs).to(args.device) if d_config.dataset == "FFHQ": model = NCSNv2Deepest elif d_config.dataset == 'LSUN': model = NCSNv2Deeper else: model = NCSNv2 sigmoid = args.target == 'dae' model_args = [m_config, d_config, sigmoid, args.no_dilation, args.std] return model(*model_args).to(args.device)
def get_model(hparams): if hparams['model'] == 'discrete': return C4UNet(hparams['in_channels'], hparams['out_channels'], group_type=hparams['group'], N=hparams['N'], n_features=hparams['n_features'], loss_type=hparams['loss_func'], lr=hparams['lr']) elif hparams['model'] == 'standard': return UNet(hparams['in_channels'], hparams['out_channels'], n_features=hparams['n_features'], loss_type=hparams['loss_func'], lr=hparams['lr']) elif hparams['model'] == 'harmonic': return HarmonicUNet(hparams['in_channels'], hparams['out_channels'], n_features=hparams['n_features'], group_type=hparams['group'], max_freq=hparams['max_freq'], loss_type=hparams['loss_func'], lr=hparams['lr']) elif hparams['model'] == 'steerable': gspace = gspaces.Rot2dOnR2(-1, maximum_frequency=hparams['max_freq']) return SteerableCNN(gspace, hparams['in_channels'], hparams['out_channels'], n_blocks=hparams['n_blocks'], n_features=hparams['n_features'], irrep_type=hparams['irrep_type'], loss_type=hparams['loss_func'], lr=hparams['lr']) else: raise ValueError(f'Unsupported model type: {hparams["model"]}')
def prior_update_gan_noise(particles, inputs, model_no): netG = net.G(5, 3, 64) netG.load_state_dict( torch.load( '%s/netG_epoch_%d.pth' % ('/home/gtx1080/Abduallah/pix2pix.pytorch/imglog/fluid_noise', model_no))) netG.eval() netG.cuda() mu = np.load( '/home/gtx1080/Sync/Kun/30_min_ele/non_ele_whole/train/mu.npy') var = np.load( '/home/gtx1080/Sync/Kun/30_min_ele/non_ele_whole/train/var.npy') particles = np.concatenate((particles, inputs), axis=1) particles, label = normalization_test(particles, particles[:, 0:3, :, :], mu, var) particles = torch.tensor(particles, dtype=torch.float).cuda() with torch.no_grad(): particles = netG(particles) particles = np.array(particles.cpu()) particles[:, 0, :, :] = particles[:, 0, :, :] * np.sqrt(var[5, 0]) + mu[5, 0] particles[:, 1, :, :] = particles[:, 1, :, :] * np.sqrt(var[6, 0]) + mu[6, 0] particles[:, 2, :, :] = particles[:, 2, :, :] * np.sqrt(var[7, 0]) + mu[7, 0] particles[:, 0, :, :] = np.maximum(particles[:, 0, :, :], 0) return particles
def build(config): global train_dataloader global val_dataloader global test_dataloader global model global loss_func global optimizer # ========= Build Data ============== if base_config['dataset'] == 'kaggle': from data import build_kaggle_dataset train_dataloader, val_dataloader, test_dataloader = build_kaggle_dataset( base_config) elif base_config['dataset'] == 'drive': from data import build_drive_dataset train_dataloader, val_dataloader, test_dataloader = build_drive_dataset( base_config) else: _logger.error('{} dataset is not supported now'.format( base_config['dataset'])) # ======== Build Model if config['model'] == 'resnet101': from torchvision.models import resnet101 model = resnet101(num_classes=base_config['n_classes']) elif config['model'] == 'resnext101': from torchvision.models import resnext101_32x8d model = resnext101_32x8d(num_classes=base_config['n_classes']) elif config['model'] == 'densenet': from torchvision.models import densenet121 model = densenet121(num_classes=base_config['n_classes']) elif config['model'] == 'unet': from models import UNet model = UNet(num_classes=base_config['n_classes']) else: _logger.error('{} model is not supported'.format(config['model'])) model = torch.nn.DataParallel(model.cuda()) # Build optimizer if base_config['loss'] == 'ce': loss_func = torch.nn.CrossEntropyLoss().cuda() elif base_config['loss'] == 'bce': loss_func = torch.nn.BCELoss().cuda() elif base_config['loss'] == 'MSE': loss_func = torch.nn.MSELoss().cuda() else: _logger.error('{} loss is not supported'.format(config['loss'])) if config['optimizer'] == 'SGD': optimizer = torch.optim.SGD(model.parameters(), lr=config['lr'], momentum=0.9, weight_decay=5e-4) if config['optimizer'] == 'Adadelta': optimizer = torch.optim.Adadelta(model.parameters(), lr=config['lr']) if config['optimizer'] == 'Adagrad': optimizer = torch.optim.Adagrad(model.parameters(), lr=config['lr']) if config['optimizer'] == 'Adam': optimizer = torch.optim.Adam(model.parameters(), lr=config['lr']) if config['optimizer'] == 'Adamax': optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])
def __init__(self, side_length, batch_size, epochs, learning_rate, momentum_parameter, seed, image_paths, state_dict, train_val_split): self.side_length = side_length self.batch_size = batch_size self.epochs = epochs self.learning_rate = learning_rate self.momentum_parameter = momentum_parameter self.seed = seed self.image_paths = glob.glob(image_paths) self.batches = Trainer.get_number_of_batches(self.image_paths, self.batch_size) self.model = UNet() self.loader = Loader(self.side_length) self.state_dict = state_dict self.train_val_split = train_val_split self.train_size = int(np.floor((self.train_val_split * self.batches)))
def get_model(model_path, model_type): """ :param model_path: :param model_type: 'UNet', 'UNet11', 'UNet16', 'AlbuNet34' :return: """ num_classes = 1 if model_type == 'UNet11': model = UNet11(num_classes=num_classes) elif model_type == 'UNet16': model = UNet16(num_classes=num_classes) elif model_type == 'AlbuNet34': model = AlbuNet34(num_classes=num_classes) elif model_type == 'MDeNet': print('Mine MDeNet..................') model = MDeNet(num_classes=num_classes) elif model_type == 'EncDec': print('Mine EncDec..................') model = EncDec(num_classes=num_classes) elif model_type == 'hourglass': model = hourglass(num_classes=num_classes) elif model_type == 'MDeNetplus': print('load MDeNetplus..................') model = MDeNetplus(num_classes=num_classes) elif model_type == 'UNet': model = UNet(num_classes=num_classes) else: print('I am here') model = UNet(num_classes=num_classes) state = torch.load(str(model_path)) state = { key.replace('module.', ''): value for key, value in state['model'].items() } model.load_state_dict(state) if torch.cuda.is_available(): return model.cuda() model.eval() return model
def get_model(model_path, model_type, num_classes): """ :param model_path: :param model_type: 'UNet', 'UNet16', 'UNet11', 'LinkNet34', :param problem_type: 'binary', 'parts', 'instruments' :return: """ if model_type == 'UNet': model = UNet(num_classes=num_classes) else: model_name = model_list[model_type] model = model_name(num_classes=num_classes) # print(model) state = torch.load(str(model_path)) state = { key.replace('module.', ''): value for key, value in state['model'].items() } model.load_state_dict(state) if torch.cuda.is_available(): return model.cuda() model.eval() return model
def denoise_trace(trace, model=None, filepath='/home/jupyter/notebooks/checkpoints/denoise_trace.pt', return_detached=True, device=torch.device('cuda')): if model is None: model = UNet(in_channels=1, num_classes=1, out_channels=[8, 16, 32], num_conv=2, n_dim=1, kernel_size=3).to(device) model.load_state_dict(torch.load(filepath)) with torch.no_grad(): mean = trace.mean() std = trace.std() pred = model((trace-mean)/std) pred = model(pred) pred = pred * std + mean if return_detached: pred = pred.detach() for k in [k for k in locals().keys() if k!='pred']: del locals()[k] torch.cuda.empty_cache() return pred
def attention_map(mat, model=None, filepath='/home/jupyter/notebooks/checkpoints/segmentation_count_hardmask.pt', batch_size=5000, return_detached=True, device=torch.device('cuda')): if model is None: model = UNet(in_channels=1, num_classes=1, out_channels=[4, 8, 16], num_conv=2, n_dim=3, kernel_size=[3, 3, 3], same_shape=True).to(device) model.load_state_dict(torch.load(filepath)) nrow, ncol = mat.shape[1:] if batch_size*nrow*ncol > 1e7: batch_size = int(1e7 / (nrow*ncol)) with torch.no_grad(): num_batches = (mat.size(0) + batch_size - 1)//batch_size mat = torch.cat([model(mat[batch_size*i:batch_size*(i+1)]) for i in range(num_batches)], dim=0).mean(0) if return_detached: mat = mat.detach() for k in [k for k in locals().keys() if k!='mat']: del locals()[k] torch.cuda.empty_cache() return mat
def __init__(self, train_dl, val_dl): self.device = ('cuda:0' if torch.cuda.is_available() else 'cpu') self.train_dl = train_dl self.val_dl = val_dl self.loss = Loss() self.net = UNet(1).to(self.device) self.net.apply(Model._init_weights) self.criterion = self.loss.BCEDiceLoss self.optim = None self.scheduler = None self._init_optim(LR, BETAS) self.cycles = 0 self.hist = {'train': [], 'val': [], 'loss': []} utils.create_dir('./pt') utils.log_data_to_txt('train_log', f'\nUsing device {self.device}')
def run(conf, data): score = ConfusionMatrix(num_classes) for file in os.listdir(data['splits_dir']): conf['log_key'] = file split = du.load_split_json(data['splits_dir'] + sep + file) model = UNet(conf['input_channels'], conf['num_classes'], reduce_by=2) optimizer = optim.Adam(model.parameters(), lr=conf['learning_rate']) if conf['distribute']: model = torch.nn.DataParallel(model) model.float() optimizer = optim.Adam(model.module.parameters(), lr=conf['learning_rate']) try: trainer = KernelTrainer(run_conf=conf, model=model, optimizer=optimizer) if conf.get('mode') == 'train': train_loader = KernelDataset.get_loader(shuffle=True, mode='train', transforms=transforms, images=split['train'], data_conf=data, run_conf=conf) validation_loader = KernelDataset.get_loader(shuffle=True, mode='validation', transforms=transforms, images=split['validation'], data_conf=data, run_conf=conf) print('### Train Val Batch size:', len(train_loader), len(validation_loader)) # trainer.resume_from_checkpoint(parallel_trained=conf.get('parallel_trained'), key='latest') trainer.train(train_loader=train_loader, validation_loader=validation_loader) test_loader = KernelDataset.get_loader(shuffle=False, mode='test', transforms=transforms, images=split['test'], data_conf=data, run_conf=conf) trainer.resume_from_checkpoint(checkpoint_file= conf.get("checkpoint_file"), parallel_trained=conf.get('parallel_trained'), key='best') trainer.test(data_loader=test_loader, global_score=score, logger=trainer.test_logger) except Exception as e: traceback.print_exc() with open(conf.get('log_dir', 'net_logs') + os.sep + 'global_score.txt', 'w') as lg: lg.write(f'{score.precision()},{score.recall()},{score.f1()},{score.accuracy()}') lg.flush()
def run_pipeline(root_dir, model_path, img_size, batch_size, use_gpu): images_path = os.path.join(root_dir, "images") masks_path = os.path.join(root_dir, "masks") outputs_path = os.path.join(root_dir, "outputs") sizes = [] file_names = [] for f in os.listdir(images_path): im = Image.open(os.path.join(images_path, f)) sizes.append(im.size) file_names.append(f) model = UNet(num_channels=1, num_classes=2) use_gpu = use_gpu and torch.cuda.is_available() if use_gpu: model.cuda() model.load_state_dict(torch.load(model_path, map_location='cpu')) test_dataset = TestDataset(images_path, im_size=[img_size, img_size], transform=tr.ToTensor()) test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=1) print("Test Data : ", len(test_loader.dataset)) start = timer() predict(model, test_loader, batch_size, sizes, file_names, masks_path, use_gpu) end = timer() print("Prediction completed in {:0.2f}s".format(end - start)) generate_bbox(images_path, masks_path, outputs_path) end2 = timer() print("Bbox generation completed in {:0.2f}s".format(end2 - end))
def setup(opts): model = UNet(backbone=opts["backbone"], num_classes=2) if torch.cuda.is_available(): print("Using CUDA") trained_dict = torch.load(opts["checkpoint"])['state_dict'] model.load_state_dict(trained_dict, strict=False) model.cuda() else: print("Using CPU") trained_dict = torch.load(opts["checkpoint"], map_location="cpu")['state_dict'] model.load_state_dict(trained_dict, strict=False) return model
def main(): device = torch.device("cuda") inputPath = "D:/VolumeSuperResolution-InputData/sparse-rendering/cleveland60/sample00004_sparse.npy" input = torch.from_numpy(np.load(inputPath)).to(device=device, dtype=torch.float32) print("input:", input.shape) B, C, H, W = input.shape mask = torch.abs(input[:, 0:1, :, :]) print("Filled pixels:", end='') for b in range(B): print(" %d" % int(torch.sum(mask[b])), end='') print(" of %d pixels" % (H * W)) print( "Now test different u-net configuations and report the filled pixels for the first batch" ) wf = 2 with torch.no_grad(): for depth in range(1, 10): unet = UNet(in_channels=C, depth=depth, wf=wf, padding='partial', return_masks=True) unet.to(device) output, masks = unet(input, mask) print("Depth:", depth) print(" output shape:", output.shape) print(" Layer: shape + filled pixels") for i, m in enumerate(masks): print( " ", m.shape, " -> %d of %d pixels" % (int(torch.sum(m[0])), m.shape[2] * m.shape[3])) print("Save mask") img = mask[0, 0].cpu().numpy() plt.imsave("unet-test-mask.png", img, cmap='Greys')
def build_model(model_type): if model_type == 'unet': model = UNet(num_classes=2) if model_type == 'dcan': model = UNet_DCAN(num_classes=2) if model_type == 'dmtn': model = UNet_DMTN(num_classes=2) if model_type == 'psinet': model = PsiNet(num_classes=2) if model_type == 'convmcd': model = UNet_ConvMCD(num_classes=2) return model