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 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 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 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 get_model(model_path, model_type='unet11', problem_type='binary'): """ :param model_path: :param model_type: 'UNet', 'UNet16', 'UNet11', 'LinkNet34' :param problem_type: 'binary', 'parts', 'instruments' :return: """ if problem_type == 'binary': num_classes = 1 elif problem_type == 'parts': num_classes = 4 elif problem_type == 'instruments': num_classes = 8 if model_type == 'UNet16': model = UNet16(num_classes=num_classes) elif model_type == 'UNet11': model = UNet11(num_classes=num_classes) elif model_type == 'LinkNet34': model = LinkNet34(num_classes=num_classes) elif model_type == 'UNet': model = UNet(num_classes=num_classes) elif model_type == 'DLinkNet': model = D_LinkNet34(num_classes=num_classes, pretrained=True) 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 start(): parser = argparse.ArgumentParser( description='UNet + BDCLSTM for BraTS Dataset') parser.add_argument('--batch-size', type=int, default=4, metavar='N', help='input batch size for training (default: 4)') parser.add_argument('--test-batch-size', type=int, default=4, metavar='N', help='input batch size for testing (default: 4)') parser.add_argument('--train', action='store_true', default=False, help='Argument to train model (default: False)') parser.add_argument('--epochs', type=int, default=2, metavar='N', help='number of epochs to train (default: 10)') parser.add_argument('--lr', type=float, default=0.001, metavar='LR', help='learning rate (default: 0.01)') parser.add_argument('--cuda', action='store_true', default=False, help='enables CUDA training (default: False)') parser.add_argument('--log-interval', type=int, default=1, metavar='N', help='batches to wait before logging training status') parser.add_argument('--size', type=int, default=128, metavar='N', help='imsize') parser.add_argument('--load', type=str, default=None, metavar='str', help='weight file to load (default: None)') parser.add_argument('--data', type=str, default='./Data/', metavar='str', help='folder that contains data') parser.add_argument('--save', type=str, default='OutMasks', metavar='str', help='Identifier to save npy arrays with') parser.add_argument('--modality', type=str, default='flair', metavar='str', help='Modality to use for training (default: flair)') parser.add_argument('--optimizer', type=str, default='SGD', metavar='str', help='Optimizer (default: SGD)') args = parser.parse_args() args.cuda = args.cuda and torch.cuda.is_available() DATA_FOLDER = args.data # %% Loading in the model # Binary # model = UNet(num_channels=1, num_classes=2) # Multiclass model = UNet(num_channels=1, num_classes=3) if args.cuda: model.cuda() if args.optimizer == 'SGD': optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.99) if args.optimizer == 'ADAM': optimizer = optim.Adam(model.parameters(), lr=args.lr, betas=(args.beta1, args.beta2)) # Defining Loss Function criterion = DICELossMultiClass() if args.train: # %% Loading in the Dataset full_dataset = BraTSDatasetUnet(DATA_FOLDER, im_size=[args.size, args.size], transform=tr.ToTensor()) #dset_test = BraTSDatasetUnet(DATA_FOLDER, train=False, # keywords=[args.modality], im_size=[args.size,args.size], transform=tr.ToTensor()) train_size = int(0.9 * len(full_dataset)) test_size = len(full_dataset) - train_size train_dataset, validation_dataset = torch.utils.data.random_split( full_dataset, [train_size, test_size]) train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=1) validation_loader = DataLoader(validation_dataset, batch_size=args.test_batch_size, shuffle=False, num_workers=1) #test_loader = DataLoader(full_dataset, batch_size=args.test_batch_size, shuffle=False, num_workers=1) print("Training Data : ", len(train_loader.dataset)) print("Validaion Data : ", len(validation_loader.dataset)) #print("Test Data : ", len(test_loader.dataset)) loss_list = [] start = timer() for i in tqdm(range(args.epochs)): train(model, i, loss_list, train_loader, optimizer, criterion, args) test(model, validation_loader, criterion, args, validation=True) end = timer() print("Training completed in {:0.2f}s".format(end - start)) plt.plot(loss_list) plt.title("UNet bs={}, ep={}, lr={}".format(args.batch_size, args.epochs, args.lr)) plt.xlabel("Number of iterations") plt.ylabel("Average DICE loss per batch") plt.savefig("./plots/{}-UNet_Loss_bs={}_ep={}_lr={}.png".format( args.save, args.batch_size, args.epochs, args.lr)) np.save( './npy-files/loss-files/{}-UNet_Loss_bs={}_ep={}_lr={}.npy'.format( args.save, args.batch_size, args.epochs, args.lr), np.asarray(loss_list)) print("Testing Validation") test(model, validation_loader, criterion, args, save_output=True) torch.save( model.state_dict(), 'unet-multiclass-model-{}-{}-{}'.format(args.batch_size, args.epochs, args.lr)) print("Testing PDF images") test_dataset = TestDataset('./pdf_data/', im_size=[args.size, args.size], transform=tr.ToTensor()) test_loader = DataLoader(test_dataset, batch_size=args.test_batch_size, shuffle=False, num_workers=1) print("Test Data : ", len(test_loader.dataset)) test_only(model, test_loader, criterion, args) elif args.load is not None: test_dataset = TestDataset(DATA_FOLDER, im_size=[args.size, args.size], transform=tr.ToTensor()) test_loader = DataLoader(test_dataset, batch_size=args.test_batch_size, shuffle=False, num_workers=1) print("Test Data : ", len(test_loader.dataset)) model.load_state_dict(torch.load(args.load)) test_only(model, test_loader, criterion, args)
# Alpha transperency else: COLOR1 = [255, 0, 0] COLOR2 = [0, 0, 255] #------------------------------------------------------------------------------ # Create model and load weights #------------------------------------------------------------------------------ model = UNet( backbone="mobilenetv2", num_classes=2, pretrained_backbone=None ) if args.use_cuda: model = model.cuda() trained_dict = torch.load(args.checkpoint, map_location="cpu")['state_dict'] model.load_state_dict(trained_dict, strict=False) model.eval() #------------------------------------------------------------------------------ # Predict frames #------------------------------------------------------------------------------ i = 0 while(cap.isOpened()): # Read frame from camera start_time = time() _, frame = cap.read() image = cv2.transpose(frame[...,::-1]) h, w = image.shape[:2]
def train(): t.cuda.set_device(1) # n_channels:医学影像为一通道灰度图 n_classes:二分类 net = UNet(n_channels=1, n_classes=1) optimizer = t.optim.SGD(net.parameters(), lr=opt.learning_rate, momentum=0.9, weight_decay=0.0005) criterion = t.nn.BCELoss() # 二进制交叉熵(适合mask占据图像面积较大的场景) start_epoch = 0 if opt.load_model_path: checkpoint = t.load(opt.load_model_path) # 加载多GPU模型参数到 单模型上 state_dict = checkpoint['net'] new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # remove `module.` new_state_dict[name] = v net.load_state_dict(new_state_dict) # 加载模型 optimizer.load_state_dict(checkpoint['optimizer']) # 加载优化器 start_epoch = checkpoint['epoch'] # 加载训练批次 # 学习率每当到达milestones值则更新参数 if start_epoch == 0: scheduler = t.optim.lr_scheduler.MultiStepLR(optimizer, milestones=opt.milestones, gamma=0.1, last_epoch=-1) # 默认为-1 print('从头训练 ,学习率为{}'.format(optimizer.param_groups[0]['lr'])) else: scheduler = t.optim.lr_scheduler.MultiStepLR(optimizer, milestones=opt.milestones, gamma=0.1, last_epoch=start_epoch) print('加载预训练模型{}并从{}轮开始训练,学习率为{}'.format( opt.load_model_path, start_epoch, optimizer.param_groups[0]['lr'])) # 网络转移到GPU上 if opt.use_gpu: net = t.nn.DataParallel(net, device_ids=opt.device_ids) # 模型转为GPU并行 net.cuda() cudnn.benchmark = True # 定义可视化对象 vis = Visualizer(opt.env) train_data = NodeDataSet(train=True) val_data = NodeDataSet(val=True) test_data = NodeDataSet(test=True) # 数据集加载器 train_dataloader = DataLoader(train_data, opt.batch_size, shuffle=True, num_workers=opt.num_workers) val_dataloader = DataLoader(val_data, opt.batch_size, shuffle=True, num_workers=opt.num_workers) test_dataloader = DataLoader(test_data, opt.test_batch_size, shuffle=False, num_workers=opt.num_workers) for epoch in range(opt.max_epoch - start_epoch): print('开始 epoch {}/{}.'.format(start_epoch + epoch + 1, opt.max_epoch)) epoch_loss = 0 # 每轮判断是否更新学习率 scheduler.step() # 迭代数据集加载器 for ii, (img, mask) in enumerate( train_dataloader): # pytorch0.4写法,不再将tensor封装为Variable # 将数据转到GPU if opt.use_gpu: img = img.cuda() true_masks = mask.cuda() masks_pred = net(img) # 经过sigmoid masks_probs = t.sigmoid(masks_pred) # 损失 = 二进制交叉熵损失 + dice损失 loss = criterion(masks_probs.view(-1), true_masks.view(-1)) # 加入dice损失 if opt.use_dice_loss: loss += dice_loss(masks_probs, true_masks) epoch_loss += loss.item() if ii % 2 == 0: vis.plot('训练集loss', loss.item()) # 优化器梯度清零 optimizer.zero_grad() # 反向传播 loss.backward() # 更新参数 optimizer.step() # 当前时刻的一些信息 vis.log("epoch:{epoch},lr:{lr},loss:{loss}".format( epoch=epoch, loss=loss.item(), lr=optimizer.param_groups[0]['lr'])) vis.plot('每轮epoch的loss均值', epoch_loss / ii) # 保存模型、优化器、当前轮次等 state = { 'net': net.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch } t.save(state, opt.checkpoint_root + '{}_unet.pth'.format(epoch)) # ============验证=================== net.eval() # 评价函数:Dice系数 Dice距离用于度量两个集合的相似性 tot = 0 for jj, (img_val, mask_val) in enumerate(val_dataloader): img_val = img_val true_mask_val = mask_val if opt.use_gpu: img_val = img_val.cuda() true_mask_val = true_mask_val.cuda() mask_pred = net(img_val) mask_pred = (t.sigmoid(mask_pred) > 0.5).float() # 阈值为0.5 # 评价函数:Dice系数 Dice距离用于度量两个集合的相似性 tot += dice_loss(mask_pred, true_mask_val).item() val_dice = tot / jj vis.plot('验证集 Dice损失', val_dice) # ============验证召回率=================== # 每10轮验证一次测试集召回率 if epoch % 10 == 0: result_test = [] for kk, (img_test, mask_test) in enumerate(test_dataloader): # 测试 unet分割能力,故 不使用真值mask if opt.use_gpu: img_test = img_test.cuda() mask_pred_test = net(img_test) # [1,1,512,512] probs = t.sigmoid(mask_pred_test).squeeze().squeeze().cpu( ).detach().numpy() # [512,512] mask = probs > opt.out_threshold result_test.append(mask) # 得到 测试集所有预测掩码,计算二维召回率 vis.plot('测试集二维召回率', getRecall(result_test).getResult()) net.train()
def predict(): net = UNet(n_channels=1, n_classes=1) net.eval() # 将多GPU模型加载为CPU模型 if opt.load_model_path: checkpoint = t.load(opt.load_model_path) state_dict = checkpoint['net'] new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # remove `module.` new_state_dict[name] = v net.load_state_dict(new_state_dict) # 加载模型 print('加载预训练模型{}'.format(opt.load_model_path)) if opt.use_gpu: net.cuda() test_data = NodeDataSet(test=True) test_dataloader = DataLoader(test_data, opt.test_batch_size, shuffle=False, num_workers=opt.num_workers) for ii, full_img in enumerate(test_dataloader): img_test = full_img[0][0].unsqueeze( 0) # 第一个[0] 取 原图像的一个batch,第二个[0]指batch为1 if opt.use_gpu: img_test = img_test.cuda() with t.no_grad(): # pytorch0.4版本写法 output = net(img_test) probs = t.sigmoid(output).squeeze(0) full_mask = probs.squeeze().cpu().numpy() # ===========================================下面方法可能未考虑 一通道图像 # if opt.use_dense_crf: # full_mask = dense_crf(np.array(full_img).astype(np.uint8), full_mask) mask = full_mask > opt.out_threshold # 预测mask值都太小,最大0.01 # # 可视化1 # plt.imsave(opt.save_test_dir+str(10000+ii)+'full_img.jpg', full_img[0][0].squeeze(0),cmap = cm.gray) #保存原图 # plt.imsave(opt.save_test_dir+str(10000+ii)+'mask.jpg', mask,cmap = cm.gray) #保存mask # plt.imsave(opt.save_test_dir+str(10000+ii)+'full_mask.jpg', full_img[0][0].squeeze(0).squeeze(0).numpy() * mask,cmap = cm.gray) #保存mask之后的原图 # 可视化2 # # 多子图显示原图和mask # plt.subplot(1,3,1) # plt.title('origin') # plt.imshow(full_img[0][0].squeeze(0),cmap='Greys_r') # # plt.subplot(1, 3, 2) # plt.title('mask') # plt.imshow(mask,cmap='Greys_r') # # plt.subplot(1, 3, 3) # plt.title('origin_after_mask') # plt.imshow( full_img[0][0].squeeze(0).squeeze(0).numpy() * mask,cmap='Greys_r') # # plt.show() # 保存mask为npy np.save('/home/bobo/data/test/test8/' + str(10000 + ii) + '_mask.npy', mask) print('测试完毕')
class Anonymizer: @classmethod def check_for_numpy(cls, tensor): if isinstance(tensor, torch.Tensor): tensor = tensor.detach().cpu().numpy() return tensor @classmethod def get_number_of_batches(cls, image_paths, batch_size): batches = len(image_paths) / batch_size if not batches.is_integer(): batches = math.floor(batches) + 1 return int(batches) @classmethod def apply_mask(cls, image, mask): return np.multiply(image, mask) @classmethod def anonymize_image(cls, image, mask): image = Anonymizer.check_for_numpy(image) image = image.reshape(image.shape[-2:]) image = np.float32(image) mask = Anonymizer.check_for_numpy(mask) mask = mask.reshape(mask.shape[-2:]) mask = np.uint8(mask) im = Anonymizer.apply_mask(image, mask) cv2.imwrite("sanity_mask.jpg", 255 * mask) cv2.imwrite("sanity_image.jpg", image) cv2.imwrite("sanity_join.jpg", im) mask = Editor.invert_mask(mask) mask = np.uint8(mask) im = cv2.inpaint(im, mask, 10, cv2.INPAINT_TELEA) cv2.imwrite("sanity_anon.jpg", im) return im 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 process_batch(self, batch): # Grab a batch, shuffled according to the provided seed. Note that # i-th image: samples[i][0], i-th mask: samples[i][1] samples = Loader.get_batch(self.image_paths, self.batch_size, batch, None) samples.astype(float) # Cast samples into torch.FloatTensor for interaction with U-Net samples = torch.from_numpy(samples) samples = samples.float() # Cast into a CUDA tensor, if GPUs are available if torch.cuda.is_available(): samples = samples.cuda() # Isolate images and their masks samples_images = samples[:, 0] samples_masks = samples[:, 1] # Reshape for interaction with U-Net samples_images = samples_images.unsqueeze(1) source = samples_images samples_masks = samples_masks.unsqueeze(1) # Run inputs through the model output = self.model(samples_images) # Clamp the target for proper interaction with BCELoss target = torch.clamp(samples_masks, min=0, max=1) del samples return source, output, target def anonymize(self): if not os.path.isdir(self.write_path): print("Making output directory") os.mkdir(self.write_path) count = 0 for batch in range(self.batches): source, output, target = self.process_batch(batch) source = Anonymizer.check_for_numpy(source) source = source.reshape(source.shape[-2:]) source = np.float32(source) binary_mask = Editor.make_binary_mask_from_torch( output[0, :, :, :], 1.0) inverted_binary_mask = Editor.invert_mask( binary_mask) # Now a numpy array instead of torch tensor anonymized_image = Anonymizer.anonymize_image( source, inverted_binary_mask) cv2.imwrite(self.write_path + "/orig_" + str(count) + ".jpg", source) cv2.imwrite(self.write_path + "/anon_" + str(count) + ".jpg", anonymized_image) count += 1 del batch, target, output, binary_mask, anonymized_image def set_cuda(self): if torch.cuda.is_available(): self.model = self.model.cuda() def set_weights(self): if torch.cuda.is_available(): buffered_state_dict = torch.load("weights/" + self.state_dict) else: buffered_state_dict = torch.load( "weights/" + self.state_dict, map_location=lambda storage, loc: storage) self.model.load_state_dict(buffered_state_dict) self.model.eval()
def main(): use_cuda = torch.cuda.is_available() device = torch.device('cuda' if use_cuda else 'cpu') unet = UNet(input_dim, label_dim) unet.load_state_dict( torch.load('./checkpoints_1_19/checkpoint_30.pth', map_location='cpu')) unet.eval() if use_cuda: unet.cuda() criterion = nn.MSELoss() """ if sys.platform.startswith('win'): num_workers = 0 # 0表示不用额外的进程来加速读取数据 else: num_workers = 4 # 读取训练数据集 batch_size = 512 #准备数据 mnist_test_dataset_with_noise = MyMnistDataSet.MyMnistDataSet(root_dir='./mnist_dataset_noise', label_root_dir='./mnist_dataset', type_name='test', transform=transforms.ToTensor()) test_data_loader_with_noise = torch.utils.data.DataLoader(mnist_test_dataset_with_noise, batch_size, shuffle=False, num_workers=num_workers) #遍历数据已有模型进行reference test_loss_sum, batch_count, start_time = 0.0, 0, time.time() for X, y in test_data_loader_with_noise: X = X.to(device) y = y.to(device) y_hat = unet(X) l = criterion(y_hat, y) test_loss_sum += l.cpu().item() batch_count += 1 print('predict: batch_cout %d, test loss %.4f, time %.1f sec' % (batch_count, test_loss_sum / batch_count, time.time() - start_time)) """ transform = transforms.Compose([ transforms.CenterCrop(256), transforms.ToTensor(), ]) dataset = SpectralDataSet( root_dir= '/mnt/liguanlin/DataSets/lowlight_hyperspectral_datasets/band_splited_dataset', type_name='test', transform=transform) dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False) count = 0 mean_psnr = 0 mean_ssim = 0 total_psnr = 0 total_ssim = 0 goundtruth_total_pnsr = 0 goundtruth_total_ssim = 0 for real, labels in tqdm(dataloader): #print('real.shape', real.shape) #print('labels.shape', labels.shape) cur_batch_size = len(real) # Flatten the image real = real.to(device) labels = labels.to(device) pred = unet(real) print(pred.shape) print(real.shape) save_image(pred, 'pred.png', nrow=2) save_image(labels, 'labels.png', nrow=2) pred_numpy = pred.detach().cpu().numpy() label_numpy = labels.detach().cpu().numpy() origin_numpy = real.detach().cpu().numpy() print(pred_numpy.shape) for i in range(pred_numpy.shape[0]): numpy_img = pred_numpy[i].reshape( (pred_numpy.shape[2], pred_numpy.shape[3])) pred_numpy_img = numpy_img * 1023 pred_numpy_img = pred_numpy_img.astype(np.int16) image = Image.fromarray(pred_numpy_img) iamge_name = "./test_results/pred/" + str(i) + ".png" image.save(iamge_name) #pred_numpy_img_8bit = numpy_img * 255 label_numpy_img = label_numpy[i].reshape( (label_numpy.shape[2], label_numpy.shape[3])) label_numpy_img = label_numpy_img * 1023 label_numpy_img = label_numpy_img.astype(np.int16) label_image = Image.fromarray(label_numpy_img) label_iamge_name = "./test_results/label/" + str(i) + ".png" label_image.save(label_iamge_name) origin_numpy_img = origin_numpy[i].reshape( (origin_numpy.shape[2], origin_numpy.shape[3])) origin_numpy_img = origin_numpy_img * 1023 origin_numpy_img = origin_numpy_img.astype(np.int16) origin_image = Image.fromarray(origin_numpy_img) origin_iamge_name = "./test_results/original/" + str(i) + ".png" origin_image.save(origin_iamge_name) count += 1 total_psnr += caculate_psnr_16bit(pred_numpy_img, origin_numpy_img) total_ssim += caculate_ssim_16bit(pred_numpy_img, label_numpy_img) goundtruth_total_pnsr += caculate_psnr_16bit( label_numpy_img, origin_numpy_img) goundtruth_total_ssim += caculate_ssim_16bit( label_numpy_img, pred_numpy_img) if (count == 4): break mean_psnr = total_psnr / count mean_ssim = total_ssim / count print("count = ", count) print("mean_psnr = ", mean_psnr) print("mean_ssim = ", mean_ssim) gound_truth_mean_psnr = goundtruth_total_pnsr / count gound_truth_mean_ssim = goundtruth_total_ssim / count print("gound_truth_mean_psnr = ", gound_truth_mean_psnr) print("gound_truth_mean_ssim = ", gound_truth_mean_ssim)
class Runner(object): 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 _init_device(self, device, out_device) -> str: if device == 'cpu': self.device = torch.device('cpu') self.out_device = torch.device('cpu') self.str_device = 'cpu' return 'cpu' # device type if type(device) == int: device = [device] elif type(device) == str: device = [int(device[-1])] else: # sequence of devices if type(device[0]) == int: device = device else: device = [int(d[-1]) for d in device] # out_device type if type(out_device) == int: out_device = torch.device(f'cuda:{out_device}') else: out_device = torch.device(out_device) self.device = torch.device(f'cuda:{device[0]}') self.out_device = out_device if len(device) > 1: self.model = nn.DataParallel(self.model, device_ids=device, output_device=out_device) self.str_device = ', '.join([f'cuda:{d}' for d in device]) else: self.str_device = str(self.device) self.model.cuda(device[0]) self.criterion.cuda(out_device) if self.sigmoid: self.sigmoid.cuda(device[0]) torch.cuda.set_device(device[0]) return 'cuda' def calc_loss(self, y: Tensor, out: Tensor, Ts: Union[List[int], int]) -> Tensor: """ :param y: (B, T) or (T,) :param out: (B, T) or (T,) :param Ts: length B list or int :return: """ assert self.class_weight is not None weight = (y > 0).float() * self.class_weight[1].item() weight += (y == 0).float() * self.class_weight[0].item() if y.dim() == 1: # if batch_size == 1 y = (y, ) out = (out, ) weight = (weight, ) Ts = (Ts, ) loss = torch.zeros(1, device=self.out_device) for ii, T in enumerate(Ts): loss_no_red = self.criterion(out[ii:ii + 1, ..., :T], y[ii:ii + 1, :T]) loss += (loss_no_red * weight[ii:ii + 1, :T]).sum() / T return loss def predict(self, out_np: ndarray, Ts: Union[List[int], int]) \ -> Tuple[List[ndarray], List]: """ peak-picking prediction :param out_np: (B, T) or (T,) :param Ts: length B list or int :return: boundaries, thresholds boundaries: length B list of boundary interval ndarrays thresholds: length B list of threshold values """ if out_np.ndim == 1: # if batch_size == 1 out_np = (out_np, ) Ts = (Ts, ) boundaries = [] thresholds = [] for item, T in zip(out_np, Ts): candid_idx = [] for idx in range(1, T - 1): i_first = max(idx - self.T_6s, 0) i_last = min(idx + self.T_6s + 1, T) if item[idx] >= np.amax(item[i_first:i_last]): candid_idx.append(idx) boundary_idx = [] threshold = np.mean(item[candid_idx]) for idx in candid_idx: if item[idx] > threshold: boundary_idx.append(idx) boundary_interval = np.array( [[0] + boundary_idx, boundary_idx + [T]], dtype=np.float64).T boundary_interval *= self.frame2time boundaries.append(boundary_interval) thresholds.append(threshold) return boundaries, thresholds @staticmethod def evaluate(reference: Union[List[ndarray], ndarray], prediction: Union[List[ndarray], ndarray]): """ :param reference: length B list of ndarray or just ndarray :param prediction: length B list of ndarray or just ndarray :return: (3,) ndarray """ if isinstance(reference, ndarray): # if batch_size == 1 reference = (reference, ) result = np.zeros(3) for item_truth, item_pred in zip(reference, prediction): mir_result = mir_eval.segment.detection(item_truth, item_pred, trim=True) result += np.array(mir_result) return result # Running model for train, test and validation. def run(self, dataloader, mode: str, epoch: int): self.model.train() if mode == 'train' else self.model.eval() if mode == 'test': state_dict = torch.load(Path(self.writer.logdir, f'{epoch}.pt')) if isinstance(self.model, nn.DataParallel): self.model.module.load_state_dict(state_dict) else: self.model.load_state_dict(state_dict) path_test_result = Path(self.writer.logdir, f'test_{epoch}') os.makedirs(path_test_result, exist_ok=True) else: path_test_result = None avg_loss = 0. avg_eval = 0. all_thresholds = dict() print() pbar = tqdm(dataloader, desc=f'{mode} {epoch:3d}', postfix='-', dynamic_ncols=True) for i_batch, (x, y, intervals, Ts, ids) in enumerate(pbar): # data n_batch = len(Ts) if hasattr(Ts, 'len') else 1 x = x.to(self.device) # B, C, F, T x = dataloader.dataset.normalization.normalize_(x) y = y.to(self.out_device) # B, T # forward out = self.model(x) # B, C, 1, T out = out[..., 0, 0, :] # B, T # loss if mode != 'test': if mode == 'valid': with torch.autograd.detect_anomaly(): loss = self.calc_loss(y, out, Ts) else: loss = self.calc_loss(y, out, Ts) else: loss = 0 out_np = self.sigmoid(out).detach().cpu().numpy() prediction, thresholds = self.predict(out_np, Ts) eval_result = self.evaluate(intervals, prediction) if mode == 'train': # backward self.optimizer.zero_grad() loss.backward() self.optimizer.step() self.scheduler.batch_step() loss = loss.item() elif mode == 'valid': loss = loss.item() if i_batch == 0: # save only the 0-th data id_0, T_0 = ids[0], Ts[0] out_np_0 = out_np[0, :T_0] pred_0, truth_0 = prediction[0][1:, 0], intervals[0][1:, 0] t_axis = np.arange(T_0) * self.frame2time fig = draw_lineplot(t_axis, out_np_0, pred_0, truth_0, id_0) self.writer.add_figure(f'{mode}/out', fig, epoch) np.save(Path(self.writer.logdir, f'{id_0}_{epoch}.npy'), out_np_0) np.save( Path(self.writer.logdir, f'{id_0}_{epoch}_pred.npy'), pred_0) if epoch == 0: np.save(Path(self.writer.logdir, f'{id_0}_truth.npy'), truth_0) else: # save all test data for id_, item_truth, item_pred, item_out, threshold, T \ in zip(ids, intervals, prediction, out_np, thresholds, Ts): np.save(path_test_result / f'{id_}_truth.npy', item_truth) np.save(path_test_result / f'{id_}.npy', item_out[:T]) np.save(path_test_result / f'{id_}_pred.npy', item_pred) all_thresholds[str(id_)] = threshold str_eval = np.array2string(eval_result / n_batch, precision=3) pbar.set_postfix_str(f'{loss / n_batch:.3f}, {str_eval}') avg_loss += loss avg_eval += eval_result avg_loss = avg_loss / len(dataloader.dataset) avg_eval = avg_eval / len(dataloader.dataset) if mode == 'test': np.savez(path_test_result / f'thresholds.npz', **all_thresholds) return avg_loss, avg_eval def step(self, valid_f1: float, epoch: int): """ :param valid_f1: :param epoch: :return: test epoch or 0 """ last_restart = self.scheduler.last_restart self.scheduler.step() # scheduler.last_restart can be updated if epoch == self.scheduler.last_restart: if valid_f1 < self.f1_last_restart: return last_restart else: self.f1_last_restart = valid_f1 torch.save(self.model.module.state_dict(), Path(self.writer.logdir, f'{epoch}.pt')) return 0
class Tester: @classmethod def partition_masks(cls, output, target): # Partition the union of the output and target into a true positive mask, # a false positive mask, and a false negative mask true_positive_mask = torch.min(output, target) false_positive_mask = output - true_positive_mask false_negative_mask = target - true_positive_mask return true_positive_mask, false_positive_mask, false_negative_mask @classmethod def get_partition_measures(cls, output, target): true_positive_mask, false_positive_mask, false_negative_mask = Tester.partition_masks(output, target) tp = torch.sum(true_positive_mask) / (torch.sum(true_positive_mask) + torch.sum(false_positive_mask)) fp = torch.sum(false_positive_mask) / (torch.sum(true_positive_mask) + torch.sum(false_positive_mask)) fn = torch.sum(false_negative_mask) / (torch.sum(true_positive_mask) + torch.sum(false_negative_mask)) return tp, fp, fn @classmethod def get_dice(cls, output, target): tp, fp, fn = Tester.get_partition_measures(output, target) if tp + fp + fn == 0: return -1 dice = (2*tp)/(2*tp + fp + fn) if math.isnan(dice): return 0 return dice.item() @classmethod def get_intersection_over_union(cls, output, target): tp, fp, fn = Tester.get_partition_measures(output, target) if tp + fp + fn == 0: return -1 iou = tp / (tp + fp + fn) if math.isnan(iou): return 0 return iou.item() @classmethod def get_accuracy(cls, output, target): tp, fp, fn = Tester.get_partition_measures(output, target) if tp + fp == 0: return -1 accuracy = tp / (tp + fp) if math.isnan(accuracy): return 0 return accuracy.item() @classmethod def get_recall(cls, output, target): tp, fp, fn = Tester.get_partition_measures(output, target) if tp + fn == 0: return -1 recall = tp / (tp + fn) if math.isnan(recall): return 0 return recall.item() @classmethod def get_number_of_batches(cls, image_paths, batch_size): batches = len(image_paths) / batch_size if not batches.is_integer(): batches = math.floor(batches) + 1 return int(batches) @classmethod def evaluate_loss(cls, criterion, output, target): loss_1 = criterion(output, target) loss_2 = 1 - Tester.get_intersection_over_union(output, target) loss = loss_1 + 0.1 * loss_2 return loss 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 set_cuda(self): if torch.cuda.is_available(): self.model = self.model.cuda() def set_seed(self): if self.seed is not None: np.random.seed(self.seed) def process_batch(self, batch): # Grab a batch, shuffled according to the provided seed. Note that # i-th image: samples[i][0], i-th mask: samples[i][1] samples = Loader.get_batch(self.image_paths, self.batch_size, batch, self.seed) samples.astype(float) # Cast samples into torch.FloatTensor for interaction with U-Net samples = torch.from_numpy(samples) samples = samples.float() # Cast into a CUDA tensor, if GPUs are available if torch.cuda.is_available(): samples = samples.cuda() # Isolate images and their masks samples_images = samples[:, 0] samples_masks = samples[:, 1] # Reshape for interaction with U-Net samples_images = samples_images.unsqueeze(1) samples_masks = samples_masks.unsqueeze(1) # Run inputs through the model output = self.model(samples_images) # Clamp the target for proper interaction with BCELoss target = torch.clamp(samples_masks, min=0, max=1) del samples return output, target def test_model(self): if torch.cuda.is_available(): buffered_state_dict = torch.load("weights/" + self.state_dict) else: buffered_state_dict = torch.load("weights/" + self.state_dict, map_location=lambda storage, loc: storage) self.model.load_state_dict(buffered_state_dict) self.model.eval() criterion = nn.BCELoss() perfect_accuracy_count = 0 zero_accuracy_count = 0 image_count = 0 accuracy_list = [] recall_list = [] iou_list = [] dice_list = [] losses_list = [] for batch in range(self.batches): output, target = self.process_batch(batch) loss = Tester.evaluate_loss(criterion, output, target) print("Batch:", batch) print("~~~~~~~~~~~~~~~~~~~~~~~~~~") print("~~~~~~~~~~~~~~~~~~~~~~~~~~") for i in range(0, output.shape[0]): image_count += 1 binary_mask = Editor.make_binary_mask_from_torch(output[i, :, :, :], 1.0) # Metrics accuracy = Tester.get_accuracy(binary_mask, target[i, :, :, :].cpu()) recall = Tester.get_recall(binary_mask, target[i, :, :, :].cpu()) iou = Tester.get_intersection_over_union(binary_mask, target[i, :, :, :].cpu()) dice = Tester.get_dice(binary_mask, target[i, :, :, :].cpu()) if accuracy == 1: perfect_accuracy_count += 1 if accuracy == 0: zero_accuracy_count += 1 accuracy_list.append(accuracy) recall_list.append(recall) iou_list.append(iou) dice_list.append(dice) print("Accuracy:", accuracy) print("Recall:", recall) print("IoU:", iou) print("Dice:", dice,"\n") print("Mean Accuracy:", mean(accuracy_list)) print("Mean Recall:", mean(recall_list)) print("Mean IoU:", mean(iou_list)) print("Mean Dice:", mean(dice_list)) print("~~~~~~~~~~~~~~~~~~~~~~~~~~") loss_value = loss.item() losses_list.append(loss_value) print("Test loss:", loss_value) print("~~~~~~~~~~~~~~~~~~~~~~~~~~") del output del target mean_iou = mean(iou_list) mean_accuracy = mean(accuracy_list) mean_recall = mean(recall_list) mean_dice = mean(dice_list) print("Perfect Accuracy Percentage:", perfect_accuracy_count / image_count) print("Zero Accuracy Percentage:", zero_accuracy_count / image_count) print("Mean Accuracy:", mean_accuracy) print("Mean Recall:", mean_recall) print("Mean IoU:", mean_iou) print("Mean Dice:", mean_dice)
dset_pred = UnetPred(PRED_INPUT, keywords=[args.modality], im_size=[args.size, args.size], transform=tr.ToTensor()) pred_loader = DataLoader(dset_pred, batch_size=args.test_batch_size, shuffle=False, num_workers=1) print("Prediction Data : ", len(pred_loader.dataset)) # %% Loading in the model model = UNet() if args.cuda: model.cuda() if args.optimizer == 'SGD': optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.99) if args.optimizer == 'ADAM': optimizer = optim.Adam(model.parameters(), lr=args.lr, betas=(args.beta1, args.beta2)) exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1) # Defining Loss Function criterion = DICELossMultiClass() def train(epoch, scheduler, loss_lsit):
class Trainer: @classmethod def intersection_over_union(cls, y, z): iou = (torch.sum(torch.min(y, z))) / (torch.sum(torch.max(y, z))) return iou @classmethod def get_number_of_batches(cls, image_paths, batch_size): batches = len(image_paths) / batch_size if not batches.is_integer(): batches = math.floor(batches) + 1 return int(batches) @classmethod def evaluate_loss(cls, criterion, output, target): loss_1 = criterion(output, target) loss_2 = 1 - Trainer.intersection_over_union(output, target) loss = loss_1 + 0.1 * loss_2 return loss 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 set_cuda(self): if torch.cuda.is_available(): self.model = self.model.cuda() def set_seed(self): if self.seed is not None: np.random.seed(self.seed) def process_batch(self, batch): # Grab a batch, shuffled according to the provided seed. Note that # i-th image: samples[i][0], i-th mask: samples[i][1] samples = Loader.get_batch(self.image_paths, self.batch_size, batch, self.seed) samples.astype(float) # Cast samples into torch.FloatTensor for interaction with U-Net samples = torch.from_numpy(samples) samples = samples.float() # Cast into a CUDA tensor, if GPUs are available if torch.cuda.is_available(): samples = samples.cuda() # Isolate images and their masks samples_images = samples[:, 0] samples_masks = samples[:, 1] # Reshape for interaction with U-Net samples_images = samples_images.unsqueeze(1) samples_masks = samples_masks.unsqueeze(1) # Run inputs through the model output = self.model(samples_images) # Clamp the target for proper interaction with BCELoss target = torch.clamp(samples_masks, min=0, max=1) del samples return output, target def train_model(self): self.model.train() criterion = nn.BCELoss() optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate) iteration = 0 best_iteration = 0 best_loss = 10**10 losses_train = [] losses_val = [] iou_train = [] average_iou_train = [] iou_val = [] average_iou_val = [] print("BEGIN TRAINING") print("TRAINING BATCHES:", self.train_size) print("VALIDATION BATCHES:", self.batches - self.train_size) print("BATCH SIZE:", self.batch_size) print("EPOCHS:", self.epochs) print("~~~~~~~~~~~~~~~~~~~~~~~~~~") for k in range(0, self.epochs): print("EPOCH:", k + 1) print("~~~~~~~~~~~~~~~~~~~~~~~~~~") # Train for batch in range(0, self.train_size): iteration = iteration + 1 output, target = self.process_batch(batch) loss = Trainer.evaluate_loss(criterion, output, target) print("EPOCH:", self.epochs) print("Batch", batch, "of", self.train_size) # Aggregate intersection over union scores for each element in the batch for i in range(0, output.shape[0]): binary_mask = Editor.make_binary_mask_from_torch( output[i, :, :, :], 1.0) iou = Trainer.intersection_over_union( binary_mask, target[i, :, :, :].cpu()) iou_train.append(iou.item()) print("IoU:", iou.item()) # Clear data to prevent memory overload del target del output # Clear gradients, back-propagate, and update weights optimizer.zero_grad() loss.backward() optimizer.step() # Record the loss value loss_value = loss.item() if best_loss > loss_value: best_loss = loss_value best_iteration = iteration losses_train.append(loss_value) if batch == self.train_size - 1: print("LOSS:", loss_value) print("~~~~~~~~~~~~~~~~~~~~~~~~~~") average_iou = sum(iou_train) / len(iou_train) print("Average IoU:", average_iou) average_iou_train.append(average_iou) #Visualizer.save_loss_plot(average_iou_train, "average_iou_train.png") # Validate for batch in range(self.train_size, self.batches): output, target = self.process_batch(batch) loss = Trainer.evaluate_loss(criterion, output, target) for i in range(0, output.shape[0]): binary_mask = Editor.make_binary_mask_from_torch( output[i, :, :, :], 1.0) iou = Trainer.intersection_over_union( binary_mask, target[i, :, :, :].cpu()) iou_val.append(iou.item()) print("IoU:", iou.item()) loss_value = loss.item() losses_val.append(loss_value) print("EPOCH:", self.epochs) print("VALIDATION LOSS:", loss_value) print("~~~~~~~~~~~~~~~~~~~~~~~~~~") del output del target average_iou = sum(iou_val) / len(iou_val) print("Average IoU:", average_iou) average_iou_val.append(average_iou) #Visualizer.save_loss_plot(average_iou_val, "average_iou_val.png") print("Least loss", best_loss, "at iteration", best_iteration) torch.save(self.model.state_dict(), "weights/" + self.state_dict)
def video_infer(args): cap = cv2.VideoCapture(args.video) _, frame = cap.read() H, W = frame.shape[:2] fourcc = cv2.VideoWriter_fourcc(*'DIVX') out = cv2.VideoWriter(args.output, fourcc, 30, (W,H)) font = cv2.FONT_HERSHEY_SIMPLEX # Background if args.bg is not None: BACKGROUND = cv2.imread(args.bg)[...,::-1] BACKGROUND = cv2.resize(BACKGROUND, (W,H), interpolation=cv2.INTER_LINEAR) KERNEL_SZ = 25 SIGMA = 0 # Alpha transperency else: COLOR1 = [90, 140, 154] COLOR2 = [0, 0, 0] if args.model=='unet': model = UNet(backbone=args.net, num_classes=2, pretrained_backbone=None) elif args.model=='deeplabv3_plus': model = DeepLabV3Plus(backbone=args.net, num_classes=2, pretrained_backbone=None) elif args.model=='hrnet': model = HighResolutionNet(num_classes=2, pretrained_backbone=None) if args.use_cuda: model = model.cuda() trained_dict = torch.load(args.checkpoint, map_location="cpu")['state_dict'] model.load_state_dict(trained_dict, strict=False) model.eval() while(cap.isOpened()): start_time = time() ret, frame = cap.read() if ret: image = frame[...,::-1] h, w = image.shape[:2] read_cam_time = time() # Predict mask X, pad_up, pad_left, h_new, w_new = utils.preprocessing(image, expected_size=args.input_sz, pad_value=0) preproc_time = time() with torch.no_grad(): if args.use_cuda: mask = model(X.cuda()) if mask.shape[1] != h_new: mask = F.interpolate(mask, size=(args.input_sz, args.input_sz), mode='bilinear', align_corners=True) mask = mask[..., pad_up: pad_up+h_new, pad_left: pad_left+w_new] mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=True) mask = F.softmax(mask, dim=1) mask = mask[0,1,...].cpu().numpy() else: mask = model(X) mask = mask[..., pad_up: pad_up+h_new, pad_left: pad_left+w_new] mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=True) mask = F.softmax(mask, dim=1) mask = mask[0,1,...].numpy() predict_time = time() # Draw result if args.bg is None: image_alpha = utils.draw_matting(image, mask) #image_alpha = utils.draw_transperency(image, mask, COLOR1, COLOR2) else: image_alpha = utils.draw_fore_to_back(image, mask, BACKGROUND, kernel_sz=KERNEL_SZ, sigma=SIGMA) draw_time = time() # Print runtime read = read_cam_time-start_time preproc = preproc_time-read_cam_time pred = predict_time-preproc_time draw = draw_time-predict_time total = read + preproc + pred + draw fps = 1 / pred print("read: %.3f [s]; preproc: %.3f [s]; pred: %.3f [s]; draw: %.3f [s]; total: %.3f [s]; fps: %.2f [Hz]" % (read, preproc, pred, draw, total, fps)) # Wait for interupt cv2.putText(image_alpha, "%.2f [fps]" % (fps), (10, 50), font, 1.5, (0, 255, 0), 2, cv2.LINE_AA) out.write(image_alpha[..., ::-1]) if args.watch: cv2.imshow('webcam', image_alpha[..., ::-1]) if cv2.waitKey(1) & 0xFF == ord('q'): break else: break
def video_infer(args): cap = cv2.VideoCapture(args.video) _, frame = cap.read() H, W = frame.shape[:2] fps = cap.get(cv2.CAP_PROP_FPS) out = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps, (W, H)) # Background if args.bg is not None: BACKGROUND = cv2.imread(args.bg)[..., ::-1] BACKGROUND = cv2.resize(BACKGROUND, (W, H), interpolation=cv2.INTER_LINEAR) KERNEL_SZ = 25 SIGMA = 0 # Alpha transperency else: COLOR1 = [255, 0, 0] COLOR2 = [0, 0, 255] if args.model == 'unet': model = UNet(backbone=args.net, num_classes=2, pretrained_backbone=None) elif args.model == 'deeplabv3_plus': model = DeepLabV3Plus(backbone=args.net, num_classes=2, pretrained_backbone=None) if args.use_cuda: model = model.cuda() trained_dict = torch.load(args.checkpoint, map_location="cpu")['state_dict'] model.load_state_dict(trained_dict, strict=False) model.eval() if W > H: w_new = int(args.input_sz) h_new = int(H * w_new / W) else: h_new = int(args.input_sz) w_new = int(W * h_new / H) disflow = cv2.DISOpticalFlow_create(cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST) prev_gray = np.zeros((h_new, w_new), np.uint8) prev_cfd = np.zeros((h_new, w_new), np.float32) is_init = True while (cap.isOpened()): start_time = time() ret, frame = cap.read() if ret: image = frame[..., ::-1] h, w = image.shape[:2] read_cam_time = time() # Predict mask X, pad_up, pad_left, h_new, w_new = utils.preprocessing( image, expected_size=args.input_sz, pad_value=0) preproc_time = time() with torch.no_grad(): if args.use_cuda: mask = model(X.cuda()) mask = mask[..., pad_up:pad_up + h_new, pad_left:pad_left + w_new] #mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=True) mask = F.softmax(mask, dim=1) mask = mask[0, 1, ...].cpu().numpy() #(213, 320) else: mask = model(X) mask = mask[..., pad_up:pad_up + h_new, pad_left:pad_left + w_new] #mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=True) mask = F.softmax(mask, dim=1) mask = mask[0, 1, ...].numpy() predict_time = time() # optical tracking cur_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) cur_gray = cv2.resize(cur_gray, (w_new, h_new)) scoremap = 255 * mask optflow_map = postprocess(cur_gray, scoremap, prev_gray, prev_cfd, disflow, is_init) optical_flow_track_time = time() prev_gray = cur_gray.copy() prev_cfd = optflow_map.copy() is_init = False optflow_map = cv2.GaussianBlur(optflow_map, (3, 3), 0) optflow_map = threshold_mask(optflow_map, thresh_bg=0.2, thresh_fg=0.8) img_matting = np.repeat(optflow_map[:, :, np.newaxis], 3, axis=2) bg_im = np.ones_like(img_matting) * 255 re_image = cv2.resize(image, (w_new, h_new)) comb = (img_matting * re_image + (1 - img_matting) * bg_im).astype( np.uint8) comb = cv2.resize(comb, (W, H)) comb = comb[..., ::-1] # Print runtime read = read_cam_time - start_time preproc = preproc_time - read_cam_time pred = predict_time - preproc_time optical = optical_flow_track_time - predict_time total = read + preproc + pred + optical print( "read: %.3f [s]; preproc: %.3f [s]; pred: %.3f [s]; optical: %.3f [s]; total: %.3f [s]; fps: %.2f [Hz]" % (read, preproc, pred, optical, total, 1 / pred)) out.write(comb) if args.watch: cv2.imshow('webcam', comb[..., ::-1]) if cv2.waitKey(1) & 0xFF == ord('q'): break else: break cap.release() out.release()
default='UNet', type=str, help='name of model for saving/loading weights') parser.add_argument('--exp_name', default='single_nuclei_lossless_augmentations', type=str, help='name of experiment for saving files') args = parser.parse_args() LR_SCHED = {0: args.lr, 100: args.lr * 0.1, 150: args.lr * 0.01} torch.cuda.set_device(args.gpu) cudnn.benchmark = True net = UNet() net.cuda() # set model filenames MODEL_CKPT = '../models-pytorch/best_{}_{}.pth'.format(args.model_name, args.exp_name) train_loader, valid_loader, len_train, len_valid = get_cropimg_training_loaders( imsize=args.img_size, test_size=args.valid_size, batch_size=args.batch_size, augment_prob=args.aug_prob, include_rnd=True) msk_crit = nn.BCEWithLogitsLoss().cuda() l1_crit = nn.L1Loss().cuda() optimizer = optim.Adam(net.parameters(), lr=args.lr, weight_decay=args.l2)