def SUE_TTA(model, batch: torch.tensor, last_layer: bool) -> Tuple[np.ndarray, np.ndarray]: r"""Interface of Binary Segmentation Uncertainty Estimation with Test-Time Augmentations (TTA) method for 1 2D slice. Inputs supposed to be in range [0, data_range]. Args: model: Trained model. batch: Tensor with shape (1, C, H, W). last_layer: Flag whether there is Sigmoid as a last NN layer Returns: Aleatoric and epistemic uncertainty maps with shapes equal to batch shape """ model.eval() transforms = tta.Compose( [ tta.VerticalFlip(), tta.HorizontalFlip(), tta.Rotate90(angles=[0, 180]), tta.Scale(scales=[1, 2, 4]), tta.Multiply(factors=[0.9, 1, 1.1]), ] ) predicted = [] for transformer in transforms: augmented_image = transformer.augment_image(batch) model_output = model(augmented_image) deaug_mask = transformer.deaugment_mask(model_output) prediction = torch.sigmoid( deaug_mask).cpu().detach().numpy() if last_layer else deaug_mask.cpu().detach().numpy() predicted.append(prediction) p_hat = np.array(predicted) aleatoric = calc_aleatoric(p_hat) epistemic = calc_epistemic(p_hat) return aleatoric, epistemic
def __init__(self, config): self.config = config self.classes = config.CLASS_NAME self.input_size = config.INPUT_SIZE self.binary_option = config.BINARY self.failClasses = config.FAIL_CLASSNAME self.passClasses = config.PASS_CLASSNAME self.pass_class_index = [ self.classes.index(class_) for class_ in self.passClasses ] self.fail_class_index = [ self.classes.index(class_) for class_ in self.failClasses ] self.pytorch_model = None # self.train_generator = None # self.val_generator = None # self.test_generator = None self.class_weights = None self.evaluate_generator = None self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") self.tta_rotate_opt_list = [0, 90, 180, 270] self.tta_option = tta.Compose([ tta.Rotate90( self.tta_rotate_opt_list), # For future developing with CAM # tta.HorizontalFlip(), # tta.VerticalFlip() ]) # Toggle TTA option self.tta_opt = True self.global_batch_size = self.config.BATCH_SIZE * self.config.GPU_COUNT if self.config.GPU_COUNT != 0 else self.config.BATCH_SIZE
def evaluate(val_loader, model, loss_fn, device, use_tta=False): model.eval() if use_tta: transformations = tta.Compose([ tta.Rotate90(angles=[0, 90, 180, 270]), tta.HorizontalFlip(), tta.VerticalFlip() ]) tta_model = tta.ClassificationTTAWrapper(model, transformations) correct = 0 total = 0 total_loss = 0 for i, batch in enumerate(val_loader): input_data, labels = batch input_data, labels = input_data.to(device), labels.to(device) with torch.no_grad(): if use_tta: predictions = tta_model(input_data) else: predictions = model(input_data) total_loss += loss_fn(predictions, labels).item() correct += (predictions.argmax(axis=1) == labels).sum().item() total += len(labels) torch.cuda.empty_cache() model.train() return total_loss / total, correct / total
def get_predictions(model_chosen, tta = False): model_chosen.cuda.eval() actual_values, predicted_values = [], [] if tta == True: transformation = ttach.Compose( [ ttach.HorizontalFlip(), ttach.VerticalFlip(), ttach.Rotate90(angles=[0, 90, 180, 270]) ] ) test_time_augmentation_wrapper = ttach.ClassificationTTAWrapper(model_chosen, transformation) with torch.no_grad(): for batch in loader_of_test: test_image, test_label = batch predicted_value = test_time_augmentation_wrapper(test_image.cuda()) predicted_value = torch.argmax(predicted_value, dim=1).detach().cpu().numpy() actual_values.append(test_label.cpu().numpy()) predicted_values.append(predicted_value) else: with torch.no_grad(): for batch in loader_of_test: test_image, test_label = batch predicted_value = model_chosen(test_image.cuda()) predicted_value = torch.argmax(predicted_value, dim=1).detach().cpu().numpy() actual_values.append(test_label.cpu().numpy()) predicted_values.append(predicted_value) return predicted_values
def predict(model_path, test_loader, saveFileName, iftta): ## predict model = initialize_model(num_classes=176) # create model and load weights from checkpoint model = model.to(device) model.load_state_dict(torch.load(model_path)) if iftta: print("Using TTA") transforms = tta.Compose( [ tta.HorizontalFlip(), tta.VerticalFlip(), tta.Rotate90(angles=[0, 180]), # tta.Scale(scales=[1, 0.3]), ] ) model = tta.ClassificationTTAWrapper(model, transforms) # Make sure the model is in eval mode. # Some modules like Dropout or BatchNorm affect if the model is in training mode. model.eval() # Initialize a list to store the predictions. predictions = [] # Iterate the testing set by batches. for batch in tqdm(test_loader): imgs = batch with torch.no_grad(): logits = model(imgs.to(device)) # Take the class with greatest logit as prediction and record it. predictions.extend(logits.argmax(dim=-1).cpu().numpy().tolist()) preds = [] for i in predictions: preds.append(num_to_class[i]) test_data = pd.read_csv('leaves_data/test.csv') test_data['label'] = pd.Series(preds) submission = pd.concat([test_data['image'], test_data['label']], axis=1) submission.to_csv(saveFileName, index=False) print("Done!!!!!!!!!!!!!!!!!!!!!!!!!!!")
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_compose_1(): transform = tta.Compose([ tta.HorizontalFlip(), tta.VerticalFlip(), tta.Rotate90(angles=[0, 90, 180, 270]), tta.Scale(scales=[1, 2, 4], interpolation="nearest"), ]) assert len( transform) == 2 * 2 * 4 * 3 # all combinations for aug parameters dummy_label = torch.ones(2).reshape(2, 1).float() dummy_image = torch.arange(2 * 3 * 4 * 5).reshape(2, 3, 4, 5).float() dummy_model = lambda x: {"label": dummy_label, "mask": x} for augmenter in transform: augmented_image = augmenter.augment_image(dummy_image) model_output = dummy_model(augmented_image) deaugmented_mask = augmenter.deaugment_mask(model_output["mask"]) deaugmented_label = augmenter.deaugment_label(model_output["label"]) assert torch.allclose(deaugmented_mask, dummy_image) assert torch.allclose(deaugmented_label, dummy_label)
[A.Normalize(mean=mean, std=std, max_pixel_value=max_value), ToTensorV2()]) _, data_loader, _ = get_train_val_loaders( train_ds, val_ds, train_transforms=transforms, val_transforms=transforms, batch_size=batch_size, num_workers=num_workers, val_batch_size=batch_size, pin_memory=True, ) prepare_batch = inference_prepare_batch_f32 # Image denormalization function to plot predictions with images img_denormalize = partial(denormalize, mean=mean, std=std) #################### Model #################### model = FPN(encoder_name='se_resnext50_32x4d', classes=2, encoder_weights=None) run_uuid = "5230c20f609646cb9870a211036ea5cb" weights_filename = "best_model_67_val_miou_bg=0.7574240313552584.pth" has_targets = True tta_transforms = tta.Compose([ tta.Rotate90(angles=[90, -90, 180]), ])
_, test = get_fold_filelist(csv_file, K=fold_K, fold=fold_index) test_img_list = [img_path+sep+i[0] for i in test] if mask_path is not None: test_mask_list = [mask_path+sep+i[0] for i in test] else: test_img_list = get_filelist_frompath(img_path,'PNG') if mask_path is not None: test_mask_list = [mask_path + sep + i.split(sep)[-1] for i in test_img_list] # 构建两个模型 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))
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
label = self.label_df.iloc[index,1:].values.astype('float') if self.transforms: image = self.transforms(image=image)['image'] / 255.0 return image, label base_transforms = { 'test' : albumentations.Compose([ albumentations.pytorch.ToTensorV2(), ]), } tta_transforms = tta.Compose( [ tta.Rotate90(angles=[0, 90, 180, 270]), # tta.Scale(scales=[0.9,1.0,1.1]), ## batch1로 했었음 ] ) def main(): device = "cuda:0" if torch.cuda.is_available() else "cpu" parser = argparse.ArgumentParser() parser.add_argument('--image_path', type=str, default="./data/test_dirty_mnist_2nd/") parser.add_argument('--label_path', type=str, default="./data/sample_submission.csv") parser.add_argument('--weight_path', type=str, default='./save/kfold_202119/') parser.add_argument('--out_path', type=str, default='./save/kfold_202119/') parser.add_argument('--model', type=str, default='efficientnet-b8')
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 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)