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)
