File size: 5,660 Bytes
3c8ff2e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 |
"""
Script for image reconstruction inference with pre-trained models
Author: Patrick Ebel (github/PatrickTUM), based on the scripts of
Vivien Sainte Fare Garnot (github/VSainteuf)
License: MIT
"""
import os
import sys
import json
import pprint
import argparse
from parse_args import create_parser
import torch
dirname = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(dirname))
from src import utils
from src.model_utils import get_model, load_checkpoint
from train_reconstruct import (
iterate,
save_results,
prepare_output,
import_from_path,
seed_packages,
)
from data.dataLoader import SEN12MSCR, get_pairedS1
from torch.utils.tensorboard import SummaryWriter
parser = create_parser(mode="test")
test_config = parser.parse_args()
# grab the PID so we can look it up in the logged config for server-side process management
test_config.pid = os.getpid()
# related to flag --use_custom:
# define custom target S2 patches (these will be mosaiced into a single sample), and fetch associated target S1 patches as well as input data
# (TODO: keeping this hard-coded until a more convenient way to pass it as an argument comes about ...)
targ_s2 = [
f"ROIs1868/73/S2/14/s2_ROIs1868_73_ImgNo_14_2018-06-21_patch_{pdx}.tif"
for pdx in [171, 172, 173, 187, 188, 189, 203, 204, 205]
]
# load previous config from training directories
# if no custom path to config file is passed, try fetching config file at default location
conf_path = (
os.path.join(
dirname, test_config.weight_folder, test_config.experiment_name, "conf.json"
)
if not test_config.load_config
else test_config.load_config
)
if os.path.isfile(conf_path):
with open(conf_path) as file:
model_config = json.loads(file.read())
t_args = argparse.Namespace()
# do not overwrite the following flags by their respective values in the config file
no_overwrite = [
"pid",
"device",
"resume_at",
"trained_checkp",
"res_dir",
"weight_folder",
"root1",
"root2",
"root3",
"max_samples_count",
"batch_size",
"display_step",
"plot_every",
"export_every",
"input_t",
"region",
"min_cov",
"max_cov",
]
conf_dict = {
key: val for key, val in model_config.items() if key not in no_overwrite
}
for key, val in vars(test_config).items():
if key in no_overwrite:
conf_dict[key] = val
t_args.__dict__.update(conf_dict)
config = parser.parse_args(namespace=t_args)
else:
config = test_config # otherwise, keep passed flags without any overwriting
config = utils.str2list(config, ["encoder_widths", "decoder_widths", "out_conv"])
if config.pretrain:
config.batch_size = 32
experime_dir = os.path.join(config.res_dir, config.experiment_name)
if not os.path.exists(experime_dir):
os.makedirs(experime_dir)
with open(os.path.join(experime_dir, "conf.json"), "w") as file:
file.write(json.dumps(vars(config), indent=4))
# seed everything
seed_packages(config.rdm_seed)
if __name__ == "__main__":
pprint.pprint(config)
# instantiate tensorboard logger
writer = SummaryWriter(os.path.join(config.res_dir, config.experiment_name))
if config.use_custom:
print("Testing on custom data samples")
# define a dictionary for the custom sample, with customized ROI and time points
custom = [
{
"input": {
"S1": [
get_pairedS1(targ_s2, config.root1, mod="s1", time=tdx)
for tdx in range(0, 3)
],
"S2": [
get_pairedS1(targ_s2, config.root1, mod="s2", time=tdx)
for tdx in range(0, 3)
],
},
"target": {
"S1": [get_pairedS1(targ_s2, config.root1, mod="s1")],
"S2": [targ_s2],
},
}
]
def main(config):
device = torch.device(config.device)
prepare_output(config)
model = get_model(config)
model = model.to(device)
config.N_params = utils.get_ntrainparams(model)
print(f"TOTAL TRAINABLE PARAMETERS: {config.N_params}\n")
# print(model)
# get data loader
if config.pretrain:
dt_test = SEN12MSCR(
config.root3,
split="test",
region=config.region,
sample_type=config.sample_type,
)
else:
pass
dt_test = torch.utils.data.Subset(
dt_test, range(0, min(config.max_samples_count, len(dt_test)))
)
test_loader = torch.utils.data.DataLoader(
dt_test, batch_size=config.batch_size, shuffle=False
)
# Load weights
ckpt_n = f"_epoch_{config.resume_at}" if config.resume_at > 0 else ""
load_checkpoint(config, config.weight_folder, model, f"model{ckpt_n}")
# Inference
print("Testing . . .")
model.eval()
_, test_img_metrics = iterate(
model,
data_loader=test_loader,
config=config,
writer=writer,
mode="test",
epoch=1,
device=device,
)
print(f"\nTest image metrics: {test_img_metrics}")
save_results(
test_img_metrics,
os.path.join(config.res_dir, config.experiment_name),
split="test",
)
print(
f"\nLogged test metrics to path {os.path.join(config.res_dir, config.experiment_name)}"
)
if __name__ == "__main__":
main(config)
exit()
|