File size: 3,154 Bytes
d4700c7 | 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 | """Does the diffusion model keep a tumour it is given? Real FLAIR latents with tumours are noised to step
T, denoised by the model, decoded; the bright-voxel fraction (tumour_check.py) before and after, plus a figure."""
import json
import os
import sys
from pathlib import Path
import matplotlib
import torch
matplotlib.use("Agg")
import matplotlib.pyplot as plt # noqa: E402
LDM = Path.home() / "scratch/ldm1305"
APP = Path.home() / "sandbox/londonaicentre/flip-wt-1305-ldm/fl-tutorials/nvflare/image_synthesis/latent_diffusion_model/app_files"
sys.path.insert(0, str(APP))
os.chdir(APP)
import models # noqa: E402
from modality import one_hot_condition # noqa: E402
from monai.networks.schedulers import DDPMScheduler # noqa: E402
torch.manual_seed(0)
model = models.get_model().eval()
blob = torch.load(LDM / "ws/flip_fedavg/server/simulate_job/FL_global_model.pt", map_location="cpu", weights_only=False)
model.load_state_dict(blob["model"] if isinstance(blob.get("model"), dict) else blob)
scale = json.loads((LDM / "config.used.json").read_text())["LATENT_SCALE_FACTOR"]
real = torch.load(LDM / "tumour/e5_FLAIR.pt", map_location="cpu", weights_only=False)["real"][[0, 3, 4, 6]] # labelled tumours
def bright(v, tau=1.5):
v = v.float().flatten()
lo, hi = torch.quantile(v[::7], torch.tensor([0.01, 0.995]))
u = ((v - lo) / (hi - lo + 1e-8)).clamp(0, 1)
brain = u > 0.1
return float((brain & (u > tau * u[brain].median())).sum() / brain.sum())
cond = one_hot_condition(torch.zeros(len(real), dtype=torch.long), 4, "cpu") # FLAIR
rows, out = {"real": real}, {"real": [bright(v) for v in real]}
with torch.no_grad():
z0 = model.autoencoder.encode_stage_2_inputs(real) * scale
rows["recon"] = model.autoencoder.decode_stage_2_outputs(z0 / scale)
for T in (250, 500, 750):
sched = DDPMScheduler(num_train_timesteps=1000, schedule="scaled_linear_beta", clip_sample=False,
prediction_type="epsilon", beta_start=0.0015, beta_end=0.0195)
sched.set_timesteps(1000)
z = sched.add_noise(z0, torch.randn_like(z0), torch.full((len(real),), T))
for t in range(T, -1, -1):
eps = model.diffusion_model(z, timesteps=torch.full((len(real),), t), context=cond)
z, _ = sched.step(eps, t, z)
rows[f"T={T}"] = model.autoencoder.decode_stage_2_outputs(z / scale)
print(T, flush=True)
for k, vols in rows.items():
out[k] = [bright(v) for v in vols]
print(json.dumps({k: [round(x, 3) for x in v] for k, v in out.items()}))
(LDM / "tumour/sdedit.json").write_text(json.dumps(out))
f, ax = plt.subplots(len(real), len(rows), figsize=(2.2 * len(rows), 2.3 * len(real)), facecolor="k")
for j, (k, vols) in enumerate(rows.items()):
for i, v in enumerate(vols):
v = v[0].float()
lo, hi = torch.quantile(v.flatten()[::7], torch.tensor([0.01, 0.995]))
ax[i, j].imshow(v[:, :, v.shape[2] // 2 + 4].T.flip(0), cmap="gray", vmin=lo, vmax=hi)
ax[i, j].axis("off")
ax[0, j].set_title(k, color="w", fontsize=10)
plt.tight_layout()
plt.savefig(LDM / "tumour/sdedit.png", dpi=80, facecolor="k")
|