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
}
