forked from educating-dip/diffusion_models_dev_project
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathrun_adapted_sampling.py
141 lines (123 loc) · 6.52 KB
/
run_adapted_sampling.py
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
import yaml
import argparse
import torch
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
from itertools import islice
from src import (get_standard_sde, PSNR, SSIM, get_standard_dataset, get_data_from_ground_truth, get_standard_ray_trafo,
get_standard_score, get_standard_configs, get_standard_path, get_standard_adapted_sampler)
parser = argparse.ArgumentParser(description='conditional sampling')
parser.add_argument('--dataset', default='walnut', help='test-dataset', choices=['walnut', 'lodopab', 'ellipses', 'mayo', 'aapm'])
parser.add_argument('--model', default='openai_unet', help='select unet arch.', choices=['dds_unet', 'openai_unet'])
parser.add_argument('--base_path', default='/localdata/AlexanderDenker/score_based_baseline', help='path to model configs')
parser.add_argument('--model_learned_on', default='lodopab', help='model-checkpoint to load', choices=['lodopab', 'ellipses', 'aapm'])
parser.add_argument('--method', default='naive', choices=['naive', 'dps', 'dds'])
parser.add_argument('--version', default=1, help="version of the model")
parser.add_argument('--noise_level', default=0.01, help="rel. additive gaussian noise.")
parser.add_argument('--add_corrector_step', action='store_true')
parser.add_argument('--ema', action='store_true')
parser.add_argument('--num_steps', default=50)
parser.add_argument('--penalty', default=1, help='reg. penalty used for ``naive'' and ``dps'' only.')
parser.add_argument('--tv_penalty', default=1e-6, help='reg. used for ``adapatation''.')
parser.add_argument('--eta', default=0.85, help='reg. used for ``dds'' weighting stochastic and deterministic noise.')
parser.add_argument('--sde', default='vesde', choices=['vpsde', 'vesde', 'ddpm'])
parser.add_argument('--adaptation', default='lora', choices=['decoder', 'full', 'vdkl', 'lora'])
parser.add_argument('--num_optim_step', default=10, help='num. of optimization steps taken per sampl. step')
parser.add_argument('--adapt_freq', default=1, help='freq. of adaptation step in sampl.')
parser.add_argument('--lora_include_blocks', default=['input_blocks','middle_block','output_blocks','out'], nargs='+', help='lora kwargs impl. of arch. blocks included')
parser.add_argument('--lr', default=1e-3, help='learning rate for adaptation')
parser.add_argument('--lora_rank', default=4, help='lora kwargs impl. of rank')
parser.add_argument('--add_cg', action='store_true', help='do DDS steps after adaptation.')
parser.add_argument('--cg_iter', default=1, help='Number of CG steps for DDS update.')
parser.add_argument('--gamma', default=0.01, help='reg. used for ``dds''.')
parser.add_argument('--load_path', help='path to ddpm model.')
parser.add_argument('--dc_type', default="cg", choices=["cg", "gd", "none"], help="use cg/gd in adaptation (or none at all)")
parser.add_argument('--stddev', default=None, help="noise_level")
parser.add_argument('--early_stopping_pct', default=1.0)
def coordinator(args):
config, dataconfig = get_standard_configs(args, base_path=args.base_path)
#dataconfig.data.stddev = float(args.noise_level)
try:
save_root = get_standard_path(args, run_type="adapt", data_part=dataconfig.data.part)
except AttributeError:
save_root = get_standard_path(args, run_type="adapt")
print("save to: ", save_root)
save_root.mkdir(parents=True, exist_ok=True)
if config.seed is not None:
torch.manual_seed(config.seed) # for reproducible noise in simulate
sde = get_standard_sde(config=config)
score = get_standard_score(config=config, sde=sde, use_ema=args.ema, model_type=args.model)
score = score.to(config.device).eval()
ray_trafo = get_standard_ray_trafo(config=dataconfig)
ray_trafo = ray_trafo.to(device=config.device)
#from odl.operator.oputils import power_method_opnorm
#print("OPERATOR NORM: ", power_method_opnorm(ray_trafo.ray_trafo_op_fun.operator))
dataset = get_standard_dataset(config=dataconfig, ray_trafo=ray_trafo)
dataconfig.data.validation.num_images = len(dataset)
_psnr, _ssim = [], []
for i, data_sample in enumerate(islice(dataset, dataconfig.data.validation.num_images)):
if config.seed is not None:
torch.manual_seed(config.seed + i) # for reproducible noise in simulate
if len(data_sample) == 3:
observation, ground_truth, filtbackproj = data_sample
ground_truth = ground_truth.to(device=config.device)
observation = observation.to(device=config.device)
filtbackproj = filtbackproj.to(device=config.device)
else:
if len(data_sample) == 1 and args.dataset == "ellipses" and dataconfig.data.part == "test":
data_sample = data_sample[0]
ground_truth, observation, filtbackproj = get_data_from_ground_truth(
ground_truth=data_sample.to(device=config.device),
ray_trafo=ray_trafo,
white_noise_rel_stddev=dataconfig.data.stddev
)
logg_kwargs = {'log_dir': save_root, 'num_img_in_log': 1, 'sample_num': i, 'ground_truth': ground_truth, 'filtbackproj': filtbackproj}
sampler = get_standard_adapted_sampler(
args=args,
config=config,
score=score,
sde=sde,
device=config.device,
observation = observation,
ray_trafo = ray_trafo
)
recon = sampler.sample(logg_kwargs=logg_kwargs, logging=True)
recon = torch.clamp(recon, 0)
torch.save( {'recon': recon.cpu().squeeze(), 'ground_truth': ground_truth.cpu().squeeze()},
str(save_root / f'recon_{i}_info.pt') )
im = Image.fromarray(recon.cpu().squeeze().numpy()*255.).convert("L")
im.save(str(save_root / f'recon_{i}.png'))
score = get_standard_score(config=config, sde=sde, use_ema=args.ema, model_type=args.model)
score = score.to(config.device)
score.eval()
print(f'reconstruction of sample {i}')
psnr = PSNR(recon[0, 0].cpu().numpy(), ground_truth[0, 0].cpu().numpy())
ssim = SSIM(recon[0, 0].cpu().numpy(), ground_truth[0, 0].cpu().numpy())
_psnr.append(psnr)
_ssim.append(ssim)
print('PSNR:', psnr)
print('SSIM:', ssim)
"""
_, (ax1, ax2, ax3) = plt.subplots(1,3)
ax1.imshow(ground_truth[0,0,:,:].detach().cpu(), cmap='gray')
ax1.axis('off')
ax1.set_title('Ground truth')
ax2.imshow(torch.clamp(recon[0,0,:,:], 0, 1).detach().cpu(), cmap='gray')
ax2.axis('off')
ax2.set_title('Adaptation Sampling')
ax3.imshow(filtbackproj[0,0,:,:].detach().cpu(), cmap='gray')
ax3.axis('off')
ax3.set_title('FBP')
plt.savefig(str(save_root/f'info_{i}.png'))
"""
report = {}
report.update(dict(dataconfig.items()))
report.update(vars(args))
report["PSNR"] = float(np.mean(_psnr))
report["SSIM"] = float(np.mean(_ssim))
with open(save_root / 'report.yaml', 'w') as file:
yaml.dump(report, file)
if __name__ == '__main__':
args = parser.parse_args()
coordinator(args)