def main(args): # set GPUS os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu # define where stithced and not stitched test data is located src_sliced_dir = os.path.join(args.test_dir, "images") # prepare output (prediction) dirs dst_sliced_dir = os.path.join(args.dst_dir, "images") os.makedirs(dst_sliced_dir, exist_ok=True) # -------------------------------------------------- # define model # -------------------------------------------------- device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("Available devices:", device) # loading trained models models = [model_from_config(config_path) for config_path in args.configs] # create ensemble model = EnsembleModel(models) # add test time augmentations (flipping and rotating input image) model = ttach.SegmentationTTAWrapper(model, ttach.aliases.d4_transform(), merge_mode='mean') # create Multi GPU model if number of GPUs is more than one n_gpus = len(args.gpu.split(",")) if n_gpus > 1: gpus = list(range(n_gpus)) model = torch.nn.DataParallel(model, gpus) print("Done loading...") model.to(device) # -------------------------------------------------- # start evaluation # -------------------------------------------------- runner = GPUNormRunner(model, model_device=device) model.eval() # predict not stitched data print("Predicting initial files") test_dataset = TestSegmentationDataset(src_sliced_dir, transform_name='test_transform_1', dst_dir=dst_sliced_dir) test_dataloader = torch.utils.data.DataLoader( test_dataset, batch_size=args.batch_size, pin_memory=True, num_workers=8, ) for batch in tqdm(test_dataloader): ids = batch['id'] predictions = runner.predict_on_batch(batch)['mask'] for image_id, prediction in zip(ids, predictions): prediction = prediction.round().int().cpu().numpy().astype("uint8") profile = test_dataset.read_image_profile(image_id) dst_path = os.path.join(dst_sliced_dir, image_id) cv2.imwrite(dst_path, prediction[0])
def main(args): # set GPUS os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu # -------------------------------------------------- # define model # -------------------------------------------------- device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("Available devices:", device) # loading trained models models = [model_from_config(config_path) for config_path in args.configs] # create ensemble model = EnsembleModel(models) # add test time augmentations (flipping and rotating input image) if args.tta: model = ttach.SegmentationTTAWrapper(model, ttach.aliases.d4_transform(), merge_mode='mean') # create Multi GPU model if number of GPUs is more than one n_gpus = len(args.gpu.split(",")) if n_gpus > 1: gpus = list(range(n_gpus)) model = torch.nn.DataParallel(model, gpus) print("Done loading...") model.to(device) # -------------------------------------------------- # start evaluation # -------------------------------------------------- runner = GPUNormRunner(model, model_device=device) model.eval() # predict big tif files predictor = TorchTifPredictor( runner=runner, sample_size=1024, cut_edge=256, batch_size=args.batch_size, count=1, NBITS=1, compress=None, driver="GTiff", blockxsize=256, blockysize=256, tiled=True, BIGTIFF='IF_NEEDED', ) print(f"Inference for {args.src_path}") predictor(args.src_path, args.dst_path)
def predict_image(model, img_crop): with torch.no_grad(): img_crop = torch.from_numpy(img_crop.transpose( 2, 0, 1)).float().unsqueeze(0).cuda() #.to(device) tta_model = tta.SegmentationTTAWrapper(model, transforms_tta) img_out = tta_model(img_crop)[0][0].cpu().detach().numpy() return img_out
def create_models(args): models = [] for encoder_type, ckp in zip(args.encoder_types.split(','), args.ckps.split(',')): model = create_model(encoder_type, ckp=ckp, act='sigmoid')[0].cuda() model = tta.SegmentationTTAWrapper(model, tta.aliases.d4_transform(), merge_mode='mean') if torch.cuda.device_count() > 1: model = nn.DataParallel(model) model.eval() models.append(model) return models
def test_time_aug(net, merge_mode='mean'): """ More operations please assess to this url: https://github.com/qubvel/ttach """ print("Using the test time augmentation! [Default: HorizontalFlip]") trans = tta.Compose([ tta.HorizontalFlip(), # tta.Rotate90(angles=[0, 180]), # tta.Scale(scales=[1, 2]), # tta.Multiply(factors=[0.9, 1, 1.1]), ]) net = tta.SegmentationTTAWrapper(net, trans, merge_mode=merge_mode) return net
def test_main(cfg): # config dataset_cfg = cfg.dataset_cfg test_cfg = cfg.test_cfg device = cfg.device if test_cfg.dataset == 'val_dataset': dataset = LandDataset(DIR=dataset_cfg.val_dir, mode='val', input_channel=dataset_cfg.input_channel, transform=dataset_cfg.val_transform) elif test_cfg.dataset == 'test_dataset': dataset = LandDataset(dataset_cfg.test_dir, mode='test', input_channel=dataset_cfg.input_channel, transform=dataset_cfg.test_transform) else: raise Exception('没有配置数据集!') def _init_fn(): np.random.seed(cfg.random_seed) dataloader = DataLoader(dataset, batch_size=test_cfg.batch_size, shuffle=False, num_workers=test_cfg.num_workers, worker_init_fn=_init_fn()) # 加载模型,预测结果 model = torch.load(test_cfg.check_point_file, map_location=device).to( device) # device参数传在里面,不然默认是先加载到cuda:0,to之后再加载到相应的device上 # model = model.module #并行训练的话需要加上这行 if test_cfg.tta_mode: model = tta.SegmentationTTAWrapper(model, tta.aliases.d4_transform(), merge_mode='mean') # 预测结果 if test_cfg.is_predict: # 获取数据集中样本的序号 sample_index_list = dataset.index_list predict(model=model, dataloader=dataloader, out_dir=test_cfg.out_dir, device=device, sample_index_list=sample_index_list) # 评估模型 if test_cfg.is_evaluate: loss_func = nn.CrossEntropyLoss().to(device) evaluate_model(model, dataloader, loss_func, device, cfg.num_classes)
def init_model(): path = os.path.join(os.path.dirname(__file__), 'net.pth') model = DeepLab(output_stride=16, class_num=17, pretrained=False, bn_momentum=0.1, freeze_bn=False) model.load_state_dict(torch.load(path)) transforms = tta.Compose([ tta.HorizontalFlip(), tta.Rotate90(angles=[0, 180]), ]) model = tta.SegmentationTTAWrapper(model, transforms) model = model.cuda() return model
def test(loader_test, model, args, logger): model = tta.SegmentationTTAWrapper(model, tta.aliases.d4_transform(), merge_mode="mean") model.eval() if args.local_rank == 0: loader_test = tqdm(loader_test, total=cfg.TEST.epoch_iters) with torch.no_grad(): for img, mask, info in loader_test: img = img.cuda() mask = mask.cuda() # import ipdb; ipdb.set_trace() pred = model(img) save_result(info, pred)
def main(): DATA_DIR = './input/test_images_png/' output_dir = './output/' x_valid_dir = DATA_DIR y_valid_dir = DATA_DIR ENCODER = 'inceptionv4' ENCODER_WEIGHTS = 'imagenet' CLASSES = ['coastline'] ACTIVATION = 'sigmoid' # could be None for logits or 'softmax2d' for multicalss segmentation DEVICE = 'cuda' preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS) # load best saved checkpoint best_model = torch.load('./best_model_Unet_resnet18.pth') # tta_model = tta.SegmentationTTAWrapper(best_model, tta.aliases.d4_transform(), merge_mode='mean') transforms = tta.Compose( [ tta.HorizontalFlip(), tta.Rotate90(angles=[0, 180]) ] ) tta_model = tta.SegmentationTTAWrapper(best_model, transforms) # create test dataset test_dataset = Dataset( x_valid_dir, # x_test_dir y_valid_dir, # y_test_dir augmentation=get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn), classes=CLASSES, ) # test_dataloader = DataLoader(test_dataset) # # # evaluate model on test set # test_epoch = smp.utils.train.ValidEpoch( # model=best_model, # loss=loss, # metrics=metrics, # device=DEVICE, # ) # # logs = test_epoch.run(test_dataloader) # test dataset without transformations for image visualization test_dataset_vis2 = Dataset( x_valid_dir, y_valid_dir, # x_test_dir, y_test_dir, augmentation=get_validation_augmentation(), classes=CLASSES, ) for i in range(len(test_dataset)): n = np.random.choice(len(test_dataset)) image_vis = test_dataset_vis2[i][0].astype('uint8') image, gt_mask = test_dataset[i] file_name = test_dataset.images_fps[i] base_name = os.path.basename(file_name) gt_mask = gt_mask.squeeze() x_tensor = torch.from_numpy(image).to(DEVICE).unsqueeze(0) # pr_mask = best_model.predict(x_tensor) pr_mask = tta_model.forward(x_tensor) # pr_mask = (pr_mask.squeeze().cpu().numpy().round()) pr_mask = pr_mask.squeeze().to('cpu').detach().numpy().copy() pr_mask = (pr_mask*255).astype(np.uint8) ret, pr_mask = cv.threshold(pr_mask, 1, 255, cv.THRESH_BINARY) visualize( image=image_vis, ground_truth_mask=gt_mask, predicted_mask=pr_mask ) org_image = cv.imread(file_name) h = org_image.shape[0] w = org_image.shape[1] pr_mask = cv.resize(pr_mask, (w, h)) cv.imwrite(output_dir + base_name, pr_mask)
def testing(num_split, class_params, encoder, decoder): """ 测试推理 """ import gc torch.cuda.empty_cache() gc.collect() sub = "./data/Clouds_Classify/sample_submission.csv" sub = pd.read_csv(open(sub)) sub.head() sub['label'] = sub['Image_Label'].apply(lambda x: x.split('_')[1]) sub['im_id'] = sub['Image_Label'].apply(lambda x: x.split('_')[0]) preprocessing_fn = smp.encoders.get_preprocessing_fn(encoder, 'imagenet') if decoder == 'unet': model = smp.Unet( encoder_name=encoder, encoder_weights='imagenet', classes=4, activation=None, ) else: model = smp.FPN( encoder_name=encoder, encoder_weights='imagenet', classes=4, activation=None, ) test_ids = [id for id in os.listdir(test_imgs_folder)] test_dataset = CloudDataset( df=sub, transforms=get_validation_augmentation(), datatype='test', img_ids=test_ids, preprocessing=get_preprocessing(preprocessing_fn)) test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=2) loaders = {"test": test_loader} logdir = "./logs/log_{}_{}/log_{}".format(encoder, decoder, num_split) encoded_pixels = [] ###############使用pytorch TTA预测#################### use_TTA = True checkpoint_path = logdir + '/checkpoints/best.pth' runner_out = [] model.load_state_dict(torch.load(checkpoint_path)['model_state_dict']) #使用tta预测 if use_TTA: transforms = tta.Compose([ tta.HorizontalFlip(), tta.VerticalFlip(), tta.Scale(scales=[5 / 6, 1, 7 / 6]), ]) tta_model = tta.SegmentationTTAWrapper(model, transforms, merge_mode='mean') else: tta_model = model tta_model = tta_model.cuda() tta_model.eval() with torch.no_grad(): for i, data in enumerate(tqdm.tqdm(loaders['test'])): img, _ = data img = img.cuda() batch_preds = tta_model(img).cpu().numpy() runner_out.extend(batch_preds) runner_out = np.array(runner_out) for i, output in tqdm.tqdm(enumerate(runner_out)): for j, probability in enumerate(output): if probability.shape != (350, 525): probability = cv2.resize(probability, dsize=(525, 350), interpolation=cv2.INTER_LINEAR) logit = sigmoid(probability) predict, num_predict = post_process(logit, class_params[j][0], class_params[j][1]) if num_predict == 0: encoded_pixels.append('') else: r = mask2rle(predict) encoded_pixels.append(r) sub['EncodedPixels'] = encoded_pixels sub.to_csv('./sub/{}_{}/tta_submission_{}.csv'.format( encoder, decoder, num_split), columns=['Image_Label', 'EncodedPixels'], index=False)
def validation(valid_ids, num_split, encoder, decoder): """ 模型验证,并选择后处理参数 """ train = "./data/Clouds_Classify/train.csv" # Data overview train = pd.read_csv(open(train)) train.head() train['label'] = train['Image_Label'].apply(lambda x: x.split('_')[1]) train['im_id'] = train['Image_Label'].apply(lambda x: x.split('_')[0]) ENCODER = encoder ENCODER_WEIGHTS = 'imagenet' if decoder == 'unet': model = smp.Unet( encoder_name=ENCODER, encoder_weights=ENCODER_WEIGHTS, classes=4, activation=None, ) else: model = smp.FPN( encoder_name=ENCODER, encoder_weights=ENCODER_WEIGHTS, classes=4, activation=None, ) preprocessing_fn = smp.encoders.get_preprocessing_fn( ENCODER, ENCODER_WEIGHTS) num_workers = 4 valid_bs = 32 valid_dataset = CloudDataset( df=train, transforms=get_validation_augmentation(), datatype='valid', img_ids=valid_ids, preprocessing=get_preprocessing(preprocessing_fn)) valid_loader = DataLoader(valid_dataset, batch_size=valid_bs, shuffle=False, num_workers=num_workers) loaders = {"valid": valid_loader} logdir = "./logs/log_{}_{}/log_{}".format(encoder, decoder, num_split) valid_masks = [] probabilities = np.zeros((len(valid_ids) * 4, 350, 525)) ############### TTA预测 #################### use_TTA = True checkpoint_path = logdir + '/checkpoints/best.pth' runner_out = [] model.load_state_dict(torch.load(checkpoint_path)['model_state_dict']) if use_TTA: transforms = tta.Compose([ tta.HorizontalFlip(), tta.VerticalFlip(), tta.Scale(scales=[5 / 6, 1, 7 / 6]), ]) tta_model = tta.SegmentationTTAWrapper(model, transforms, merge_mode='mean') else: tta_model = model tta_model = tta_model.cuda() tta_model.eval() with torch.no_grad(): for i, data in enumerate(tqdm.tqdm(loaders['valid'])): img, _ = data img = img.cuda() batch_preds = tta_model(img).cpu().numpy() runner_out.extend(batch_preds) runner_out = np.array(runner_out) ######################END########################## for i, ((_, mask), output) in enumerate(tqdm.tqdm(zip(valid_dataset, runner_out))): for m in mask: if m.shape != (350, 525): m = cv2.resize(m, dsize=(525, 350), interpolation=cv2.INTER_LINEAR) valid_masks.append(m) for j, probability in enumerate(output): if probability.shape != (350, 525): probability = cv2.resize(probability, dsize=(525, 350), interpolation=cv2.INTER_LINEAR) probabilities[i * 4 + j, :, :] = probability # Find optimal values print('searching for optimal param...') params_0 = [[35, 76], [12000, 19001]] params_1 = [[35, 76], [12000, 19001]] params_2 = [[35, 76], [12000, 19001]] params_3 = [[35, 76], [8000, 15001]] param = [params_0, params_1, params_2, params_3] for class_id in range(4): par = param[class_id] attempts = [] for t in range(par[0][0], par[0][1], 5): t /= 100 for ms in range(par[1][0], par[1][1], 2000): masks = [] print('==> searching [class_id:%d threshold:%.3f ms:%d]' % (class_id, t, ms)) for i in tqdm.tqdm(range(class_id, len(probabilities), 4)): probability = probabilities[i] predict, _ = post_process(sigmoid(probability), t, ms) masks.append(predict) d = [] for i, j in zip(masks, valid_masks[class_id::4]): if (i.sum() == 0) & (j.sum() == 0): d.append(1) else: d.append(dice(i, j)) attempts.append((t, ms, np.mean(d))) attempts_df = pd.DataFrame(attempts, columns=['threshold', 'size', 'dice']) attempts_df = attempts_df.sort_values('dice', ascending=False) attempts_df.to_csv( './params/{}_{}_par/params_{}/tta_params_{}.csv'.format( encoder, decoder, num_split, class_id), columns=['threshold', 'size', 'dice'], index=False)
classes=1, activation=None, decoder_attention_type='scse', encoder_depth=5, decoder_channels=[1024, 512, 256, 128, 64], decoder_use_batchnorm=True) net.to(cfg.device) logger.info('load from {}'.format(cfg.test_model_path)) pretrained_dict = torch.load(cfg.test_model_path) net.load_state_dict(pretrained_dict) transforms = tta.Compose([ tta.HorizontalFlip(), tta.VerticalFlip(), # tta.Rotate90(angles=[0, 90, 180, 270]), ]) net = tta.SegmentationTTAWrapper(net, transforms, merge_mode='mean') preprocessing_fn = smp.encoders.get_preprocessing_fn( cfg.ENCODER, 'imagenet') preprocessing = get_preprocessing(preprocessing_fn) result_dir = os.path.join(cfg.result_path, cfg.name) if not os.path.exists(result_dir): os.makedirs(result_dir) cudnn.benchmark = True net.eval() with torch.no_grad(): file_name_list = sorted([(name.split('/')[-1]).split('.')[0] for name in glob( os.path.join(cfg.dataset_dir, cfg.problem,
def main(): parser = argparse.ArgumentParser() parser.add_argument('--encoder', type=str, default='efficientnet-b0') parser.add_argument('--model', type=str, default='unet') parser.add_argument('--loc', type=str) parser.add_argument('--data_folder', type=str, default='../input/') parser.add_argument('--batch_size', type=int, default=2) parser.add_argument('--optimize', type=bool, default=False) parser.add_argument('--tta_pre', type=bool, default=False) parser.add_argument('--tta_post', type=bool, default=False) parser.add_argument('--merge', type=str, default='mean') parser.add_argument('--min_size', type=int, default=10000) parser.add_argument('--thresh', type=float, default=0.5) parser.add_argument('--name', type=str) args = parser.parse_args() encoder = args.encoder model = args.model loc = args.loc data_folder = args.data_folder bs = args.batch_size optimize = args.optimize tta_pre = args.tta_pre tta_post = args.tta_post merge = args.merge min_size = args.min_size thresh = args.thresh name = args.name if model == 'unet': model = smp.Unet(encoder_name=encoder, encoder_weights='imagenet', classes=4, activation=None) if model == 'fpn': model = smp.FPN( encoder_name=encoder, encoder_weights='imagenet', classes=4, activation=None, ) if model == 'pspnet': model = smp.PSPNet( encoder_name=encoder, encoder_weights='imagenet', classes=4, activation=None, ) if model == 'linknet': model = smp.Linknet( encoder_name=encoder, encoder_weights='imagenet', classes=4, activation=None, ) preprocessing_fn = smp.encoders.get_preprocessing_fn(encoder, 'imagenet') test_df = get_dataset(train=False) test_df = prepare_dataset(test_df) test_ids = test_df['Image_Label'].apply( lambda x: x.split('_')[0]).drop_duplicates().values test_dataset = CloudDataset( df=test_df, datatype='test', img_ids=test_ids, transforms=valid1(), preprocessing=get_preprocessing(preprocessing_fn)) test_loader = DataLoader(test_dataset, batch_size=bs, shuffle=False) val_df = get_dataset(train=True) val_df = prepare_dataset(val_df) _, val_ids = get_train_test(val_df) valid_dataset = CloudDataset( df=val_df, datatype='train', img_ids=val_ids, transforms=valid1(), preprocessing=get_preprocessing(preprocessing_fn)) valid_loader = DataLoader(valid_dataset, batch_size=bs, shuffle=False) model.load_state_dict(torch.load(loc)['model_state_dict']) class_params = { 0: (thresh, min_size), 1: (thresh, min_size), 2: (thresh, min_size), 3: (thresh, min_size) } if optimize: print("OPTIMIZING") print(tta_pre) if tta_pre: opt_model = tta.SegmentationTTAWrapper( model, tta.Compose([ tta.HorizontalFlip(), tta.VerticalFlip(), tta.Rotate90(angles=[0, 180]) ]), merge_mode=merge) else: opt_model = model tta_runner = SupervisedRunner() print("INFERRING ON VALID") tta_runner.infer( model=opt_model, loaders={'valid': valid_loader}, callbacks=[InferCallback()], verbose=True, ) valid_masks = [] probabilities = np.zeros((4 * len(valid_dataset), 350, 525)) for i, (batch, output) in enumerate( tqdm( zip(valid_dataset, tta_runner.callbacks[0].predictions["logits"]))): _, mask = batch for m in mask: if m.shape != (350, 525): m = cv2.resize(m, dsize=(525, 350), interpolation=cv2.INTER_LINEAR) valid_masks.append(m) for j, probability in enumerate(output): if probability.shape != (350, 525): probability = cv2.resize(probability, dsize=(525, 350), interpolation=cv2.INTER_LINEAR) probabilities[(i * 4) + j, :, :] = probability print("RUNNING GRID SEARCH") for class_id in range(4): print(class_id) attempts = [] for t in range(30, 70, 5): t /= 100 for ms in [7500, 10000, 12500, 15000, 175000]: masks = [] for i in range(class_id, len(probabilities), 4): probability = probabilities[i] predict, num_predict = post_process( sigmoid(probability), t, ms) masks.append(predict) d = [] for i, j in zip(masks, valid_masks[class_id::4]): if (i.sum() == 0) & (j.sum() == 0): d.append(1) else: d.append(dice(i, j)) attempts.append((t, ms, np.mean(d))) attempts_df = pd.DataFrame(attempts, columns=['threshold', 'size', 'dice']) attempts_df = attempts_df.sort_values('dice', ascending=False) print(attempts_df.head()) best_threshold = attempts_df['threshold'].values[0] best_size = attempts_df['size'].values[0] class_params[class_id] = (best_threshold, best_size) del opt_model del tta_runner del valid_masks del probabilities gc.collect() if tta_post: model = tta.SegmentationTTAWrapper(model, tta.Compose([ tta.HorizontalFlip(), tta.VerticalFlip(), tta.Rotate90(angles=[0, 180]) ]), merge_mode=merge) else: model = model print(tta_post) runner = SupervisedRunner() runner.infer( model=model, loaders={'test': test_loader}, callbacks=[InferCallback()], verbose=True, ) encoded_pixels = [] image_id = 0 for i, image in enumerate(tqdm(runner.callbacks[0].predictions['logits'])): for i, prob in enumerate(image): if prob.shape != (350, 525): prob = cv2.resize(prob, dsize=(525, 350), interpolation=cv2.INTER_LINEAR) predict, num_predict = post_process(sigmoid(prob), class_params[image_id % 4][0], class_params[image_id % 4][1]) if num_predict == 0: encoded_pixels.append('') else: r = mask2rle(predict) encoded_pixels.append(r) image_id += 1 test_df['EncodedPixels'] = encoded_pixels test_df.to_csv(name, columns=['Image_Label', 'EncodedPixels'], index=False)
def main(MODEL_TYPE, LEARNING_RATE, CROP_SIZE, UPSCALE_FACTOR, NUM_EPOCHS, BATCH_SIZE, IMAGE_DIR, LAST_EPOCH, MODEL_NAME='', TV_LOSS_RATE=1e-3): global net, history, lr_img, hr_img, test_lr_img, test_sr_img, valid_hr_img, valid_sr_img, valid_lr_img, scheduler, optimizer train_set = TrainDatasetFromFolder(IMAGE_DIR, crop_size=CROP_SIZE, upscale_factor=UPSCALE_FACTOR) train_loader = DataLoader(dataset=train_set, num_workers=8, batch_size=BATCH_SIZE, shuffle=True) n_iter = (len(train_set) // BATCH_SIZE + 1) * NUM_EPOCHS net = eval(f"{MODEL_TYPE}({UPSCALE_FACTOR})") criterion = TotalLoss(TV_LOSS_RATE) optimizer = optim.RAdam(net.parameters(), lr=LEARNING_RATE) scheduler = schedule.StepLR(optimizer, int(n_iter * 0.3), gamma=0.5, last_epoch=LAST_EPOCH) if LAST_EPOCH == -1: scheduler = GradualWarmupScheduler(optimizer, 1, n_iter // 50, after_scheduler=scheduler) # plot_scheduler(scheduler, n_iter) if MODEL_NAME: net.load_state_dict(torch.load('epochs/' + MODEL_NAME)) print(f"# Loaded model: [{MODEL_NAME}]") print(f'# {MODEL_TYPE} parameters:', sum(param.numel() for param in net.parameters())) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") net.to(device) criterion.to(device) tta_transform = tta.Compose([ tta.HorizontalFlip(), tta.VerticalFlip() ]) # Train model: tta_net = tta.SegmentationTTAWrapper(net, tta_transform) history = [] img_test = plt.imread(r'data\testing_lr_images\09.png') img_valid = plt.imread(r'data\valid_hr_images\t20.png') test_lr_img = torch.from_numpy(img_test.transpose(2, 0, 1)).unsqueeze(0).to(device) valid_hr_img = ToTensor()(img_valid) valid_lr_img = valid_hr_transform(img_valid.shape, UPSCALE_FACTOR)(img_valid).to(device) for epoch in range(1, NUM_EPOCHS + 1): train_bar = tqdm(train_loader) running_results = {'batch_sizes': 0, 'loss': 0, 'psnr': 0} # Train a epoch: net.train() for lr_img, hr_img in train_bar: running_results['batch_sizes'] += BATCH_SIZE optimizer.zero_grad() lr_img, hr_img = MixUp()(lr_img, hr_img, lr_img, hr_img) lr_img = lr_img.type(torch.FloatTensor).to(device) hr_img = hr_img.type(torch.FloatTensor).to(device) sr_img = tta_net(lr_img) loss = criterion(sr_img, hr_img) loss.backward() optimizer.step() scheduler.step() running_results['loss'] += loss.item() * BATCH_SIZE running_results['psnr'] += psnr(hr_img, sr_img) * BATCH_SIZE train_bar.set_description(desc='[%d/%d] Loss: %.4f, PSNR: %.4f' % ( epoch, NUM_EPOCHS, running_results['loss'] / running_results['batch_sizes'], running_results['psnr'] / running_results['batch_sizes'] )) # Save model parameters: psnr_now = running_results['psnr'] / running_results['batch_sizes'] filename = f'epochs/{MODEL_TYPE}_x%d_epoch=%d_PSNR=%.4f.pth' % (UPSCALE_FACTOR, epoch, psnr_now) torch.save(net.state_dict(), filename) history.append(running_results) # Test model: if epoch % 5 == 0: with torch.no_grad(): net.eval() # Plot up-scaled testing image: test_sr_img = net(test_lr_img) plot_hr_lr(test_sr_img, test_lr_img) # Compute PSNR of validation image: valid_sr_img = net(valid_lr_img) psnr_valid = psnr(valid_hr_img, valid_sr_img) # Print PSNR: print('\n' + '-' * 50) print(f"PSNR of Validation = {psnr_valid}") print('-' * 50 + '\n') torch.save(optimizer.state_dict(), f'optimizer_{MODEL_TYPE}_epoch={NUM_EPOCHS}.pth') torch.save(scheduler.state_dict(), f'scheduler_{MODEL_TYPE}_epoch={NUM_EPOCHS}.pth')
file=f) optimizer.param_groups[0]['lr'] *= 0.97 if max_score < valid_logs['iou_score']: max_score = valid_logs['iou_score'] torch.save(model, f"models/saved/{config['model_name']}.pth") print('Model saved!', file=f) patience = 0 else: patience += 1 if patience == config['patience']: break tta_model = tta.SegmentationTTAWrapper( torch.load(f"models/saved/{config['model_name']}.pth"), tta.aliases.d4_transform(), merge_mode='mean') original_res = [] res = [] for file in df['Id'].values: ds = gdal.Open(f"{config['images_path']}{file}.tif") IMG1 = np.array([ds.GetRasterBand(i).ReadAsArray() for i in range(1, 4)]) IMG2 = np.array([ds.GetRasterBand(i).ReadAsArray() for i in range(5, 8)]) IMG1, IMG2 = normalize(IMG1, IMG2, config['img_channels']) new_shape = generate_new_shape(IMG1, config['img_size'], config['img_channels']) res_mask = np.zeros((new_shape[0], new_shape[1])) if config['img_channels'] == 3:
def predict(best_threshold, min_size, device, transforms): test_dataset = SegmentationDataset(data_folder=config_main['path_to_data'], transforms=AUGMENTATIONS_TEST, phase='test') test_loader = DataLoader(test_dataset, batch_size=config['batch_size'], shuffle=False, num_workers=16, drop_last=False) models = [] for weight in glob.glob(os.path.join(config['weights'], config['name'], 'cosine/') + "*.pth"): model = pydoc.locate(config['model'])(**config['model_params']) model.load_state_dict(torch.load(weight)) model = model.to(device) model.eval() models.append(model) if len(config['cls_predict_test']) > 0: print("Use classification model results") cls_df = pd.read_csv(config['cls_predict_test']) cls_df['is_mask_empty'] = cls_df['EncodedPixels'].map(lambda x: 1 if x==0 else 0) cls_df.index = cls_df.Image_Label.values cls_df.drop_duplicates(inplace=True) else: cls_df = None predictions = [] image_names = [] with torch.no_grad(): for i, batch in enumerate(tqdm(test_loader)): fnames = batch["filename"] images = batch["image"].to(device) batch_preds = np.zeros((images.size(0), 4, TRAIN_SHAPE[0], TRAIN_SHAPE[1]), dtype=np.float32) batch_preds_test_shape = np.zeros((images.size(0), 4, TEST_SHAPE[0], TEST_SHAPE[1]), dtype=np.float32) if config['type'] == 'crop': for model in models: if config['TTA'] == 'true': model = tta.SegmentationTTAWrapper(model, transforms) tmp_batch_preds = np.zeros((images.size(0), 4, TRAIN_SHAPE[0], TRAIN_SHAPE[1]), dtype=np.float32) for step in np.arange(0, TRAIN_SHAPE[1], 384)[:-1]: tmp_pred = torch.sigmoid(model(images[:, :, :, step:step + 448])).cpu().numpy() tmp_batch_preds[:, :, :, step:step + 448] += tmp_pred tmp_batch_preds[:, :, :, 384:384 + 64] /= 2 tmp_batch_preds[:, :, :, 2 * 384:2 * 384 + 64] /= 2 tmp_batch_preds[:, :, :, 3 * 384:3 * 384 + 64] /= 2 batch_preds += tmp_batch_preds else: for model in models: if config['TTA'] == 'true': model = tta.SegmentationTTAWrapper(model, transforms) batch_preds += torch.sigmoid(model(images)).cpu().numpy() batch_preds = batch_preds / len(models) for i in range(batch_preds.shape[0]): tmp = cv2.resize(np.moveaxis(batch_preds[i], 0, -1), (TEST_SHAPE[1], TEST_SHAPE[0])) batch_preds_test_shape[i] = np.moveaxis(tmp, -1, 0) for fname, preds in zip(fnames, batch_preds_test_shape): for cls, pred in enumerate(preds): if cls_df is not None: if cls_df.loc[fname + f"_{inv_map[cls]}"]['is_mask_empty'] == 1: pred = np.zeros((TEST_SHAPE[0], TEST_SHAPE[1])) else: pred = post_process(pred, best_threshold, min_size, cls, use_dense_crf=config['use_dense_crf'], image=cv2.imread(test_dataset.images[i]) if config['use_dense_crf']=='true' else None, use_dilations=config['use_dilations'], use_poligonization=config['use_poligonization']) else: pred = post_process(pred, best_threshold, min_size, cls, use_dense_crf=config['use_dense_crf'], image=cv2.imread(test_dataset.images[i]) if config['use_dense_crf']=='true' else None, use_dilations=config['use_dilations'], use_poligonization=config['use_poligonization']) rle = mask2rle(pred) name = fname + f"_{inv_map[cls]}" image_names.append(name) predictions.append(rle) df = pd.DataFrame() df["Image_Label"] = image_names df["EncodedPixels"] = predictions df.to_csv(os.path.join(config['weights'], config['name'], "submission.csv"), index=False)
def search_threshold(device, transforms): val_dataset = SegmentationDataset(data_folder=config_main['path_to_data'], transforms=AUGMENTATIONS_TEST, phase='val', fold=config['fold'], activation=config_main['activation']) if len(config['cls_predict_val']) > 0: val_dataset.start_value = 0.1 val_dataset.delta = 0.0 val_dataset.update_empty_mask_ratio(0) val_loader = DataLoader(val_dataset, batch_size=config['batch_size'], shuffle=False, num_workers=16, drop_last=False) if len(config['cls_predict_val']) > 0: print("Use classification model results") cls_df = pd.read_csv(config['cls_predict_val']) cls_df['is_mask_empty'] = cls_df['label'].map(lambda x: 1 if x==0 else 0) cls_df.index = cls_df.Image_Label.values cls_df.drop_duplicates(inplace=True) else: cls_df = None models = [] for weight in glob.glob(os.path.join(config['weights'], config['name'], 'cosine/') + "*.pth"): model = pydoc.locate(config['model'])(**config['model_params']) model.load_state_dict(torch.load(weight)) model = model.to(device) model.eval() models.append(model) print(f"Use {len(models)} models.") assert len(models) > 0, "Models not loaded" masks, predicts = [], [] with torch.no_grad(): for i, batch in enumerate(tqdm(val_loader)): fnames = batch["filename"] images = batch["image"].to(device) mask = batch['mask'].cpu().numpy() mask_pred_shape = np.zeros((images.size(0), 4, TEST_SHAPE[0], TEST_SHAPE[1]), dtype=np.float32) batch_preds = np.zeros((images.size(0), 4, TRAIN_SHAPE[0], TRAIN_SHAPE[1]), dtype=np.float32) batch_preds_test_shape = np.zeros((images.size(0), 4, TEST_SHAPE[0], TEST_SHAPE[1]), dtype=np.float32) if config['type'] == 'crop': for model in models: if config['TTA'] == 'true': model = tta.SegmentationTTAWrapper(model, transforms) tmp_batch_preds = np.zeros((images.size(0), 4, TRAIN_SHAPE[0], TRAIN_SHAPE[1]), dtype=np.float32) for step in np.arange(0, TRAIN_SHAPE[1], 384)[:-1]: tmp_pred = torch.sigmoid(model(images[:,:,:,step:step+448])).cpu().numpy() tmp_batch_preds[:,:,:,step:step+448] += tmp_pred tmp_batch_preds[:,:,:,384:384+64] /= 2 tmp_batch_preds[:,:,:,2*384:2*384+64] /= 2 tmp_batch_preds[:,:,:,3*384:3*384+64] /= 2 batch_preds += tmp_batch_preds else: for model in models: if config['TTA'] == 'true': model = tta.SegmentationTTAWrapper(model, transforms) batch_preds += torch.sigmoid(model(images)).cpu().numpy() batch_preds = batch_preds / len(models) for i in range(batch_preds.shape[0]): tmp = cv2.resize(np.moveaxis(batch_preds[i], 0, -1), (TEST_SHAPE[1], TEST_SHAPE[0])) batch_preds_test_shape[i] = np.moveaxis(tmp, -1, 0) tmp = cv2.resize(np.moveaxis(mask[i], 0, -1), (TEST_SHAPE[1], TEST_SHAPE[0])) mask_pred_shape[i] = np.moveaxis(tmp, -1, 0) for num_file in range(batch_preds_test_shape.shape[0]): for cls in range(config_main['n_classes']): if cls_df is not None: if cls_df.loc[fnames[num_file] + f"_{inv_map[cls]}"]['is_mask_empty'] == 1: batch_preds_test_shape[num_file, cls] = np.zeros((TEST_SHAPE[0], TEST_SHAPE[1])) predicts.append(batch_preds_test_shape) masks.append(mask_pred_shape) predicts = np.vstack(predicts) masks = np.vstack(masks) print("Search threshold ...") thresholds = np.arange(0.1, 1.0, 0.1) if config['channel_threshold'] == 'true': best_threshold = [] for channel in range(4): scores = [] for threshold in tqdm(thresholds): score = dice_coef_numpy(preds=(predicts>threshold).astype(int), trues=masks, channel=channel) print(f"{threshold} - {score}") scores.append(score) best_score = np.max(scores) print(f"Best threshold - {thresholds[np.argmax(scores)]}, best score - {best_score}") print(f"Scores: {scores}") best_threshold.append(thresholds[np.argmax(scores)]) print(f"Best thresholds - {best_threshold}") else: scores = [] for threshold in tqdm(thresholds): score = dice_coef_numpy(preds=(predicts > threshold).astype(int), trues=masks) print(f"{threshold} - {score}") scores.append(score) best_score = np.max(scores) best_threshold = thresholds[np.argmax(scores)] print(f"Best threshold - {best_threshold}, best score - {best_score}") print(f"Scores: {scores}") #best_threshold = [0.8, 0.8, 0.9, 0.7] print("Search min_size threshold ...") thresholds = np.arange(0, 20000, 1000) scores = [] for threshold in tqdm(thresholds): tmp = predicts.copy() for i in range(tmp.shape[0]): for j in range(tmp.shape[1]): tmp[i,j] = post_process(tmp[i,j], best_threshold, threshold, j, use_dense_crf=config['use_dense_crf'], image=cv2.imread(val_dataset.images[i]) if config['use_dense_crf']=='true' else None, use_dilations=config['use_dilations'], use_poligonization=config['use_poligonization']) score = dice_coef_numpy(preds=tmp, trues=masks) print(f"{threshold} - {score}") scores.append(score) best_score = np.max(scores) best_min_size_threshold = thresholds[np.argmax(scores)] print(f"Best min_size threshold - {best_min_size_threshold}, best score - {best_score}") print(f"Scores: {scores}") return best_threshold, best_min_size_threshold
def predict_test(CropStage=False, TestStage=True, toMask=True, toZip=True, newTH=0.05): # os.environ["CUDA_VISIBLE_DIVICES"] ="1" root_path = '/media/totem_disk/totem/weitang/project' # model = smp.Unet('se_resnext101_32x4d', activation=None).cuda() # i_size=512 # i_scale=0.25 # dir_model = root_path + '/model/unet_se_resnext101_32x4d_2_1_best.pth' # model.load_state_dict(torch.load(dir_model)['state_dict']) model = smp.Unet('densenet161', activation=None).cuda() i_size = 512 i_scale = 0.25 dir_model = root_path + '/model/unet_densenet161_2_1_best_0.73.pth' model.load_state_dict(torch.load(dir_model)['state_dict']) tta_transforms = tta.Compose([ tta.HorizontalFlip(), # tta.Scale(scales=[1,2,4]) # tta.Rotate90(angles=[0,180]) ]) tta_model = tta.SegmentationTTAWrapper(model, tta_transforms, merge_mode='mean') # model = smp.Unet('resnet34', activation=None).cuda() # dir_model = root_path + '/model/unet_resnet34_1_1_best.pth' # 裁切测试集路径 # crop_test_images_path = root_path + '/temp_data_test/0.4crop_test_set1024' crop_test_images_path = root_path + '/temp_data_test/crop_test_set' test_path_list = glob.glob( '/media/totem_disk/totem/weitang/competition/test2/test/*jpg') print("Total {} images for testing.".format(len(test_path_list))) # 裁切程序 if CropStage == True: print("Stage 1: ") #crop images crop(test_path_list, crop_test_images_path, scale=i_scale, image_size=i_size, mode="test") # prob_save_path = root_path + '/temp_data_test/resprob1024' prob_save_path = root_path + '/temp_data_test/dense101_t' crop_predict = root_path + '/temp_data_test/predict512_resnet101_t' # crop_predict = root_path + '/temp_data_test/crop_predict_1024' if TestStage == True: print("Stage 2: ") #predict cropped images test_images_path_list = glob.glob(crop_test_images_path + '/*.jpg') os.makedirs(crop_predict, exist_ok=True) test_loader = get_test_loader(test_images_path_list, image_size=i_size, batch_size=2) test(test_loader, crop_predict, model=tta_model) print("Stage 3: ") #merge predicted images # prob_save_path = root_path + '/temp_data_test/prob' # 缩放倍率:0.25,即将原图*0.25再进行裁切 os.makedirs(prob_save_path, exist_ok=True) merge_hot_pic(test_path_list, crop_predict, i_scale, prob_save_path) mask_save_path = root_path + '/temp_data_test/mask/' if toMask == True: print("Stage 4: ") #convert probs to masks os.makedirs(mask_save_path, exist_ok=True) prob_to_mask(prob_save_path, mask_save_dir=mask_save_path, th=newTH, pad_white=True) if toZip == True: print("Stage 5: ") #zip masks zf = zipfile.ZipFile(f'{root_path}/result/result.zip', 'w') for i in glob.glob(f"{mask_save_path}/*.png"): basename = os.path.split(i)[1] zf.write(i, f'result/{basename}') zf.close()
def get_val_logits(valid_ids, num_split, encoder, decoder): # valid train = "./data/Clouds_Classify/train.csv" # Data overview train = pd.read_csv(open(train)) train.head() train['label'] = train['Image_Label'].apply(lambda x: x.split('_')[1]) train['im_id'] = train['Image_Label'].apply(lambda x: x.split('_')[0]) ENCODER = encoder ENCODER_WEIGHTS = 'imagenet' if decoder == 'unet': #建立模型 model = smp.Unet( encoder_name=ENCODER, encoder_weights=ENCODER_WEIGHTS, classes=4, activation=None, ) else: model = smp.FPN( encoder_name=ENCODER, encoder_weights=ENCODER_WEIGHTS, classes=4, activation=None, ) preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS) num_workers = 4 valid_bs = 1 valid_dataset = CloudDataset(df=train, transforms = get_validation_augmentation(), datatype='valid', img_ids=valid_ids, preprocessing=get_preprocessing(preprocessing_fn)) valid_loader = DataLoader(valid_dataset, batch_size=valid_bs, shuffle=False, num_workers=num_workers) loaders = {"valid": valid_loader} logdir = "./logs/log_{}_{}/log_{}".format(encoder, decoder, num_split) print('predicting for validation data...') ###############使用pytorch TTA预测#################### use_TTA = True checkpoint_path = logdir + '/checkpoints/best.pth' model.load_state_dict(torch.load(checkpoint_path)['model_state_dict']) #使用tta预测 if use_TTA: transforms = tta.Compose([ tta.HorizontalFlip(), tta.VerticalFlip(), tta.Scale(scales=[5/6, 1, 7/6]), ]) tta_model = tta.SegmentationTTAWrapper(model, transforms, merge_mode='mean') else: tta_model = model tta_model = tta_model.cuda() tta_model.eval() with torch.no_grad(): for _, data in enumerate(tqdm.tqdm(loaders['valid'])): img, _, img_name = data img = img.cuda() batch_preds = tta_model(img).cpu().numpy() batch_preds = batch_preds.astype(np.float16) save_dir = './logits/valid/' + encoder + '_' + decoder + '/log_{}'.format(num_split) if not os.path.exists(save_dir): os.makedirs(save_dir) file_name = img_name[0].split('.')[0] + '.plk' file_path = os.path.join(save_dir, file_name) with open(file_path, 'wb') as wf: plk.dump(batch_preds, wf)
def get_test_logits(encoder, decoder): ''' 预测并保存测试集logits ''' sub = "./data/Clouds_Classify/sample_submission.csv" sub = pd.read_csv(open(sub)) sub['label'] = sub['Image_Label'].apply(lambda x: x.split('_')[1]) sub['im_id'] = sub['Image_Label'].apply(lambda x: x.split('_')[0]) #建立模型 preprocessing_fn = smp.encoders.get_preprocessing_fn(encoder, 'imagenet') if decoder == 'unet': model = smp.Unet( encoder_name=encoder, encoder_weights='imagenet', classes=4, activation=None, ) else: model = smp.FPN( encoder_name=encoder, encoder_weights='imagenet', classes=4, activation=None, ) #载入模型参数 logdir = "./logs/log_{}_{}/log_{}".format(encoder, 'fpn', 4) checkpoint_path = logdir + '/checkpoints/best.pth' model.load_state_dict(torch.load(checkpoint_path)['model_state_dict']) #使用tta预测 use_TTA = True if use_TTA: print('using TTA!!!') transforms = tta.Compose([ tta.HorizontalFlip(), tta.VerticalFlip(), tta.Scale(scales=[5/6, 1, 7/6]), ]) tta_model = tta.SegmentationTTAWrapper(model, transforms, merge_mode='mean') else: tta_model = model test_ids = [id for id in os.listdir(test_imgs_folder)] test_dataset = CloudDataset(df=sub, transforms = get_validation_augmentation(), datatype='test', img_ids=test_ids, preprocessing=get_preprocessing(preprocessing_fn)) test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=2) tta_model = tta_model.cuda() tta_model.eval() with torch.no_grad(): for i, data in enumerate(tqdm.tqdm(test_loader)): img, _, img_name = data img = img.cuda() batch_preds = tta_model(img).cpu().numpy() batch_preds = batch_preds.astype(np.float16) save_dir = './logits/test/' + encoder + '_fpn' + '/log_{}'.format(4) if not os.path.exists(save_dir): os.makedirs(save_dir) #保存为pickle file_name = img_name[0].split('.')[0] + '.plk' file_path = os.path.join(save_dir, file_name) with open(file_path, 'wb') as wf: plk.dump(batch_preds, wf)
def main(): fold_path = args.fold_path fold_num = args.fold_num model_name = args.model_name train_csv = args.train_csv sub_csv = args.sub_csv encoder = args.encoder num_workers = args.num_workers batch_size = args.batch_size log_path = args.log_path is_tta = args.is_tta test_batch_size = args.test_batch_size attention_type = args.attention_type print(log_path) train = pd.read_csv(train_csv) train['label'] = train['Image_Label'].apply(lambda x: x.split('_')[-1]) train['im_id'] = train['Image_Label'].apply(lambda x: x.replace('_' + x.split('_')[-1], '')) val_fold = pd.read_csv(f'{fold_path}/valid_file_fold_{fold_num}.csv') valid_ids = np.array(val_fold.file_name) attention_type = None if attention_type == 'None' else attention_type encoder_weights = 'imagenet' if model_name == 'Unet': model = smp.Unet( encoder_name=encoder, encoder_weights=encoder_weights, classes=CLASS, activation='softmax', attention_type=attention_type, ) if model_name == 'Linknet': model = smp.Linknet( encoder_name=encoder, encoder_weights=encoder_weights, classes=CLASS, activation='softmax', ) if model_name == 'FPN': model = smp.FPN( encoder_name=encoder, encoder_weights=encoder_weights, classes=CLASS, activation='softmax', ) if model_name == 'ORG': model = Linknet_resnet18_ASPP( ) preprocessing_fn = smp.encoders.get_preprocessing_fn(encoder, encoder_weights) valid_dataset = CloudDataset(df=train, datatype='valid', img_ids=valid_ids, transforms=get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn)) valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers) loaders = {"infer": valid_loader} runner = SupervisedRunner() checkpoint = torch.load(f"{log_path}/checkpoints/best.pth") model.load_state_dict(checkpoint['model_state_dict']) model.eval() transforms = tta.Compose( [ tta.HorizontalFlip(), tta.VerticalFlip(), ] ) model = tta.SegmentationTTAWrapper(model, transforms) runner.infer( model=model, loaders=loaders, callbacks=[InferCallback()], ) callbacks_num = 0 valid_masks = [] probabilities = np.zeros((valid_dataset.__len__() * CLASS, IMG_SIZE[0], IMG_SIZE[1])) # ======== # val predict # for batch in tqdm(valid_dataset): # クラスごとの予測値 _, mask = batch for m in mask: m = resize_img(m) valid_masks.append(m) for i, output in enumerate(tqdm(runner.callbacks[callbacks_num].predictions["logits"])): for j, probability in enumerate(output): probability = resize_img(probability) # 各クラスごとにprobability(予測値)が取り出されている。jは0~3だと思う。 probabilities[i * CLASS + j, :, :] = probability # ======== # search best size and threshold # class_params = {} for class_id in range(CLASS): attempts = [] for threshold in range(20, 90, 5): threshold /= 100 for min_size in [10000, 15000, 20000]: masks = class_masks(class_id, probabilities, threshold, min_size) dices = class_dices(class_id, masks, valid_masks) attempts.append((threshold, min_size, np.mean(dices))) attempts_df = pd.DataFrame(attempts, columns=['threshold', 'size', 'dice']) attempts_df = attempts_df.sort_values('dice', ascending=False) print(attempts_df.head()) best_threshold = attempts_df['threshold'].values[0] best_size = attempts_df['size'].values[0] class_params[class_id] = (best_threshold, best_size) # ======== # gc # torch.cuda.empty_cache() gc.collect() # ======== # predict # sub = pd.read_csv(sub_csv) sub['label'] = sub['Image_Label'].apply(lambda x: x.split('_')[-1]) sub['im_id'] = sub['Image_Label'].apply(lambda x: x.replace('_' + x.split('_')[-1], '')) test_ids = sub['Image_Label'].apply(lambda x: x.split('_')[0]).drop_duplicates().values test_dataset = CloudDataset(df=sub, datatype='test', img_ids=test_ids, transforms=get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn)) encoded_pixels = get_test_encoded_pixels(test_dataset, runner, class_params, test_batch_size) sub['EncodedPixels'] = encoded_pixels # ======== # val dice # val_Image_Label = [] for i, row in val_fold.iterrows(): val_Image_Label.append(row.file_name + '_Fish') val_Image_Label.append(row.file_name + '_Flower') val_Image_Label.append(row.file_name + '_Gravel') val_Image_Label.append(row.file_name + '_Sugar') val_encoded_pixels = get_test_encoded_pixels(valid_dataset, runner, class_params, test_batch_size) val = pd.DataFrame(val_encoded_pixels, columns=['EncodedPixels']) val['Image_Label'] = val_Image_Label sub.to_csv(f'./sub/sub_{model_name}_fold_{fold_num}_{encoder}.csv', columns=['Image_Label', 'EncodedPixels'], index=False) val.to_csv(f'./val/val_{model_name}_fold_{fold_num}_{encoder}.csv', columns=['Image_Label', 'EncodedPixels'], index=False)
heatmap = utils.detach(runner.callbacks["infer"].heatmap[0]) plt.figure(figsize=(20, 9)) plt.imshow(heatmap, cmap="hot", interpolation="nearest") plt.show() # ### Advanced: test-time augmentations (TTA) # # There is [ttach](https://github.com/qubvel/ttach) is a new awesome library for test-time augmentation for segmentation or classification tasks. # In[ ]: # D4 makes horizontal and vertical flips + rotations for [0, 90, 180, 270] angels. # and then merges the result masks with merge_mode="mean" tta_model = tta.SegmentationTTAWrapper(model, tta.aliases.d4_transform(), merge_mode="mean") tta_runner = SupervisedRunner(model=tta_model, device=utils.get_device(), input_key="image") # In[ ]: infer_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=num_workers) batch = next(iter(infer_loader))
def test_tta(self, mode='train', unet_path=None): """Test model & Calculate performances.""" print(char_color('@,,@ %s with TTA' % (mode))) if not unet_path is None: if os.path.isfile(unet_path): checkpoint = torch.load(unet_path) self.unet.load_state_dict(checkpoint['state_dict']) self.myprint('Successfully Loaded from %s' % (unet_path)) self.unet.train(False) self.unet.eval() if mode == 'train': data_lodear = self.train_loader elif mode == 'test': data_lodear = self.test_loader elif mode == 'valid': data_lodear = self.valid_loader acc = 0. # Accuracy SE = 0. # Sensitivity (Recall) SP = 0. # Specificity PC = 0. # Precision DC = 0. # Dice Coefficient IOU = 0. # IOU length = 0 # model pre for each image detail_result = [] # detail_result = [id, acc, SE, SP, PC, dsc, IOU] with torch.no_grad(): for i, sample in enumerate(data_lodear): (image_paths, images, GT) = sample images_path = list(image_paths) images = images.to(self.device) GT = GT.to(self.device) tta_trans = tta.Compose([ tta.VerticalFlip(), tta.HorizontalFlip(), tta.Rotate90(angles=[0, 180]) ]) tta_model = tta.SegmentationTTAWrapper(self.unet, tta_trans) SR = tta_model(images) # SR = self.unet(images) SR = F.sigmoid(SR) if self.save_image: images_all = torch.cat((images, SR, GT), 0) torchvision.utils.save_image( images_all.data.cpu(), os.path.join(self.result_path, 'images', '%s_%d_image.png' % (mode, i)), nrow=self.batch_size) SR = SR.data.cpu().numpy() GT = GT.data.cpu().numpy() for ii in range(SR.shape[0]): SR_tmp = SR[ii, :].reshape(-1) GT_tmp = GT[ii, :].reshape(-1) tmp_index = images_path[ii].split(sep)[-1] tmp_index = int(tmp_index.split('.')[0][:]) SR_tmp = torch.from_numpy(SR_tmp).to(self.device) GT_tmp = torch.from_numpy(GT_tmp).to(self.device) result_tmp = np.array([ tmp_index, get_accuracy(SR_tmp, GT_tmp), get_sensitivity(SR_tmp, GT_tmp), get_specificity(SR_tmp, GT_tmp), get_precision(SR_tmp, GT_tmp), get_DC(SR_tmp, GT_tmp), get_IOU(SR_tmp, GT_tmp) ]) acc += result_tmp[1] SE += result_tmp[2] SP += result_tmp[3] PC += result_tmp[4] DC += result_tmp[5] IOU += result_tmp[6] detail_result.append(result_tmp) length += 1 accuracy = acc / length sensitivity = SE / length specificity = SP / length precision = PC / length disc = DC / length iou = IOU / length detail_result = np.array(detail_result) if (self.save_detail_result ): # detail_result = [id, acc, SE, SP, PC, dsc, IOU] excel_save_path = os.path.join(self.result_path, mode + '_pre_detial_result.xlsx') writer = pd.ExcelWriter(excel_save_path) detail_result = pd.DataFrame(detail_result) detail_result.to_excel(writer, mode, float_format='%.5f') writer.save() writer.close() return accuracy, sensitivity, specificity, precision, disc, iou
def main(args): # set GPUS os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu # define where stithced and not stitched test data is located src_sliced_dir = os.path.join(args.test_dir, "sliced") src_stitched_dir = os.path.join(args.test_dir, "stitched") # prepare output (prediction) dirs dst_sliced_dir = os.path.join(args.dst_dir, "sliced") dst_stitch_dir = os.path.join(args.dst_dir, "stitched") os.makedirs(dst_sliced_dir, exist_ok=True) os.makedirs(dst_stitch_dir, exist_ok=True) # -------------------------------------------------- # define model # -------------------------------------------------- device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("Available devices:", device) # loading trained models models = [model_from_config(config_path) for config_path in args.configs] # create ensemble model = EnsembleModel(models) # add test time augmentations (flipping and rotating input image) model = ttach.SegmentationTTAWrapper(model, ttach.aliases.d4_transform(), merge_mode='mean') # create Multi GPU model if number of GPUs is more than one n_gpus = len(args.gpu.split(",")) if n_gpus > 1: gpus = list(range(n_gpus)) model = torch.nn.DataParallel(model, gpus) print("Done loading...") model.to(device) # -------------------------------------------------- # start evaluation # -------------------------------------------------- runner = GPUNormRunner(model, model_device=device) model.eval() # predict big tif files predictor = TorchTifPredictor( runner=runner, sample_size=1024, cut_edge=256, batch_size=args.batch_size, count=1, NBITS=1, compress=None, driver="GTiff", blockxsize=256, blockysize=256, tiled=True, BIGTIFF='IF_NEEDED', ) print("Inference for big tif files...") filesnames = os.listdir(src_stitched_dir) for filename in filesnames: _src_path = os.path.join(src_stitched_dir, filename) _dst_path = os.path.join(dst_stitch_dir, filename) predictor(_src_path, _dst_path) print("Slicing big tif files to original test size...") df = pd.read_csv(args.test_csv) df = df[df.cluster_id != -1] cluster_ids = df.cluster_id.unique() for id in cluster_ids: _df = df[df.cluster_id == id] _args = [(dst_stitch_dir, dst_sliced_dir, row) for i, row in _df.iterrows()] with Pool(12) as p: with tqdm(p.imap(slice_to_tiles, _args), total=len(_args)) as pp: for _ in pp: pass # predict not stitched data print("Predicting small tif files...") test_dataset = TestSegmentationDataset(src_sliced_dir, transform_name='test_transform_1') test_dataloader = torch.utils.data.DataLoader( test_dataset, batch_size=args.batch_size, pin_memory=True, num_workers=8, ) for batch in tqdm(test_dataloader): ids = batch['id'] predictions = runner.predict_on_batch(batch)['mask'] for image_id, prediction in zip(ids, predictions): prediction = prediction.round().int().cpu().numpy().astype("uint8") profile = test_dataset.read_image_profile(image_id) dst_path = os.path.join(dst_sliced_dir, image_id) write_tile(dst_path, prediction, profile, compress="lzw", driver="GTiff")
] models = [] for pth in model_pth: model = torch.load(pth, map_location=lambda storage, loc: storage) model.to(device) model.eval() models.append(model) print(len(models)) # print(model_pth[-1]) # model = torch.load(model_pth[-1], map_location=lambda storage, loc: storage) # model.to(device) # model.eval() tta_model = tta.SegmentationTTAWrapper(model, tta.aliases.hflip_transform(), merge_mode='mean') # model.to(device) # class_params = {0: (0.65, 10000), 1: (0.65, 10000), 2: (0.65, 10000), 3: (0.60, 10000)}#Efficient-B0 # class_params = {0: (0.6, 30000), 1: (0.65, 20000), 2: (0.6, 15000), 3: (0.65, 10000)} # Efficient-B2 # class_params = {0: (0.65, 20000), 1: (0.6, 15000), 2: (0.65, 15000), 3: (0.6, 10000)} # se_resnext50_32x4d # class_params = {0: (0.7, 10000), 1: (0.6, 15000), 2: (0.6, 20000), 3: (0.65, 15000)} #inceptionresnetv2 # class_params = {0: (0.65, 20000), 1: (0.65, 20000), 2: (0.65, 20000), 3: (0.6, 15000)} # resnet50 # class_params = {0: (0.55, 25000), 1: (0.55, 20000), 2: (0.55, 20000), 3: (0.55, 20000)} # se_resnext50_32x4d, resnet50, efficientnet-b2 # class_params = {0: (0.65, 20000), 1: (0.65, 25000), 2: (0.65, 20000), 3: (0.65, 15000)} # se_resnext50_32x4d, resnet50, efficientnet-b2, se-resnet50 # class_params = {0: (0.6, 25000), 1: (0.65, 20000), 2: (0.6, 25000), 3: (0.65, 20000)} # efficientB2 0.6467 # class_params = {0: (0.65, 30000), 1: (0.65, 15000), 2: (0.6, 15000), 3: (0.65, 10000)} # se_resnext50_32x4d # class_params = {0: (0.6, 25000), 1: (0.65, 20000), 2: (0.65, 20000), 3: (0.65, 15000)} # resnet50 # class_params = {0: (0.6, 25000), 1: (0.65, 20000), 2: (0.65, 20000), 3: (0.65, 15000)} # se_resnet50 # class_params = {0: (0.65, 20000), 1: (0.65, 20000), 2: (0.6, 20000), 3: (0.65, 15000)} # Unet_resnet34 # 0.6522
def main(): input_dir = './input/test_images_png/' pre_detection_dir = './input/test_images_predetection/' output_dir = './output/' x_valid_dir = input_dir y_valid_dir = input_dir predict_interval_width = 256 # 画像を幅方向に何ピクセルごとに予測するか crop_size = 512 crop_size_half = int(crop_size/2) ENCODER = 'densenet121' ENCODER_WEIGHTS = 'imagenet' CLASSES = ['coastline'] ACTIVATION = 'sigmoid' # could be None for logits or 'softmax2d' for multicalss segmentation DEVICE = 'cuda' preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS) # load best saved checkpoint best_model = torch.load('./runs/Nov01_14-42-50_ashida_ENCODER_densenet121_LR_0.0001/best_model.pth') # tta_model = tta.SegmentationTTAWrapper(best_model, tta.aliases.d4_transform(), merge_mode='mean') transforms = tta.Compose( [ tta.HorizontalFlip(), # tta.Rotate90(angles=[0, 180]) ] ) tta_model = tta.SegmentationTTAWrapper(best_model, transforms) # create test dataset test_dataset = Dataset( x_valid_dir, # x_test_dir pre_detection_dir, y_valid_dir, # y_test_dir augmentation=get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn), classes=CLASSES, ) # test_dataloader = DataLoader(test_dataset) # # # evaluate model on test set # test_epoch = smp.utils.train.ValidEpoch( # model=best_model, # loss=loss, # metrics=metrics, # device=DEVICE, # ) # # logs = test_epoch.run(test_dataloader) # test dataset without transformations for image visualization test_dataset_vis2 = Dataset( x_valid_dir, pre_detection_dir, y_valid_dir, # x_test_dir, y_test_dir, augmentation=get_validation_augmentation(), classes=CLASSES, ) for n, i in enumerate(range(len(test_dataset))): image_vis = test_dataset_vis2[i][0].astype('uint8') image, pre_detection_image, gt_mask = test_dataset[i] file_name = test_dataset.images_fps[i] base_name = os.path.basename(file_name) # パッチに分けて予測する h = pre_detection_image.shape[0] # 回転したあとは変わるので注意 w = pre_detection_image.shape[1] # 回転したあとは変わるので注意 # 画像が縦長な場合は90度回転させる if h > w: # image = torch.rot90(image, 1, [1, 2]) image = image.transpose(1, 2, 0) image = cv.rotate(image, cv.ROTATE_90_CLOCKWISE) image = image.transpose(2, 0, 1) pre_detection_image = cv.rotate(pre_detection_image, cv.ROTATE_90_CLOCKWISE) rotate_90 = True else: rotate_90 = False # 画像を何分割で予測するか計算する image_len = pre_detection_image.shape[1] split_num = image_len // predict_interval_width # 岡田さんロジックで事前検出した海岸線位置のインデックスを取得する predict_index = np.argmax(pre_detection_image, axis=0) predict_image = np.zeros((pre_detection_image.shape[0], pre_detection_image.shape[1])) # 画像を分割して海岸線を予測する for j in range(1, split_num): width_idx = j*predict_interval_width tmp = pre_detection_image[predict_index[width_idx], width_idx] y = predict_index[width_idx] x = width_idx if y == 0: continue ys = y - crop_size_half xs = x - crop_size_half if ys < 0: ys = 0 elif ys + crop_size > pre_detection_image.shape[0]: ys = pre_detection_image.shape[0] - crop_size if xs < 0: xs = 0 elif xs + crop_size > pre_detection_image.shape[1]: xs = pre_detection_image.shape[0] - crop_size crop_image = image[:, ys:ys+crop_size, xs:xs+crop_size] x_tensor = torch.from_numpy(crop_image).to(DEVICE).unsqueeze(0) pr_mask = tta_model.forward(x_tensor) pr_mask = pr_mask.squeeze().to('cpu').detach().numpy().copy() # pr_mask = (pr_mask * 255).astype(np.uint8) # ret, pr_mask = cv.threshold(pr_mask, 1, 255, cv.THRESH_BINARY) tmp_predict_image = np.zeros((pre_detection_image.shape[0], pre_detection_image.shape[1])) tmp_predict_image[ys:ys+crop_size, xs:xs+crop_size] = pr_mask predict_image = predict_image + tmp_predict_image gt_mask = gt_mask.squeeze() if rotate_90: predict_image = cv.rotate(predict_image, cv.ROTATE_90_COUNTERCLOCKWISE) pre_detection_image = cv.rotate(pre_detection_image, cv.ROTATE_90_COUNTERCLOCKWISE) # x_tensor = torch.from_numpy(image).to(DEVICE).unsqueeze(0) # pr_mask = best_model.predict(x_tensor) # pr_mask = tta_model.forward(x_tensor) # pr_mask = (pr_mask.squeeze().cpu().numpy().round()) # pr_mask = pr_mask.squeeze().to('cpu').detach().numpy().copy() predict_image = np.clip(predict_image, None, 1) predict_image = (predict_image*255).astype(np.uint8) # ret, predict_image = cv.threshold(predict_image, 1, 255, cv.THRESH_BINARY) # visualize( # image=image_vis, # ground_truth_mask=gt_mask, # predicted_mask=predict_image # ) org_image = cv.imread(file_name) h = org_image.shape[0] w = org_image.shape[1] predict_image = cv.resize(predict_image, (w, h)) cv.imwrite(output_dir + base_name, predict_image) print(str(n) + '/' + str(len(test_dataset)))
# 构建两个模型 with torch.no_grad(): # tta设置 tta_trans = tta.Compose([ tta.VerticalFlip(), tta.HorizontalFlip(), tta.Rotate90(angles=[0,180]), ]) # 构建模型 # cascade1 model_cascade1 = smp.DeepLabV3Plus(encoder_name="efficientnet-b6", encoder_weights=None, in_channels=1, classes=1) model_cascade1.to(device) model_cascade1.load_state_dict(torch.load(weight_c1)) if c1_tta: model_cascade1 = tta.SegmentationTTAWrapper(model_cascade1, tta_trans,merge_mode='mean') model_cascade1.eval() # cascade2 model_cascade2 = smp.DeepLabV3Plus(encoder_name="efficientnet-b6", encoder_weights=None, in_channels=1, classes=1) # model_cascade2 = smp.Unet(encoder_name="efficientnet-b6", encoder_weights=None, in_channels=1, classes=1, encoder_depth=5, decoder_attention_type='scse') # model_cascade2 = smp.PAN(encoder_name="efficientnet-b6",encoder_weights='imagenet', in_channels=1, classes=1) model_cascade2.to(device) model_cascade2.load_state_dict(torch.load(weight_c2)) if c2_tta: model_cascade2 = tta.SegmentationTTAWrapper(model_cascade2, tta_trans,merge_mode='mean') model_cascade2.eval() # 指标 IOU_list = [] DSC_list = []
import pandas as pd import torch from torch.utils.data import DataLoader from osgeo import gdal from src.models.siamese_unet import SCSeResneXt, ResneXt from src.dataset import Satellite from src.utils import * with open(sys.argv[1], 'r') as f: config = json.load(f) df = pd.read_csv(config['sample_submission_path']) best_model = torch.load(f"models/saved/{config['model_name']}.pth") tta_model = tta.SegmentationTTAWrapper(best_model, tta.aliases.d4_transform(), merge_mode='mean') original_res = [] res = [] for file in os.listdir(config['mask_path']): ds = gdal.Open(f"{config['images_path']}{file}") IMG1 = np.array([ds.GetRasterBand(i).ReadAsArray() for i in range(1, 4)]) IMG2 = np.array([ds.GetRasterBand(i).ReadAsArray() for i in range(5, 8)]) IMG1, IMG2 = normalize(IMG1, IMG2, config['img_channels']) new_shape = generate_new_shape(IMG1, config['img_size'], config['img_channels']) res_mask = np.zeros((new_shape[0], new_shape[1])) if config['img_channels'] == 3:
Return: transform: albumentations.Compose """ _transform = [ albu.Normalize(), albu.Lambda(image=to_tensor, mask=to_tensor), ] return albu.Compose(_transform) model_seg = torch.load(path_to_segmentation_model) #TTA tta_model = tta.SegmentationTTAWrapper(model_seg, tta.aliases.hflip_transform(), merge_mode='sum') segmentation_preprocessing = get_preprocessing() results = [] score = [] for id in sample_df['id']: print(id) image = cv2.imread(test_data_path + str(id) + '.png') image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) #Classification sample = classification_preprocessing(image=image) image_to_classify = sample['image']
1, n_iter // 20, after_scheduler=schedulerD) # plot_scheduler(schedulerG, n_iter, "Gen LR Scheduler") # plot_scheduler(schedulerD, n_iter * NUM_DISCRIMINATOR_TRAIN, "Dis LR Scheduler") device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") gen.to(device) dis.to(device) criterionG.to(device) criterionD.to(device) tta_transform = tta.Compose([tta.HorizontalFlip(), tta.VerticalFlip()]) tta_gen = tta.SegmentationTTAWrapper(gen, tta_transform) tta_dis = tta.ClassificationTTAWrapper(dis, tta_transform) history = [] img_test = plt.imread(r'data\testing_lr_images\09.png') img_valid = plt.imread(r'data\valid_hr_images\t20.png') for epoch in range(1, NUM_EPOCHS + 1): train_bar = tqdm(train_loader) running_results = { 'batch_sizes': 0, 'lossG': 0, 'lossD': 0, 'gradeHR': 0, 'gradeSR': 0, 'psnr': 0 }