| """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 |
|
|
| 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 |
| from modality import one_hot_condition |
| from monai.networks.schedulers import DDPMScheduler |
|
|
| 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]] |
|
|
|
|
| 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") |
| 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") |
|
|
