def get_model(
height,
width,
channels,
model_path,
version=0
): #Gets a new, untrained unet, or a pretrained model. Returns Keras model object
inputs = Input((height, width, channels))
if (version == 1): #Untrained u-net implementations (see models.py)
base = models.get_unet(inputs, 1)
elif (version == 2):
base = models.get_unet2(inputs, 1)
elif (version == 3):
base = models.get_unet3(inputs, 1)
else:
print("Getting pretrained model from %s" % model_path)
model = load_model(model_path)
return model
reshape = Reshape((-1, 1))(base)
act = Activation('relu')(reshape)
model = Model(inputs=inputs, outputs=act)
model.compile(optimizer=Adadelta(), loss="binary_crossentropy")
return model
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-m', '--model')
args = parser.parse_args()
model = models.get_unet(input_shape=(1280, 1280, 3),
pool_cnt=7,
filter_cnt=8)
if args.model:
model.load_weights(args.model, by_name=True)
train_x = HDF5Matrix(DATA_FILE, "train/x", end=4080)
train_y = HDF5Matrix(DATA_FILE, "train/y", end=4080)
train_it = HDF5MatrixIterator(train_x,
train_y,
batch_size=BATCH_SIZE,
preprocess=preprocess_train,
shuffle=True)
train_cnt = len(train_x)
val_x = HDF5Matrix(DATA_FILE, "train/x", start=4080)
val_y = HDF5Matrix(DATA_FILE, "train/y", start=4080)
val_it = HDF5MatrixIterator(val_x,
val_y,
batch_size=BATCH_SIZE,
preprocess=preprocess_val)
val_cnt = len(val_x)
chkpt = ModelCheckpoint(filepath='model.hdf5',
monitor='val_dice_coef',
verbose=1,
save_best_only=True,
save_weights_only=True,
mode='max')
early_stop = EarlyStopping(monitor='val_dice_coef',
patience=5,
verbose=1,
mode='max')
model.fit_generator(generator=train_it,
steps_per_epoch=math.ceil(train_cnt / BATCH_SIZE),
epochs=200,
callbacks=[chkpt, early_stop],
validation_data=val_it,
validation_steps=math.ceil(val_cnt / BATCH_SIZE),
max_q_size=20,
workers=4)
def main():
no_of_parts = 5 #for saving "train\" folder images to numpy
batch_size = 1 #batch normalization is used. 1 is best
predict_while_training = True
if len(os.listdir('batches')) == 0:
print("Found no batches. Saving Numpy from folder 'train\\'")
save_numpy(no_of_parts)
if len(os.listdir('checkpoint')) == 0:
print("Found no save points. Creating new model")
model = models.get_unet(img_rows=256, img_cols=256, dimensions=1)
else:
newest = check_newest()
print("Found Save Point. Using {}".format(newest))
model = load_model('checkpoint\\{}'.format(newest))
train(model, batch_size, predict_while_training)
def main(image_dir, label_dir, checkpoint_path, output_dir, landmarks_path,
df_path, batch_size, num_workers, multi_gpu):
import torch
import torchio as tio
import models
import datasets
import engine
import utils
fps = get_paths(image_dir)
lfps = get_paths(label_dir)
assert len(fps) == len(lfps)
# key must be 'image' as in get_test_transform
subjects = [
tio.Subject(image=tio.ScalarImage(fp), label=tio.LabelMap(lfp))
for (fp, lfp) in zip(fps, lfps)
]
transform = datasets.get_test_transform(landmarks_path)
dataset = tio.SubjectsDataset(subjects, transform)
checkpoint = torch.load(checkpoint_path)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = models.get_unet().to(device)
if multi_gpu:
model = torch.nn.DataParallel(model)
model.module.load_state_dict(checkpoint['model'])
else:
model.load_state_dict(checkpoint['model'])
output_dir = Path(output_dir)
model.eval()
torch.set_grad_enabled(False)
loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers)
output_dir.mkdir(parents=True)
evaluator = engine.Evaluator()
df = evaluator.infer(model, loader, output_dir)
df.to_csv(df_path)
med, iqr = 100 * utils.get_median_iqr(df.Dice)
print(f'{med:.1f} ({iqr:.1f})')
return 0
def get_prediction(data, model_name):
w, h = data.shape[:2]
split = []
hstart = 0
if ('model_trees' in model_name) or ('model_crops'
in model_name) or ('model_road'
in model_name):
model = get_unet()
else:
model = get_model_incept()
model.load_weights(model_name)
for i in range(1, 9):
hend = (h / 8) * i
wstart = 0
temp = []
for j in range(1, 9):
wend = (w / 8) * j
test = data[int(hstart):int(hend), int(wstart):int(wend)]
img = cv2.resize(test, (224, 224), cv2.INTER_AREA)
img = img / np.max(img)
predicted = model.predict(np.reshape(img, (-1, 224, 224, 3)))
predicted = predicted[0]
temp.append(np.round(predicted))
wstart = wend
split.append(
np.concatenate([
temp[0], temp[1], temp[2], temp[3], temp[4], temp[5], temp[6],
temp[7]
],
axis=1))
hstart = hend
pred = np.concatenate([
split[0], split[1], split[2], split[3], split[4], split[5], split[6],
split[7]
],
axis=0)
return pred
def main(input_path, checkpoint_path, output_dir, landmarks_path, batch_size, num_workers, resample):
import torch
from tqdm import tqdm
import torchio as tio
import models
import datasets
fps = get_paths(input_path)
subjects = [tio.Subject(image=tio.ScalarImage(fp)) for fp in fps] # key must be 'image' as in get_test_transform
transform = tio.Compose((
tio.ToCanonical(),
datasets.get_test_transform(landmarks_path),
))
if resample:
transform = tio.Compose((
tio.Resample(),
transform,
# tio.CropOrPad((264, 268, 144)), # ################################# for BITE?
))
dataset = tio.SubjectsDataset(subjects, transform)
checkpoint = torch.load(checkpoint_path)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = models.get_unet().to(device)
model.load_state_dict(checkpoint['model'])
output_dir = Path(output_dir)
model.eval()
torch.set_grad_enabled(False)
loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)
output_dir.mkdir(exist_ok=True, parents=True)
for batch in tqdm(loader):
inputs = batch['image'][tio.DATA].float().to(device)
seg = model(inputs).softmax(dim=1)[:, 1:].cpu() > 0.5
for tensor, affine, path in zip(seg, batch['image'][tio.AFFINE], batch['image'][tio.PATH]):
image = tio.LabelMap(tensor=tensor, affine=affine.numpy())
path = Path(path)
out_path = output_dir / path.name.replace('.nii', '_seg_cnn.nii')
image.save(out_path)
return 0
# xtrain = xtrain[:200]
# ytrain = ytrain[:200]
# xval = xval[:50]
# yval = yval[:50]
H, W = 224, 224
lr = 3e-3
loss = dice_loss
BATCH_SIZE = 5
EPOCHS = 5
print('getting model...')
model = get_unet(H, W, lr, loss)
# In[5]:
datagen = ImageDataGenerator(zoom_range=0.1,
width_shift_range=0.1,
height_shift_range=0.1,
vertical_flip=True)
datagen.fit(xtrain)
# In[6]:
# define callbacks
lr_plat = ReduceLROnPlateau(monitor='val_dice_coef',
factor=0.2,
patience=5,
return 1 - muti_dice_coef(y_true, y_pred)
height, width = 512, 512
nb_epoch = 200
batch_size = 1
model_name = "kits2019"
model_filename = "saved_models/{}.h5".format(model_name)
print(model_filename)
if not os.path.exists('saved_models'): os.mkdir('saved_models')
model = models.get_unet(height,
width,
loss=[muti_dice_coef_loss],
optimizer=Adam(lr=1e-4),
metrics=[
muti_jacc_coef, 'accuracy', sensitivity,
specificity, muti_dice_coef
],
channels=1,
num_class=1)
if os.path.exists("saved_models/{}_1.h5".format(model_name)):
print('loading model')
model = load_model("saved_models/{}_1.h5".format(model_name),
custom_objects={
'muti_dice_coef_loss': muti_dice_coef_loss,
'muti_jacc_coef': muti_jacc_coef,
'muti_dice_coef': muti_dice_coef,
'sensitivity': sensitivity,
'specificity': specificity
})
# print("Mean R, G, B, M", np.mean(train_img_1[:,0,:,:]), np.mean(train_img_1[:,1,:,:]),
# np.mean(train_img_1[:,2,:,:]), np.mean(train_mask_1[:,0,:,:]) )
# train_img[:,0,:,:] -= 187
# train_img[:,1,:,:] -= 171
# train_img[:,2,:,:] -= 182
# train_img = train_img/128;
# train_img = train_img - 128;
# train_img = train_img/128;
# print train_img
print(train_img.shape, train_mask.shape, len(train_img_list),
train_img.shape[0])
model = mo.get_unet(num_color_component, filter_size)
# if os.path.isfile(WEIGHT_OLD_CACHE_PATH):
# model.load_weights(WEIGHT_OLD_CACHE_PATH)
# else:
# print("No model loaded!")
num_images = train_img.shape[0]
num_train_images = int(num_images * train_factor)
num_valid_images = num_images - num_train_images
print("Num images total, train, valid ", num_images, num_train_images,
num_valid_images)
train_x = train_img[0:num_train_images]
valid_x = train_img[num_train_images:num_images]
train_y = train_mask[0:num_train_images]
valid_y = train_mask[num_train_images:num_images]
downscale_factor = 2
train_factor = 0.9
batch_size = 128
nb_epoch = 10
TEST_OUTPUT_BASE_PATH = DATA_CACHE_BASE_PATH + "output_64_1_17_ovr_lpd_7/"
#[test_list, test_file_list] = dl.get_test_images(TEST_PATH)
scaled_dimension = dimension / downscale_factor
#[test_scaled, test_file_list, test_file_dimension] = dl.get_test_images_dwn_scled(TEST_PATH, scaled_dimension)
[test_scaled, test_file_list,
test_file_dimension] = dl.get_test_images_ovr_lap(TEST_PATH, scaled_dimension)
#[test_scaled, test_file_list, test_file_dimension] = dl.get_test_images_ovr_lap_off(TEST_PATH, scaled_dimension)
print("Read test images", len(test_scaled), test_scaled[0].shape)
model = mo.get_unet(num_color_component, scaled_dimension)
model.load_weights(WEIGHT_CACHE_PATH)
print("Model Loaded")
test_y = model.predict(test_scaled, batch_size=batch_size, verbose=1)
print("Prediction Ends ", test_y.shape)
#dl.dump_test_images_dwn_scled(TEST_OUTPUT_BASE_PATH, test_file_list, test_file_dimension, test_y, scaled_dimension)#, TEST_PATH)
dl.dump_test_images_ovr_lpd(TEST_OUTPUT_BASE_PATH, test_file_list,
test_file_dimension, test_y,
scaled_dimension) #, TEST_PATH)
#dl.dump_test_images_ovr_lpd_off(TEST_OUTPUT_BASE_PATH, test_file_list, test_file_dimension, test_y, scaled_dimension)#, TEST_PATH)
print("Test image dumped")
def main():
print("Starting DFC2021 baseline training script at %s" % (str(datetime.datetime.now())))
#-------------------
# Setup
#-------------------
assert os.path.exists(args.input_fn)
if os.path.isfile(args.output_dir):
print("A file was passed as `--output_dir`, please pass a directory!")
return
if os.path.exists(args.output_dir) and len(os.listdir(args.output_dir)):
if args.overwrite:
print("WARNING! The output directory, %s, already exists, we might overwrite data in it!" % (args.output_dir))
else:
print("The output directory, %s, already exists and isn't empty. We don't want to overwrite and existing results, exiting..." % (args.output_dir))
return
else:
print("The output directory doesn't exist or is empty.")
os.makedirs(args.output_dir, exist_ok=True)
if torch.cuda.is_available():
device = torch.device("cuda:%d" % args.gpu)
else:
print("WARNING! Torch is reporting that CUDA isn't available, exiting...")
return
np.random.seed(args.seed)
torch.manual_seed(args.seed)
#-------------------
# Load input data
#-------------------
input_dataframe = pd.read_csv(args.input_fn)
image_fns = input_dataframe["image_fn"].values
label_fns = input_dataframe["label_fn"].values
groups = input_dataframe["group"].values
dataset = StreamingGeospatialDataset(
imagery_fns=image_fns, label_fns=label_fns, groups=groups, chip_size=CHIP_SIZE, num_chips_per_tile=NUM_CHIPS_PER_TILE, windowed_sampling=False, verbose=False,
image_transform=image_transforms, label_transform=label_transforms, nodata_check=nodata_check
)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
num_workers=NUM_WORKERS,
pin_memory=True,
)
num_training_batches_per_epoch = int(len(image_fns) * NUM_CHIPS_PER_TILE / args.batch_size)
print("We will be training with %d batches per epoch" % (num_training_batches_per_epoch))
#-------------------
# Setup training
#-------------------
if args.model == "unet":
model = models.get_unet()
elif args.model == "fcn":
model = models.get_fcn()
else:
raise ValueError("Invalid model")
model = model.to(device)
optimizer = optim.AdamW(model.parameters(), lr=0.001, amsgrad=True)
criterion = nn.CrossEntropyLoss()
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, "min")
print("Model has %d parameters" % (utils.count_parameters(model)))
#-------------------
# Model training
#-------------------
training_task_losses = []
num_times_lr_dropped = 0
model_checkpoints = []
temp_model_fn = os.path.join(args.output_dir, "most_recent_model.pt")
for epoch in range(args.num_epochs):
lr = utils.get_lr(optimizer)
training_losses = utils.fit(
model,
device,
dataloader,
num_training_batches_per_epoch,
optimizer,
criterion,
epoch,
)
scheduler.step(training_losses[0])
model_checkpoints.append(copy.deepcopy(model.state_dict()))
if args.save_most_recent:
torch.save(model.state_dict(), temp_model_fn)
if utils.get_lr(optimizer) < lr:
num_times_lr_dropped += 1
print("")
print("Learning rate dropped")
print("")
training_task_losses.append(training_losses[0])
if num_times_lr_dropped == 4:
break
#-------------------
# Save everything
#-------------------
save_obj = {
'args': args,
'training_task_losses': training_task_losses,
"checkpoints": model_checkpoints
}
save_obj_fn = "results.pt"
with open(os.path.join(args.output_dir, save_obj_fn), 'wb') as f:
torch.save(save_obj, f)
def main(
input_path,
checkpoint_path,
output_dir,
landmarks_path,
num_iterations,
csv_path,
batch_size,
num_workers,
gpu,
threshold,
augmentation,
save_volumes,
interpolation,
std_noise,
):
import torch
import pandas as pd
import numpy as np
import torchio as tio
from tqdm import tqdm
import models
device = torch.device(
'cuda' if torch.cuda.is_available() and gpu else 'cpu')
checkpoint = torch.load(checkpoint_path, map_location=device)
model = models.get_unet().to(device)
model.load_state_dict(checkpoint['model'])
output_dir = Path(output_dir)
model.eval()
torch.set_grad_enabled(False)
fps = get_paths(input_path)
mean_dir = output_dir / 'mean'
std_dir = output_dir / 'std'
# entropy_dir = output_dir / 'entropy'
mean_dir.mkdir(parents=True, exist_ok=True)
std_dir.mkdir(parents=True, exist_ok=True)
# entropy_dir.mkdir(parents=True, exist_ok=True)
records = []
progress = tqdm(fps, unit='subject')
for fp in progress:
subject_id = fp.name[:4]
progress.set_description(subject_id)
image = tio.ScalarImage(fp)
subject = tio.Subject(
image=image) # key must be 'image' as in get_test_transform
transform = get_transform(augmentation, landmarks_path)
dataset = tio.SubjectsDataset(num_iterations * [subject],
transform=transform)
loader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
num_workers=num_workers,
collate_fn=lambda x: x,
)
all_results = []
for subjects_list_batch in tqdm(loader, leave=False, unit='batch'):
inputs = torch.stack([
subject.image.data for subject in subjects_list_batch
]).float().to(device)
with torch.cuda.amp.autocast():
segs = model(inputs).softmax(dim=1)[:, 1:].cpu()
iterable = list(zip(subjects_list_batch, segs))
for subject, seg in tqdm(iterable, leave=False, unit='subject'):
subject.image.set_data(seg)
inverse_transform = subject.get_inverse_transform(warn=False)
inverse_transforms = inverse_transform.transforms
first = inverse_transforms[0]
if hasattr(first, 'image_interpolation'
) and first.image_interpolation != 'linear':
first.image_interpolation = 'linear' # force interp to be lin so probs stay in [0,1]
subject_back = inverse_transform(subject)
result = subject_back.image.data
assert np.count_nonzero(
result.numpy() < 0) == 0, 'neg values found in result'
if threshold:
result = (result >= 0.5).float()
all_results.append(result)
result = torch.stack(all_results)
volumes = result.sum(dim=(-3, -2, -1)).numpy()
mean_volumes = volumes.mean()
std_volumes = volumes.std()
volume_variation_coefficient = std_volumes / mean_volumes
q1, q3 = np.percentile(volumes, (25, 75))
quartile_coefficient_of_dispersion = (q3 - q1) / (q3 + q1)
record = dict(
Subject=subject_id,
VolumeMean=mean_volumes,
VolumeSTD=std_volumes,
VVC=volume_variation_coefficient,
Q1=q1,
Q3=q3,
QCD=quartile_coefficient_of_dispersion,
)
if save_volumes:
for i, volume in enumerate(volumes):
record[f'Volume_{i}'] = volume
records.append(record)
mean = result.mean(dim=0)
std = result.std(dim=0)
# entropy = utils.get_entropy(result)
mean_image = tio.ScalarImage(tensor=mean, affine=image.affine)
std_image = tio.ScalarImage(tensor=std, affine=image.affine)
# entropy_image = tio.ScalarImage(tensor=entropy, affine=image.affine)
mean_path = mean_dir / fp.name.replace('.nii', '_mean.nii')
std_path = std_dir / fp.name.replace('.nii', '_std.nii')
# entropy_path = entropy_dir / fp.name.replace('.nii', '_entropy.nii')
mean_image.save(mean_path)
std_image.save(std_path)
# entropy_image.save(entropy_path)
# So it's updated during execution
df = pd.DataFrame.from_records(records)
df.to_csv(csv_path)
return 0
'sgd': SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)
}
loss = loss_options[loss_param]
optimizer = optimizer_options[optimizer_param]
model_filename = "../saved_models/2018_final/{}.h5".format(model_name)
print(model_filename)
if do_ensemble == False:
print('Create model')
if model == 'unet':
# model = models.Unet(height,width, loss=loss, optimizer = optimizer, metrics = metrics, fc_size = fc_size, channels=n_channels)
model = models.get_unet(height,
width,
loss=loss,
optimizer=optimizer,
metrics=metrics,
channels=n_channels)
elif model == 'unet2':
model = models.Unet2(height,
width,
loss=loss,
optimizer=optimizer,
metrics=metrics,
fc_size=fc_size,
channels=n_channels)
elif model == 'vgg':
model = models.VGG16(height,
width,
pretrained=False,
freeze_pretrained=False,
def main(
input_path,
checkpoint_path,
output_dir,
landmarks_path,
num_iterations,
csv_path,
num_workers,
gpu,
threshold,
interpolation,
):
import torch
import pandas as pd
import numpy as np
import torchio as tio
from tqdm import tqdm, trange
import utils
import models
import datasets
device = torch.device(
'cuda' if torch.cuda.is_available() and gpu else 'cpu')
checkpoint = torch.load(checkpoint_path, map_location=device)
model = models.get_unet().to(device)
model.load_state_dict(checkpoint['model'])
output_dir = Path(output_dir)
model.eval()
utils.enable_dropout(model)
torch.set_grad_enabled(False)
fps = get_paths(input_path)
mean_dir = output_dir / 'mean'
std_dir = output_dir / 'std'
entropy_dir = output_dir / 'entropy'
mean_dir.mkdir(parents=True, exist_ok=True)
std_dir.mkdir(parents=True, exist_ok=True)
entropy_dir.mkdir(parents=True, exist_ok=True)
records = []
progress = tqdm(fps, unit='subject')
for fp in progress:
subject_id = fp.name[:4]
progress.set_description(subject_id)
image = tio.ScalarImage(fp)
subject = tio.Subject(
image=image) # key must be 'image' as in get_test_transform
preprocess = datasets.get_test_transform(landmarks_path)
preprocessed = preprocess(subject)
inputs = preprocessed.image.data.float()[np.newaxis].to(device)
all_results = []
for _ in trange(num_iterations, leave=False):
with torch.cuda.amp.autocast():
segs = model(inputs).softmax(dim=1)[0, 1:]
all_results.append(segs.cpu())
result = torch.stack(all_results)
volumes = result.sum(dim=(-3, -2, -1)).numpy()
mean_volumes = volumes.mean()
std_volumes = volumes.std()
volume_variation_coefficient = std_volumes / mean_volumes
q1, q3 = np.percentile(volumes, (25, 75))
quartile_coefficient_of_dispersion = (q3 - q1) / (q3 + q1)
records.append(
dict(
Subject=subject_id,
VolumeMean=mean_volumes,
VolumeSTD=std_volumes,
VVC=volume_variation_coefficient,
Q1=q1,
Q3=q3,
QCD=quartile_coefficient_of_dispersion,
))
crop: tio.Crop = preprocessed.history[-1]
pad = crop.inverse()
assert np.count_nonzero(
result.numpy() < 0) == 0, 'neg values found in result'
mean = result.mean(dim=0)
assert np.count_nonzero(
mean.numpy() < 0) == 0, 'neg values found in mean'
std = result.std(dim=0)
# entropy = utils.get_entropy(result)
mean_image = tio.ScalarImage(tensor=mean,
affine=preprocessed.image.affine)
std_image = tio.ScalarImage(tensor=std,
affine=preprocessed.image.affine)
# entropy_image = tio.ScalarImage(tensor=entropy, affine=preprocessed.image.affine)
mean_path = mean_dir / fp.name.replace('.nii', '_mean.nii')
std_path = std_dir / fp.name.replace('.nii', '_std.nii')
# entropy_path = entropy_dir / fp.name.replace('.nii', '_entropy.nii')
pad(mean_image).save(mean_path)
pad(std_image).save(std_path)
# pad(entropy_image).save(entropy_path)
# So it's updated while it runs
df = pd.DataFrame.from_records(records)
df.to_csv(csv_path)
return 0
def main():
print("Starting DFC2021 model inference script at %s" %
(str(datetime.datetime.now())))
#-------------------
# Setup
#-------------------
assert os.path.exists(args.input_fn)
assert os.path.exists(args.model_fn)
if os.path.isfile(args.output_dir):
print("A file was passed as `--output_dir`, please pass a directory!")
return
if os.path.exists(
args.output_dir) and len(os.listdir(args.output_dir)) > 0:
if args.overwrite:
print(
"WARNING! The output directory, %s, already exists, we might overwrite data in it!"
% (args.output_dir))
else:
print(
"The output directory, %s, already exists and isn't empty. We don't want to overwrite and existing results, exiting..."
% (args.output_dir))
return
else:
print("The output directory doesn't exist or is empty.")
os.makedirs(args.output_dir, exist_ok=True)
if torch.cuda.is_available():
device = torch.device("cuda:%d" % args.gpu)
else:
print(
"WARNING! Torch is reporting that CUDA isn't available, exiting..."
)
return
#-------------------
# Load model
#-------------------
if args.model == "unet":
model = models.get_unet()
elif args.model == "fcn":
model = models.get_fcn()
elif args.model == "hrnet":
model = models.get_hrnet()
else:
raise ValueError("Invalid model")
model.load_state_dict(torch.load(args.model_fn))
model = model.to(device)
#-------------------
# Run on each line in the input
#-------------------
input_dataframe = pd.read_csv(args.input_fn)
image_fns = input_dataframe["image_fn"].values
groups = input_dataframe["group"].values
for image_idx in range(len(image_fns)):
tic = time.time()
image_fn = image_fns[image_idx]
group = groups[image_idx]
print("(%d/%d) Processing %s" % (image_idx, len(image_fns), image_fn),
end=" ... ")
#-------------------
# Load input and create dataloader
#-------------------
def image_transforms(img):
if group == 0:
img = (img - utils.NAIP_2013_MEANS) / utils.NAIP_2013_STDS
elif group == 1:
img = (img - utils.NAIP_2017_MEANS) / utils.NAIP_2017_STDS
else:
raise ValueError("group not recognized")
img = np.rollaxis(img, 2, 0).astype(np.float32)
img = torch.from_numpy(img)
return img
with rasterio.open(image_fn) as f:
input_width, input_height = f.width, f.height
input_profile = f.profile.copy()
dataset = TileInferenceDataset(image_fn,
chip_size=CHIP_SIZE,
stride=CHIP_STRIDE,
transform=image_transforms,
verbose=False)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
num_workers=NUM_WORKERS,
pin_memory=True,
)
#-------------------
# Run model and organize output
#-------------------
output = np.zeros((len(utils.NLCD_CLASSES), input_height, input_width),
dtype=np.float32)
kernel = np.ones((CHIP_SIZE, CHIP_SIZE), dtype=np.float32)
kernel[HALF_PADDING:-HALF_PADDING, HALF_PADDING:-HALF_PADDING] = 5
counts = np.zeros((input_height, input_width), dtype=np.float32)
for i, (data, coords) in enumerate(dataloader):
data = data.to(device)
with torch.no_grad():
t_output = model(data)
t_output = F.softmax(t_output, dim=1).cpu().numpy()
for j in range(t_output.shape[0]):
y, x = coords[j]
output[:, y:y + CHIP_SIZE,
x:x + CHIP_SIZE] += t_output[j] * kernel
counts[y:y + CHIP_SIZE, x:x + CHIP_SIZE] += kernel
output = output / counts
output_hard = output.argmax(axis=0).astype(np.uint8)
#-------------------
# Save output
#-------------------
output_profile = input_profile.copy()
output_profile["driver"] = "GTiff"
output_profile["dtype"] = "uint8"
output_profile["count"] = 1
output_profile["nodata"] = 0
output_fn = image_fn.split("/")[
-1] # something like "546_naip-2013.tif"
output_fn = output_fn.replace("naip", "predictions")
output_fn = os.path.join(args.output_dir, output_fn)
if args.save_4lc:
output_profile["nodata"] = 4
output_fn = output_fn.replace("_predictions-", "_lcpredictions-")
output_hard = utils.NLCD_IDX_TO_REDUCED_LC_MAP[output_hard].astype(
np.uint8)
with rasterio.open(output_fn, "w", **output_profile) as f:
f.write(output_hard, 1)
f.write_colormap(1, utils.LC4_CLASS_COLORMAP)
else:
with rasterio.open(output_fn, "w", **output_profile) as f:
f.write(output_hard, 1)
f.write_colormap(1, utils.NLCD_IDX_COLORMAP)
if args.save_soft:
output = output / output.sum(axis=0, keepdims=True)
output = (output * 255).astype(np.uint8)
output_profile = input_profile.copy()
output_profile["driver"] = "GTiff"
output_profile["dtype"] = "uint8"
output_profile["count"] = len(utils.NLCD_CLASSES)
del output_profile["nodata"]
output_fn = image_fn.split("/")[
-1] # something like "546_naip-2013.tif"
output_fn = output_fn.replace("naip", "predictions-soft")
output_fn = os.path.join(args.output_dir, output_fn)
with rasterio.open(output_fn, "w", **output_profile) as f:
f.write(output)
print("finished in %0.4f seconds" % (time.time() - tic))
hyperparameters = hyperparameters_processing(hyperparameters)
folds = get_folds(hyperparameters)
autocontext_step = hyperparameters['autocontext_step']
model_predictions = [None] * len(
os.listdir(
os.path.join(os.getcwd(),
os.path.join(hyperparameters['data_path'], 'Image'))))
for s_step in range(0, autocontext_step):
folds = get_folds(hyperparameters)
for fold_num, fold in enumerate(folds):
train_images, train_masks, validation_images, validation_masks = fold
print("train images", train_images)
train_data_gen, test_data_gen = get_data_with_generator_on_the_fly(
hyperparameters, train_images, train_masks, validation_images,
validation_masks, s_step, fold_num, len(validation_images))
model = get_unet(hyperparameters)
# for i in range(10):
# batch_x_w, batch_y = train_data_gen.__getitem__(i)
# plot_triplet(batch_x_w[0][0,:,:,0], batch_x_w[1][0,:,:,0], batch_y[0,:,:,0])
# batch_x_w, batch_y = test_data_gen.__getitem__(i)
# plot_triplet(batch_x_w[0][0,:,:,0], batch_x_w[1][0,:,:,0], batch_y[0,:,:,0])
model_history = model.fit_generator(
train_data_gen,
# steps_per_epoch=total_train // hyperparameters['batch_size'],
epochs=hyperparameters['epochs'],
validation_data=test_data_gen)
plot_history(hyperparameters, model_history, task, s_step,
fold_num)
if 'save_model' in hyperparameters and hyperparameters[
'save_model']:
save_model(model, task, s_step, fold_num)
from models import get_model_incept, get_unet
import os
import tifffile as tiff
import numpy as np
trainX = []
trainY = []
for root, dirs, files in os.walk('./train'):
for dir_name in dirs:
image = tiff.imread('./train/' + str(dir_name) + '/' + str(dir_name) +
'.tif')
trainX.append(image / np.max(image))
image = tiff.imread('./train/' + str(dir_name) + '/mask.tif')
trainY.append(image / 255)
model = get_unet() # or get_model_incept
checkpoint = [
ModelCheckpoint(filepath='/content/checks',
save_weights_only=False,
monitor='val_accuracy',
mode='max',
save_best_only=True),
ReduceLROnPlateau(monitor='val_loss',
factor=0.1,
mode='min',
min_lr=0.00001,
min_delta=0.001,
verbose=1,
patience=10)
]
def get_model(self):
return setup_model(models.get_unet()) # pylint: disable=E1120
