OSCR

Dual-conditioned diffusion model with anatomical guidance for geometric distortion correction in prostate MRI.

Code ↔ Paper

5 matches between paragraphs of the paper and lines of its authors' code, computed by the harvester (lexical-v1). Click a colored paragraph or line to see its counterpart.

The 5 matches
  1. [1] § Methods › DeDistortNet model architecture and training ↔ source/train_DeDistortNet.py, lines 958–996 · score 0.71 · forward diffusion process, noisy latent, predicted, timestep, noise, training
  2. [2] § Methods › Evaluation metrics ↔ source/train_DeDistortNet.py, lines 1–51 · score 0.58 · Peak Signal, Noise Ratio, PSNR, SSIM
  3. [3] § Methods › DeDistortNet model architecture and training ↔ source/train_DeDistortNet.py, lines 53–137 · score 0.56 · CLIP image encoder, Stable Diffusion, prompts, pretrained, rotation, weights
  4. [4] § Methods › Comparison models ↔ source/train_DeDistortNet.py, lines 53–137 · score 0.54 · CLIP image encoder, Stable Diffusion, prompt, guidance, models, training
  5. [5] § Methods › DeDistortNet model architecture and training ↔ source/train_DeDistortNet.py, lines 787–866 · score 0.50 · AdamW, optimizer, channels, batch, Training, weighted

Paper

Loaded from Europe PMC by your browser, not stored by OSCR: doi.org · Europe PMC

The paper is loaded when this pane is shown.

The authors' code

Python · 1,060 lines · 44 KB · Apache-2.0 · 5 matches

  1. import os
  3. import random
  4. import argparse
  5. from pathlib import Path
  6. import json
  7. import itertools
  8. import time
  9. import logging
  10. import math
  12. import elasticdeform
  13. import accelerate
  14. import numpy as np
  15. from tqdm.auto import tqdm
  17. import matplotlib.pyplot as plt
  19. import torch
  20. import torch.nn.functional as F
  21. from torchvision import transforms
  22. import SimpleITK as sitk
  23. from transformers import CLIPImageProcessor
  24. from accelerate import Accelerator
  25. from accelerate.logging import get_logger
  26. from accelerate.utils import ProjectConfiguration, set_seed
  27. from packaging import version
  28. from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel, DDIMScheduler
  29. from diffusers import ControlNetModel
  30. from diffusers.utils.torch_utils import is_compiled_module
  31. from transformers import CLIPVisionModelWithProjection
  33. from diffusers import StableDiffusionControlNetPipeline
  35. from torchmetrics.image import StructuralSimilarityIndexMeasure
  36. from torchmetrics.image import PeakSignalNoiseRatio
  38. from scipy.ndimage import zoom
  39. from skimage import transform
  41. PSNR = PeakSignalNoiseRatio(data_range=1.0)
  42. SSIM = StructuralSimilarityIndexMeasure(data_range=1.0)
  44. print("PyTorch version:", torch.__version__)
  45. print("Is CUDA available:", torch.cuda.is_available())
  46. print("cuDNN version:", torch.backends.cudnn.version())
  47. print("Is cuDNN enabled:", torch.backends.cudnn.enabled)
  49. import wandb
  51. logger = get_logger(__name__)
  53. def log_validation(
  54. args, accelerator, weight_dtype, step, checkpoint_path
  55. ):
  56. weight_dtype = torch.float16
  57. logger.info("Running validation... ")
  59. val_dataset = MyDataset(args.val_data_json_file, size=args.resolution, displacement_rate=args.displacement_rate, random_y_squeeze_rate=args.random_y_squeeze_rate, random_rotation_degree=args.random_rotation_degree, random_shift_range=args.random_shift_range, image_root_path=args.data_root_path, clip_image_processor_path=args.image_encoder_path)
  60. val_dataloader = torch.utils.data.DataLoader(
  61. val_dataset,
  62. shuffle=False,
  63. collate_fn=collate_fn,
  64. batch_size=args.train_batch_size,
  65. num_workers=args.dataloader_num_workers,
  66. pin_memory=True
  67. )
  69. controlnet = ControlNetModel.from_pretrained(os.path.join(checkpoint_path, "controlnet"))
  70. controlnet = controlnet.to(accelerator.device, dtype=weight_dtype)
  72. vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae").to(accelerator.device, dtype=weight_dtype)
  74. noise_scheduler = DDIMScheduler(
  75. num_train_timesteps=1000,
  76. beta_start=0.00085,
  77. beta_end=0.012,
  78. beta_schedule="scaled_linear",
  79. clip_sample=False,
  80. set_alpha_to_one=False,
  81. steps_offset=1,
  82. )
  83. image_encoder = CLIPVisionModelWithProjection.from_pretrained(os.path.join(checkpoint_path, 'image_encoder')).to(accelerator.device, dtype=weight_dtype)
  84. unet = UNet2DConditionModel.from_pretrained(checkpoint_path, subfolder="unet").to(accelerator.device, dtype=weight_dtype)
  85. pipeline = StableDiffusionControlNetPipeline.from_pretrained(
  86. args.pretrained_model_name_or_path,
  87. vae=vae,
  88. unet=unet,
  89. controlnet=controlnet,
  90. safety_checker=None,
  91. feature_extractor=None,
  92. torch_dtype=weight_dtype,
  93. scheduler=noise_scheduler,
  94. )
  95. pipeline = pipeline.to(accelerator.device, dtype=weight_dtype)
  96. pipeline.set_progress_bar_config(disable=True)
  98. if args.seed is None:
  99. generator = None
  100. else:
  101. generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
  103. val_psnr = 0.0
  104. val_ssim = 0.0
  106. val_psnr_cropped = 0.0
  107. val_ssim_cropped = 0.0
  109. for v_i, val_example in enumerate(tqdm(val_dataloader)):
  110. image = val_example["images"]
  111. control_image = val_example["control_images"]
  112. clip_image = val_example["clip_images"].to(accelerator.device, dtype=weight_dtype)
  113. with torch.no_grad():
  114. encoder_hidden_states = image_encoder(clip_image).last_hidden_state
  116. with torch.autocast("cuda"):
  117. result_image = pipeline(
  118. prompt_embeds=encoder_hidden_states,
  119. image=control_image,
  120. num_inference_steps=100,
  121. guidance_scale=0.0,
  122. generator=generator,
  123. output_type="np",
  124. ).images
  126. image = image * 0.5 + 0.5
  127. image = image.cpu().float()
  128. control_image = control_image * 0.5 + 0.5
  129. control_image = control_image.cpu().float()
  130. result_image = np.transpose(result_image, (0, 3, 1, 2))
  131. result_image = torch.from_numpy(result_image).float()
  132. val_psnr += PSNR(result_image, image)
  133. val_ssim += SSIM(result_image, image)
  135. for i in range(image.shape[0]):
  136. # Find non-zero pixel indices in the first channel
  137. non_zero_indices = torch.nonzero(control_image[i, 0])
  139. # Get bounding box of non-zero region
  140. min_row, min_col = torch.min(non_zero_indices, dim=0)[0]
  141. max_row, max_col = torch.max(non_zero_indices, dim=0)[0]
  142. result_image_cropped = result_image[i:i+1].clone()
  143. result_image_cropped = result_image_cropped[:, :, min_row:max_row+1, min_col:max_col+1]
  144. image_cropped = image[i:i+1, :, min_row:max_row+1, min_col:max_col+1]
  145. val_psnr_cropped += PSNR(result_image_cropped, image_cropped)
  146. val_ssim_cropped += SSIM(result_image_cropped, image_cropped)
  148. val_psnr /= len(val_dataloader)
  149. val_ssim /= len(val_dataloader)
  151. val_psnr_cropped /= len(val_dataset)
  152. val_ssim_cropped /= len(val_dataset)
  154. logger.info(f"Validation PSNR: {val_psnr}, Cropped: {val_psnr_cropped}")
  155. logger.info(f"Validation SSIM: {val_ssim}, Cropped: {val_ssim_cropped}")
  157. wandb.log({"val_psnr": val_psnr, "val_ssim": val_ssim, "step": step, "cropped_val_psnr": val_psnr_cropped, "cropped_val_ssim": val_ssim_cropped})
  159. logger.info("Run plot validation...")
  161. val_dataset = ValDataset(args.plot_data_json_file, size=args.resolution, image_root_path=args.data_root_path, clip_image_processor_path=args.image_encoder_path)
  162. fig, axs = plt.subplots(len(val_dataset), 8, figsize=(40, 5 * len(val_dataset)))
  164. for v_i, val_example in enumerate(tqdm(val_dataset)):
  165. image = val_example["image"]
  166. control_image = val_example["control_image"]
  167. gt_seg_image = val_example["gt_seg_image"]
  168. control_image = control_image.unsqueeze(0)
  169. clip_image = val_example["clip_image"].to(accelerator.device, dtype=weight_dtype)
  170. with torch.no_grad():
  171. encoder_hidden_states = image_encoder(clip_image).last_hidden_state
  173. with torch.autocast("cuda"):
  174. result_image = pipeline(
  175. prompt_embeds=encoder_hidden_states,
  176. image=control_image,
  177. num_inference_steps=100,
  178. guidance_scale=0.0,
  179. generator=generator,
  180. output_type="np",
  181. ).images[0]
  183. channel_names = val_example["channel_names"]
  184. # Plot context images and generated images for each channel
  185. for ch in range(3):
  186. vmax = max(image[:,:,ch].max(), result_image[:,:,ch].max())
  187. axs[v_i, ch].imshow(image[:,:,ch], cmap="gray", vmin=0, vmax=vmax)
  188. axs[v_i, ch].set_title("Context Image " + channel_names[ch])
  189. axs[v_i, ch].axis("off")
  191. axs[v_i, ch+4].imshow(result_image[:,:,ch], cmap="gray", vmin=0, vmax=vmax)
  192. axs[v_i, ch+4].set_title("Generated Image " + channel_names[ch])
  193. axs[v_i, ch+4].axis("off")
  195. axs[v_i, 3].imshow(control_image[0,0], cmap="gray", vmin=-1, vmax=1)
  196. axs[v_i, 3].set_title("Structure Image")
  197. axs[v_i, 3].axis("off")
  199. axs[v_i, 7].imshow(gt_seg_image, cmap="viridis", vmin=0, vmax=3)
  200. axs[v_i, 7].set_title("Ground Truth Segmentation")
  201. axs[v_i, 7].axis("off")
  203. os.makedirs(os.path.join(args.output_dir, "validation_plots"), exist_ok=True)
  205. plt.tight_layout()
  206. plt.savefig(os.path.join(args.output_dir, "validation_plots", f"validation_{step}.png"), dpi=300)
  207. plt.close()
  208. del pipeline
  209. del controlnet
  211. return val_psnr, val_ssim, val_psnr_cropped, val_ssim_cropped
  213. class ValDataset(torch.utils.data.Dataset):
  214. def __init__(self, json_file, size=512, image_root_path="", clip_image_processor_path=None):
  215. super().__init__()
  217. self.size = size
  218. self.image_root_path = image_root_path
  220. train_min_max_file = os.path.join(self.image_root_path, "train_min_max.json")
  221. with open(train_min_max_file, "r") as f:
  222. train_min_max = json.load(f)
  224. self.b50_min = train_min_max["B50"]["prostate_min"]
  225. self.b50_max = train_min_max["B50"]["prostate_max"]
  226. self.b400_min = train_min_max["B400"]["prostate_min"]
  227. self.b400_max = train_min_max["B400"]["prostate_max"]
  228. self.b800_min = train_min_max["B800"]["prostate_min"]
  229. self.b800_max = train_min_max["B800"]["prostate_max"]
  230. self.t2_min = train_min_max["T2"]["prostate_min"]
  231. self.t2_max = train_min_max["T2"]["prostate_max"]
  233. self.data = []
  234. with open(json_file, "r") as f:
  235. for line in f:
  236. self.data.append(json.loads(line))
  238. self.transform = transforms.Compose([
  239. transforms.ToTensor(),
  240. transforms.Normalize([0.5], [0.5]),
  241. ])
  242. if clip_image_processor_path is None:
  243. self.clip_image_processor = CLIPImageProcessor()
  244. else:
  245. self.clip_image_processor = CLIPImageProcessor.from_pretrained(clip_image_processor_path)
  247. def __getitem__(self, idx):
  248. item = self.data[idx]
  249. image_file = item["image"]
  250. crop_coor = item["crop_coor"]
  252. b50_image = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(self.image_root_path, 'B50', image_file)))
  253. b400_image = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(self.image_root_path, 'B400', image_file)))
  254. b800_image = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(self.image_root_path, 'B800', image_file)))
  255. control_image = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(self.image_root_path, 'T2', image_file))).astype(np.float32)
  256. gt_seg_image = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(self.image_root_path, 'PZTZ', image_file)))
  258. # global min-max normalization
  259. b50_image = (b50_image - self.b50_min) / (self.b50_max - self.b50_min)
  260. b400_image = (b400_image - self.b400_min) / (self.b400_max - self.b400_min)
  261. b800_image = (b800_image - self.b800_min) / (self.b800_max - self.b800_min)
  263. # clip 0.0 ~ 1.0
  264. b50_image = np.clip(b50_image, 0, 1)
  265. b400_image = np.clip(b400_image, 0, 1)
  266. b800_image = np.clip(b800_image, 0, 1)
  268. image = np.stack([b50_image, b400_image, b800_image], axis=-1)
  270. channel_names = ["B50", "B400", "B800"]
  271. image = image[crop_coor['y_start']:crop_coor['y_start']+crop_coor['y_size'], crop_coor['x_start']:crop_coor['x_start']+crop_coor['x_size']]
  272. control_image = control_image[crop_coor['y_start']:crop_coor['y_start']+crop_coor['y_size'], crop_coor['x_start']:crop_coor['x_start']+crop_coor['x_size']]
  273. gt_seg_image = gt_seg_image[crop_coor['y_start']:crop_coor['y_start']+crop_coor['y_size'], crop_coor['x_start']:crop_coor['x_start']+crop_coor['x_size']]
  275. # Define target size
  276. target_size = (self.size, self.size)
  278. # Create an empty array of zeros with the target size
  279. image_padded = np.zeros((target_size[0], target_size[1], image.shape[2]), dtype=np.float32)
  280. control_image_padded = np.zeros((target_size[0], target_size[1]), dtype=np.float32)
  281. gt_seg_image_padded = np.zeros((target_size[0], target_size[1]), dtype=np.int32)
  282. # Calculate padding offsets
  283. y_offset = (target_size[0] - image.shape[0]) // 2
  284. x_offset = (target_size[1] - image.shape[1]) // 2
  285. # Place the cropped image in the center of the padded image
  286. image_padded[y_offset:y_offset+image.shape[0], x_offset:x_offset+image.shape[1], :] = image
  287. control_image_padded[y_offset:y_offset+control_image.shape[0], x_offset:x_offset+control_image.shape[1]] = control_image
  288. gt_seg_image_padded[y_offset:y_offset+gt_seg_image.shape[0], x_offset:x_offset+gt_seg_image.shape[1]] = gt_seg_image
  290. # Apply min-max normalization to each channel
  291. control_image = (control_image_padded - control_image_padded.min()) / (control_image_padded.max() - control_image_padded.min())
  292. image = image_padded
  293. control_image = control_image[..., np.newaxis]
  294. control_image = self.transform(control_image)
  296. clip_image = self.clip_image_processor(images=image, return_tensors="pt", do_rescale=False).pixel_values
  298. return {
  299. "channel_names": channel_names,
  300. "image": image,
  301. "control_image": control_image,
  302. "clip_image": clip_image,
  303. "gt_seg_image": gt_seg_image_padded
  304. }
  305. def __len__(self):
  306. return len(self.data)
  308. # Dataset
  309. class MyDataset(torch.utils.data.Dataset):
  311. def __init__(self, json_file, size=512, displacement_rate=32, random_y_squeeze_rate=0.05, random_rotation_degree=5, random_shift_range=2, i_drop_rate=0.0, image_root_path="", clip_image_processor_path=None, isTrain=False):
  312. super().__init__()
  314. self.size = size
  315. self.displacement_rate = displacement_rate
  316. self.random_y_squeeze_rate = random_y_squeeze_rate
  317. self.random_rotation_degree = random_rotation_degree
  318. self.random_shift_range = random_shift_range
  319. self.i_drop_rate = i_drop_rate
  320. self.image_root_path = image_root_path
  321. self.isTrain = isTrain
  323. train_min_max_file = os.path.join(self.image_root_path, "train_min_max.json")
  324. with open(train_min_max_file, "r") as f:
  325. train_min_max = json.load(f)
  327. self.b50_min = train_min_max["B50"]["prostate_min"]
  328. self.b50_max = train_min_max["B50"]["prostate_max"]
  329. self.b400_min = train_min_max["B400"]["prostate_min"]
  330. self.b400_max = train_min_max["B400"]["prostate_max"]
  331. self.b800_min = train_min_max["B800"]["prostate_min"]
  332. self.b800_max = train_min_max["B800"]["prostate_max"]
  333. self.t2_min = train_min_max["T2"]["prostate_min"]
  334. self.t2_max = train_min_max["T2"]["prostate_max"]
  336. self.data = []
  337. with open(json_file, "r") as f:
  338. for line in f:
  339. self.data.append(json.loads(line))
  341. self.transform = transforms.Compose([
  342. transforms.ToTensor(),
  343. transforms.Normalize([0.5], [0.5]),
  344. ])
  345. if clip_image_processor_path is None:
  346. self.clip_image_processor = CLIPImageProcessor()
  347. else:
  348. self.clip_image_processor = CLIPImageProcessor.from_pretrained(clip_image_processor_path)
  350. def __getitem__(self, idx):
  351. item = self.data[idx]
  353. image_file = item["image"]
  354. distortion_pivot = item["distortion_pivot"]
  355. crop_coor = item["crop_coor"]
  357. # read image
  358. raw_b50_image = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(self.image_root_path, 'B50', image_file))).astype(np.float32)
  359. raw_b400_image = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(self.image_root_path, 'B400', image_file))).astype(np.float32)
  360. raw_b800_image = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(self.image_root_path, 'B800', image_file))).astype(np.float32)
  361. control_raw_image = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(self.image_root_path, 'T2', image_file))).astype(np.float32)
  362. gt_seg_image = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(self.image_root_path, 'PZTZ', image_file)))
  364. # global min-max normalization
  365. raw_b50_image = (raw_b50_image - self.b50_min) / (self.b50_max - self.b50_min)
  366. raw_b400_image = (raw_b400_image - self.b400_min) / (self.b400_max - self.b400_min)
  367. raw_b800_image = (raw_b800_image - self.b800_min) / (self.b800_max - self.b800_min)
  369. # clip 0.0 ~ 1.0
  370. raw_b50_image = np.clip(raw_b50_image, 0, 1)
  371. raw_b400_image = np.clip(raw_b400_image, 0, 1)
  372. raw_b800_image = np.clip(raw_b800_image, 0, 1)
  374. image = np.stack([raw_b50_image, raw_b400_image, raw_b800_image], axis=-1)
  375. image = image[crop_coor['y_start']:crop_coor['y_start']+crop_coor['y_size'], crop_coor['x_start']:crop_coor['x_start']+crop_coor['x_size']]
  376. control_image = control_raw_image[crop_coor['y_start']:crop_coor['y_start']+crop_coor['y_size'], crop_coor['x_start']:crop_coor['x_start']+crop_coor['x_size']]
  377. gt_seg_image = gt_seg_image[crop_coor['y_start']:crop_coor['y_start']+crop_coor['y_size'], crop_coor['x_start']:crop_coor['x_start']+crop_coor['x_size']]
  379. # Define target size
  380. target_size = (self.size, self.size)
  382. # Create an empty array of zeros with the target size
  383. image_padded = np.zeros((target_size[0], target_size[1], image.shape[2]), dtype=np.float32)
  384. control_image_padded = np.zeros((target_size[0], target_size[1]), dtype=np.float32)
  385. gt_seg_image_padded = np.zeros((target_size[0], target_size[1]), dtype=np.int32)
  386. # Calculate padding offsets
  387. y_offset = (target_size[0] - image.shape[0]) // 2
  388. x_offset = (target_size[1] - image.shape[1]) // 2
  389. # Place the cropped image in the center of the padded image
  390. image_padded[y_offset:y_offset+image.shape[0], x_offset:x_offset+image.shape[1], :] = image
  391. control_image_padded[y_offset:y_offset+control_image.shape[0], x_offset:x_offset+control_image.shape[1]] = control_image
  392. gt_seg_image_padded[y_offset:y_offset+gt_seg_image.shape[0], x_offset:x_offset+gt_seg_image.shape[1]] = gt_seg_image
  394. control_image = (control_image_padded - control_image_padded.min()) / (control_image_padded.max() - control_image_padded.min())
  396. image = image_padded
  398. control_image = control_image[..., np.newaxis]
  399. image = self.transform(image)
  400. control_image = self.transform(control_image)
  401. gt_seg_image = torch.from_numpy(gt_seg_image_padded)
  403. if np.random.rand() > 0.5:
  404. displacement = self.get_random_displacement(distortion_pivot, raw_b50_image)
  405. b50_deformed = elasticdeform.deform_grid(raw_b50_image, displacement, order=3)
  406. b400_deformed = elasticdeform.deform_grid(raw_b400_image, displacement, order=3)
  407. b800_deformed = elasticdeform.deform_grid(raw_b800_image, displacement, order=3)
  408. upscaled_displacement = self.upscale_displacement(displacement, raw_b50_image.shape)
  409. jacobian_det = self.compute_jacobian(upscaled_displacement)
  411. b50_adjusted = self.adjust_pixel_values(raw_b50_image, b50_deformed, jacobian_det)
  412. b400_adjusted = self.adjust_pixel_values(raw_b400_image, b400_deformed, jacobian_det)
  413. b800_adjusted = self.adjust_pixel_values(raw_b800_image, b800_deformed, jacobian_det)
  415. image_deformed = np.stack([b50_adjusted, b400_adjusted, b800_adjusted], axis=-1)
  416. image_deformed = np.clip(image_deformed, 0, 1)
  418. image_deformed = self.random_squeeze_y_axis(image_deformed)
  419. image_deformed = self.random_rotation(image_deformed)
  420. x_shift = random.randint(-self.random_shift_range, self.random_shift_range)
  421. y_shift = random.randint(-self.random_shift_range, self.random_shift_range)
  422. image_deformed = np.roll(image_deformed, (x_shift, y_shift), axis=(1, 0))
  424. else:
  425. image_deformed = np.stack([raw_b50_image, raw_b400_image, raw_b800_image], axis=-1)
  427. image_deformed = image_deformed[crop_coor['y_start']:crop_coor['y_start']+crop_coor['y_size'], crop_coor['x_start']:crop_coor['x_start']+crop_coor['x_size']]
  428. image_deformed_padded = np.zeros((target_size[0], target_size[1], image_deformed.shape[2]), dtype=image_deformed.dtype)
  429. image_deformed_padded[y_offset:y_offset+image_deformed.shape[0], x_offset:x_offset+image_deformed.shape[1], :] = image_deformed
  431. image_deformed = image_deformed_padded
  433. clip_image = self.clip_image_processor(images=image_deformed, return_tensors="pt", do_rescale=False).pixel_values
  435. return {
  436. "image": image,
  437. "clip_image": clip_image,
  438. "control_image": control_image,
  439. "gt_seg_image": gt_seg_image,
  440. }
  442. def __len__(self):
  443. return len(self.data)
  445. def downsample_point(self, point, arr):
  446. original_size = np.array(arr.shape)
  447. target_size = original_size // self.displacement_rate
  448. return [int(round(p * t / o)) for p, o, t in zip(point, original_size, target_size)]
  450. def get_random_displacement(self, points, arr):
  451. points_downsampled = []
  452. for point in points:
  453. points_downsampled.append(self.downsample_point(point, arr))
  454. points_downsampled = np.array(points_downsampled)
  456. center_point = points_downsampled[4]
  458. min_x = points_downsampled[0][1]
  459. max_x = points_downsampled[2][1]
  460. min_y = points_downsampled[0][0]
  461. max_y = points_downsampled[6][0]
  463. displacement_size_y = arr.shape[0] // self.displacement_rate
  464. displacement_size_x = arr.shape[1] // self.displacement_rate
  465. displacement = np.zeros((2, displacement_size_y, displacement_size_x))
  466. for i in range(displacement_size_y):
  467. for j in range(displacement_size_x):
  468. if (i,j) == (center_point[0], center_point[1]):
  469. displacement[0, i, j] = np.random.randn() * 10 # randn: expansion+compression, rand: expansion only
  470. displacement[1, i, j] = np.random.randn() * 10
  471. elif (min_y < i < max_y) and (min_x < j < max_x):
  472. displacement[0, i, j] = np.random.randn() * 10
  473. displacement[1, i, j] = np.random.randn() * 10
  475. return displacement
  477. def upscale_displacement(self, disp, target_shape, order=3):
  478. factor = target_shape[0] / disp.shape[1] # assumes square image and displacement grid
  479. upscaled_disp = np.zeros((2, *target_shape))
  480. upscaled_disp[0] = zoom(disp[0], factor, order=order)
  481. upscaled_disp[1] = zoom(disp[1], factor, order=order)
  482. return upscaled_disp
  484. def compute_jacobian(self, disp):
  485. # Compute spatial gradients of displacement field
  486. grad_y_x = np.gradient(disp[0], axis=1) # ∂(y-disp)/∂x
  487. grad_y_y = np.gradient(disp[0], axis=0) # ∂(y-disp)/∂y
  488. grad_x_x = np.gradient(disp[1], axis=1) # ∂(x-disp)/∂x
  489. grad_x_y = np.gradient(disp[1], axis=0) # ∂(x-disp)/∂y
  491. # Jacobian determinant: J = (1 + ∂y/∂y)(1 + ∂x/∂x) - (∂y/∂x)(∂x/∂y)
  492. jacobian = (1 + grad_y_y) * (1 + grad_x_x) - (grad_y_x * grad_x_y)
  493. return jacobian
  495. def adjust_pixel_values(self, original_image, deformed_image, jacobian_det):
  496. return deformed_image * jacobian_det
  498. def random_squeeze_y_axis(self, image):
  499. """
  500. Randomly scales the image along the Y-axis while keeping the X and Z axes unchanged.
  502. Parameters:
  503. image (numpy.ndarray): 3D image array with shape (height, width, channels).
  504. scale_range (tuple): Tuple with min and max scale factors for the Y-axis.
  506. Returns:
  507. numpy.ndarray: Scaled image.
  508. """
  509. # Ensure image is a 3D numpy array
  510. if image.ndim != 3:
  511. raise ValueError("Input image must be a 3D numpy array.")
  513. # Randomly choose a scale factor for the Y-axis
  514. scale_y = np.random.uniform(1-self.random_y_squeeze_rate, 1+self.random_y_squeeze_rate)
  516. # Define scale factors for X, Y, Z axes
  517. scale_factors = [scale_y, 1, 1]
  519. # Rescale the image
  520. image_rescaled = transform.rescale(image, scale_factors, mode='reflect', anti_aliasing=True)
  522. return image_rescaled
  524. def random_rotation(self, image):
  525. # Randomly choose an angle for rotation
  526. max_angle = self.random_rotation_degree
  527. angle = np.random.uniform(-max_angle, max_angle)
  529. # Rotate the image
  530. image_rotated = transform.rotate(image, angle, mode='reflect')
  532. return image_rotated
  535. def collate_fn(data):
  536. images = torch.stack([example["image"] for example in data])
  537. clip_images = torch.cat([example["clip_image"] for example in data], dim=0)
  538. control_images = torch.stack([example["control_image"] for example in data])
  539. gt_seg_images = torch.stack([example["gt_seg_image"] for example in data])
  541. return {
  542. "images": images,
  543. "clip_images": clip_images,
  544. "control_images": control_images,
  545. "gt_seg_images": gt_seg_images
  546. }
  548. def parse_args():
  549. parser = argparse.ArgumentParser(description="Simple example of a training script.")
  550. parser.add_argument(
  551. "--pretrained_model_name_or_path",
  552. type=str,
  553. default=None,
  554. required=True,
  555. help="Path to pretrained model or model identifier from huggingface.co/models.",
  556. )
  557. parser.add_argument(
  558. "--data_json_file",
  559. type=str,
  560. default=None,
  561. required=True,
  562. help="Training data",
  563. )
  564. parser.add_argument(
  565. "--val_data_json_file",
  566. type=str,
  567. default=None,
  568. required=True,
  569. help="Validation data",
  570. )
  571. parser.add_argument(
  572. "--plot_data_json_file",
  573. type=str,
  574. default=None,
  575. required=True,
  576. help="Plot data",
  577. )
  578. parser.add_argument(
  579. "--data_root_path",
  580. type=str,
  581. default="",
  582. required=True,
  583. help="Training data root path",
  584. )
  585. parser.add_argument(
  586. "--image_encoder_path",
  587. type=str,
  588. default=None,
  589. required=True,
  590. help="Path to CLIP image encoder",
  591. )
  592. parser.add_argument(
  593. "--output_dir",
  594. type=str,
  595. default="sd-ipadapter-dec+controlnet",
  596. help="The output directory where the model predictions and checkpoints will be written.",
  597. )
  598. parser.add_argument(
  599. "--logging_dir",
  600. type=str,
  601. default="logs",
  602. help=(
  603. "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
  604. " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
  605. ),
  606. )
  607. parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
  608. parser.add_argument(
  609. "--resolution",
  610. type=int,
  611. default=512,
  612. help=(
  613. "The resolution for input images"
  614. ),
  615. )
  616. parser.add_argument(
  617. "--displacement_rate",
  618. type=int,
  619. default=32,
  620. help=(
  621. "The rate for displacement"
  622. ),
  623. )
  624. parser.add_argument(
  625. "--random_y_squeeze_rate",
  626. type=float,
  627. default=0.05,
  628. help=(
  629. "The rate for random y squeeze"
  630. ),
  631. )
  632. parser.add_argument(
  633. "--random_rotation_degree",
  634. type=int,
  635. default=5,
  636. help=(
  637. "The degree for random rotation"
  638. ),
  639. )
  640. parser.add_argument(
  641. "--random_shift_range",
  642. type=int,
  643. default=2,
  644. help=(
  645. "The range for random shift"
  646. ),
  647. )
  648. parser.add_argument(
  649. "--image_condition_drop_rate",
  650. type=float,
  651. default=0.1,
  652. help=(
  653. "The rate for image condition drop"
  654. ),
  655. )
  656. parser.add_argument(
  657. "--learning_rate",
  658. type=float,
  659. default=1e-4,
  660. help="Learning rate to use.",
  661. )
  662. parser.add_argument("--weight_decay", type=float, default=1e-2, help="Weight decay to use.")
  663. parser.add_argument("--num_train_epochs", type=int, default=100)
  664. parser.add_argument(
  665. "--max_train_steps",
  666. type=int,
  667. default=None,
  668. help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
  669. )
  670. parser.add_argument(
  671. "--gradient_accumulation_steps",
  672. type=int,
  673. default=1,
  674. help="Number of updates steps to accumulate before performing a backward/update pass.",
  675. )
  676. parser.add_argument(
  677. "--scale_lr",
  678. action="store_true",
  679. default=False,
  680. help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
  681. )
  682. parser.add_argument(
  683. "--train_batch_size", type=int, default=8, help="Batch size (per device) for the training dataloader."
  684. )
  685. parser.add_argument(
  686. "--dataloader_num_workers",
  687. type=int,
  688. default=0,
  689. help=(
  690. "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
  691. ),
  692. )
  693. parser.add_argument(
  694. "--checkpointing_steps",
  695. type=int,
  696. default=500,
  697. help=(
  698. "Save a checkpoint of the training state every X updates. Checkpoints can be used for resuming training via `--resume_from_checkpoint`. "
  699. "In the case that the checkpoint is better than the final trained model, the checkpoint can also be used for inference."
  700. "Using a checkpoint for inference requires separate loading of the original pipeline and the individual checkpointed model components."
  701. "See https://huggingface.co/docs/diffusers/main/en/training/dreambooth#performing-inference-using-a-saved-checkpoint for step by step"
  702. "instructions."
  703. ),
  704. )
  705. parser.add_argument(
  706. "--resume_from_checkpoint",
  707. type=str,
  708. default=None,
  709. help=(
  710. "Whether training should be resumed from a previous checkpoint. Use a path saved by"
  711. ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
  712. ),
  713. )
  714. parser.add_argument(
  715. "--mixed_precision",
  716. type=str,
  717. default=None,
  718. choices=["no", "fp16", "bf16"],
  719. help=(
  720. "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
  721. " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the"
  722. " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
  723. ),
  724. )
  725. parser.add_argument(
  726. "--report_to",
  727. type=str,
  728. default="tensorboard",
  729. help=(
  730. 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
  731. ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
  732. ),
  733. )
  734. parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
  736. args = parser.parse_args()
  737. env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
  738. if env_local_rank != -1 and env_local_rank != args.local_rank:
  739. args.local_rank = env_local_rank
  741. return args
  744. def main():
  745. args = parse_args()
  746. logging_dir = Path(args.output_dir, args.logging_dir)
  748. accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
  750. accelerator = Accelerator(
  751. gradient_accumulation_steps=args.gradient_accumulation_steps,
  752. mixed_precision=args.mixed_precision,
  753. log_with=args.report_to,
  754. project_config=accelerator_project_config,
  755. )
  757. # Make one log on every process with the configuration for debugging.
  758. logging.basicConfig(
  759. format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
  760. datefmt="%m/%d/%Y %H:%M:%S",
  761. level=logging.INFO,
  762. )
  763. logger.info(accelerator.state, main_process_only=False)
  765. wandb.init(project="DWIDistortion")
  767. if args.seed is not None:
  768. set_seed(args.seed)
  770. if accelerator.is_main_process:
  771. if args.output_dir is not None:
  772. os.makedirs(args.output_dir, exist_ok=True)
  774. # Load scheduler and models.
  775. noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
  776. vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae")
  777. unet = UNet2DConditionModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="unet", cross_attention_dim=1280, ignore_mismatched_sizes=True, low_cpu_mem_usage=False)
  778. image_encoder = CLIPVisionModelWithProjection.from_pretrained(args.image_encoder_path)
  779. # freeze parameters of models to save more memory
  780. unet.requires_grad_(False)
  781. unet.conv_out.requires_grad_(True)
  783. pretrained_state_dict = UNet2DConditionModel.from_pretrained(
  784. args.pretrained_model_name_or_path, subfolder="unet"
  785. ).state_dict()
  787. # Enable gradients for randomly initialized layers
  788. # (layers whose shapes changed due to ignore_mismatched_sizes=True)
  789. random_initialized_layers = []
  790. for name, param in unet.named_parameters():
  791. if param.shape != pretrained_state_dict[name].shape:
  792. param.requires_grad = True
  793. random_initialized_layers.append(name)
  794. print("Random initialized layers: ", random_initialized_layers)
  795. # Free up memory by deleting pretrained_state_dict
  796. del pretrained_state_dict
  798. vae.requires_grad_(False)
  799. image_encoder.requires_grad_(True)
  801. controlnet = ControlNetModel.from_unet(unet, conditioning_channels=1)
  802. controlnet.train()
  803. unet.train()
  804. vae.eval()
  805. image_encoder.train()
  807. # Taken from [Sayak Paul's Diffusers PR #6511](https://github.com/huggingface/diffusers/pull/6511/files)
  808. def unwrap_model(model):
  809. model = accelerator.unwrap_model(model)
  810. model = model._orig_mod if is_compiled_module(model) else model
  811. return model
  813. # `accelerate` 0.16.0 will have better support for customized saving
  814. if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
  815. # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
  816. def save_model_hook(models, weights, output_dir):
  817. if accelerator.is_main_process:
  819. models[0].save_pretrained(os.path.join(output_dir, "unet"))
  820. weights.pop()
  822. models[1].save_pretrained(os.path.join(output_dir, "controlnet"))
  823. weights.pop()
  825. models[2].save_pretrained(os.path.join(output_dir, "image_encoder"))
  826. weights.pop()
  828. def load_model_hook(models, input_dir):
  830. model = models.pop()
  831. load_image_encoder = CLIPVisionModelWithProjection.from_pretrained(os.path.join(input_dir, "image_encoder"))
  832. model.load_state_dict(load_image_encoder.state_dict())
  833. del load_image_encoder
  835. model = models.pop()
  836. load_controlnet = ControlNetModel.from_pretrained(input_dir, subfolder="controlnet")
  837. model.register_to_config(**load_controlnet.config)
  838. model.load_state_dict(load_controlnet.state_dict())
  839. del load_controlnet
  841. model = models.pop()
  842. load_unet = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
  843. model.load_state_dict(load_unet.state_dict())
  844. del load_unet
  846. accelerator.register_save_state_pre_hook(save_model_hook)
  847. accelerator.register_load_state_pre_hook(load_model_hook)
  849. if args.scale_lr:
  850. args.learning_rate = (
  851. args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
  852. )
  854. weight_dtype = torch.float32
  855. if accelerator.mixed_precision == "fp16":
  856. weight_dtype = torch.float16
  857. elif accelerator.mixed_precision == "bf16":
  858. weight_dtype = torch.bfloat16
  859. vae.to(accelerator.device, dtype=weight_dtype)
  860. params_to_opt = itertools.chain(
  861. image_encoder.parameters(),
  862. unet.conv_out.parameters(),
  863. (param for name, param in unet.named_parameters() if name in random_initialized_layers), # randomly initialized layers
  864. controlnet.parameters()
  865. )
  866. optimizer = torch.optim.AdamW(params_to_opt, lr=args.learning_rate, weight_decay=args.weight_decay)
  868. # dataloader
  869. train_dataset = MyDataset(args.data_json_file, size=args.resolution, displacement_rate=args.displacement_rate, random_y_squeeze_rate=args.random_y_squeeze_rate, random_rotation_degree=args.random_rotation_degree, random_shift_range=args.random_shift_range, image_root_path=args.data_root_path, clip_image_processor_path=args.image_encoder_path, isTrain=True)
  870. train_dataloader = torch.utils.data.DataLoader(
  871. train_dataset,
  872. shuffle=True,
  873. collate_fn=collate_fn,
  874. batch_size=args.train_batch_size,
  875. num_workers=args.dataloader_num_workers,
  876. pin_memory=True,
  877. drop_last=True,
  878. )
  880. # Prepare everything with our `accelerator`.
  881. unet, controlnet, image_encoder, optimizer, train_dataloader = accelerator.prepare(unet, controlnet, image_encoder, optimizer, train_dataloader)
  883. # Scheduler and math around the number of training steps.
  884. overrode_max_train_steps = False
  885. num_update_steps_per_epoch = math.ceil(
  886. len(train_dataloader) / args.gradient_accumulation_steps
  887. )
  888. if args.max_train_steps is None:
  889. args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
  890. overrode_max_train_steps = True
  891. # Train!
  892. total_batch_size = (
  893. args.train_batch_size
  894. * accelerator.num_processes
  895. * args.gradient_accumulation_steps
  896. )
  897. logger.info("***** Running training *****")
  898. logger.info(f" Num examples = {len(train_dataset)}")
  899. logger.info(f" Num Epochs = {args.num_train_epochs}")
  900. logger.info(f" Instantaneous batch size per device = {args.train_batch_size}")
  901. logger.info(
  902. f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}"
  903. )
  904. logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}")
  905. logger.info(f" Total optimization steps = {args.max_train_steps}")
  906. global_step = 0
  907. first_epoch = 0
  909. initial_global_step = 0
  911. train_losses = []
  912. train_noise_losses = []
  913. val_psnrs = []
  914. val_ssims = []
  916. progress_bar = tqdm(
  917. range(0, args.max_train_steps),
  918. initial=initial_global_step,
  919. desc="Steps",
  920. # Only show the progress bar once on each machine.
  921. disable=not accelerator.is_local_main_process,
  922. )
  925. for epoch in range(first_epoch, args.num_train_epochs):
  926. begin = time.perf_counter()
  928. epoch_train_loss = 0.0
  929. epoch_train_noise_loss = 0.0
  931. for step, batch in enumerate(train_dataloader):
  933. if accelerator.sync_gradients:
  934. if accelerator.is_main_process:
  935. if global_step % args.checkpointing_steps == 0:
  936. save_path = os.path.join(args.output_dir, f"checkpoint-last")
  937. accelerator.save_state(save_path)
  938. logger.info(f"Saved state to {save_path}")
  940. if args.val_data_json_file is not None and args.plot_data_json_file is not None:
  941. val_psnr, val_ssim, val_psnr_cropped, val_ssim_cropped = log_validation(args, accelerator, weight_dtype, global_step, save_path)
  943. if global_step!=0 and (val_psnr_cropped+val_ssim_cropped) > (max(val_psnrs)+max(val_ssims)):
  944. save_path = os.path.join(args.output_dir, f"checkpoint-best-{global_step}")
  945. accelerator.save_state(save_path)
  946. logger.info(f"Saved state to {save_path}")
  948. val_psnrs.append(float(val_psnr_cropped))
  949. val_ssims.append(float(val_ssim_cropped))
  951. load_data_time = time.perf_counter() - begin
  952. with accelerator.accumulate(unet, controlnet, image_encoder):
  953. # Convert images to latent space
  954. with torch.no_grad():
  955. latents = vae.encode(batch["images"].to(accelerator.device, dtype=weight_dtype)).latent_dist.sample()
  956. latents = latents * vae.config.scaling_factor
  958. # Sample noise that we'll add to the latents
  959. noise = torch.randn_like(latents)
  960. bsz = latents.shape[0]
  961. # Sample a random timestep for each image
  962. timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
  963. timesteps = timesteps.long().to(accelerator.device)
  965. # Add noise to the latents according to the noise magnitude at each timestep
  966. # (this is the forward diffusion process)
  967. noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps).to(accelerator.device, dtype=weight_dtype)
  969. encoder_hidden_states = image_encoder(batch["clip_images"].to(accelerator.device, dtype=weight_dtype)).last_hidden_state
  971. controlnet_image = batch["control_images"].to(accelerator.device, dtype=weight_dtype)
  973. down_block_res_samples, mid_block_res_sample = controlnet(
  974. noisy_latents,
  975. timesteps,
  976. encoder_hidden_states,
  977. controlnet_cond=controlnet_image,
  978. return_dict=False,
  979. )
  981. noise_pred = unet(
  982. noisy_latents,
  983. timesteps,
  984. encoder_hidden_states,
  985. down_block_additional_residuals=[
  986. sample.to(dtype=noisy_latents.dtype) for sample in down_block_res_samples
  987. ],
  988. mid_block_additional_residual=mid_block_res_sample.to(dtype=noisy_latents.dtype),
  989. return_dict=False,
  990. )[0]
  992. noise_loss = F.mse_loss(noise_pred.float(), noise.float(), reduction="mean")
  994. timesteps[timesteps == 0] = 1
  996. loss = noise_loss
  998. # Gather the losses across all processes for logging (if we use distributed training).
  999. avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean().item()
  1000. avg_noise_loss = accelerator.gather(noise_loss.repeat(args.train_batch_size)).mean().item()
  1002. # accumulate the average loss for each batch
  1003. epoch_train_loss += avg_loss
  1004. epoch_train_noise_loss += avg_noise_loss
  1006. # Backward pass and optimization for main model
  1007. optimizer.zero_grad()
  1008. accelerator.backward(loss, retain_graph=True)
  1009. optimizer.step()
  1011. # Checks if the accelerator has performed an optimization step behind the scenes
  1012. if accelerator.sync_gradients:
  1013. progress_bar.update(1)
  1014. progress_bar.set_postfix_str(f"Loss: {avg_loss:.3f}")
  1015. global_step += 1
  1017. if global_step >= args.max_train_steps:
  1018. break
  1020. begin = time.perf_counter()
  1021. epoch_train_loss /= len(train_dataloader)
  1022. epoch_train_noise_loss /= len(train_dataloader)
  1024. wandb.log({
  1025. "train_loss": epoch_train_loss,
  1026. "train_noise_loss": epoch_train_noise_loss,
  1027. "step": global_step,
  1028. "learning_rate": optimizer.param_groups[0]["lr"],
  1029. })
  1031. train_losses.append(epoch_train_loss)
  1032. train_noise_losses.append(epoch_train_noise_loss)
  1034. losses_data = {
  1035. "train_losses": train_losses,
  1036. "train_noise_losses": train_noise_losses,
  1037. "val_psnrs": val_psnrs,
  1038. "val_ssims": val_ssims,
  1039. }
  1040. with open(os.path.join(args.output_dir, "losses.json"), "w") as f:
  1041. json.dump(losses_data, f)
  1043. # Create the pipeline using using the trained modules and save it.
  1044. accelerator.wait_for_everyone()
  1045. if accelerator.is_main_process:
  1046. save_path = os.path.join(args.output_dir, "final")
  1047. if not os.path.isdir(save_path):
  1048. os.mkdir(save_path)
  1049. unet = unwrap_model(unet)
  1050. unet.save_pretrained(os.path.join(save_path, "unet"))
  1051. controlnet = unwrap_model(controlnet)
  1052. controlnet.save_pretrained(os.path.join(save_path, "controlnet"))
  1053. image_encoder = unwrap_model(image_encoder)
  1054. image_encoder.save_pretrained(os.path.join(save_path, "image_encoder"))
  1055. logger.info(f"Saved final model to {save_path}")
  1057. accelerator.end_training()
  1059. if __name__ == "__main__":
  1060. main()

train_DeDistortNet.py at commit 5534062, under Apache-2.0 · at the source

Overview

Authors: Inye Na1, Qi Miao2, Jonghun Kim1, Kyunghyun Sung2, Hyunjin Park1
ORCID iDs: Hyunjin Park
  1. Department of Electrical and Computer Engineering, Sungkyunkwan University,Suwon, Republic of Korea
  2. Department of Radiological Sciences, David Geffen School of Medicine, University of California, Los Angeles,Los Angeles, CA USA
Institutions: Sungkyunkwan University (South Korea); University of California, Los Angeles (United States)
Journal: European radiology experimental, volume 10, issue 1, article 66
Dates: received 5 November 2025; accepted 17 April 2026; published online 13 May 2026
Type: Research article · Language: English
License: CC BY
Identifiers: DOI 10.1186/s41747-026-00735-w · PMID 42126718 · PMCID PMC13172242 · OpenAlex W7161056552
Open access: gold, a free copy (OpenAlex)
Status: code verified
Categories: structural MRI / diffusion (modality), human (organism)
Methods: Connectivity
Keywords: Artifacts, Artificial intelligence, Deep learning, Diffusion magnetic resonance imaging, Prostate
MeSH: Diffusion Magnetic Resonance Imaging*, Image Interpretation, Computer-Assisted*, Image Processing, Computer-Assisted*, Prostate*, Prostatic Neoplasms*, Artifacts, Humans, Male (* major topic)
Topic: Prostate Cancer Diagnosis and Treatment (Pulmonary and Respiratory Medicine, Medicine), according to OpenAlex
Funding: National Research Foundation of Korea (RS-2024-00408040)
Citations: not cited yet (Europe PMC); 39 references in the paper

Abstract

Objective: Geometric distortion from susceptibility artifacts in diffusion-weighted imaging (DWI) degrades anatomical fidelity and complicates clinical prostate magnetic resonance imaging (MRI) interpretation. We developed and evaluated DeDistortNet, a generative dual-conditioned diffusion model for correcting geometric distortions in prostate DWI without experimentally acquired paired distorted-undistorted data.

Materials and methods: We utilized the public PROSTATEx dataset, divided into training (n = 135, 1,893 slices), validation (n = 4, 64 slices), and test (n = 189, 2,623 slices) cohorts stratified by distortion severity. Model training used only undistorted DWIs, with simulated distortions to generate paired examples. DeDistortNet combines contextual guidance from distorted DWIs with structural guidance from T2-weighted images to synthesize distortion-free DWIs. Performance was evaluated using quantitative analysis on simulated distortions and indirect validation on clinically distorted DWIs through anatomical concordance with T2-derived prostate masks.

Results: The cohort included prostate MRI exams from 328 male subjects. The DWIs from 139 exams were undistorted for training/validation, while 189 were distorted for testing. In simulated data, DeDistortNet restored image quality across distortion severities, improving peak signal-to-noise ratio by 36% and structural similarity by 55% in the peripheral zone (PZ) under extreme distortion. In clinically distorted data, concordance with T2-weighted references (PZ Dice similarity) improved by 37% under severe distortion and by 72% under extreme distortion. Radiologist assessment further supported improved geometric fidelity and prostate boundary delineation after correction.

Conclusion: DeDistortNet, trained on undistorted DWIs with simulated distortions, effectively corrected distortions in prostate DWI and restored anatomical fidelity, particularly in the PZ, without requiring additional acquisitions or specialized imaging protocols.

Relevance statement: By correcting severe distortions without additional imaging sequences, DeDistortNet restores anatomical fidelity in prostate diffusion-weighted imaging, particularly in the peripheral zone, enabling more reliable image interpretation and reducing the need for repeat scans.

Key Points: DeDistortNet corrects geometric distortion in prostate diffusion MRI. Our method integrates anatomical guidance from T2-weighted images for distortion-free reconstruction. Our method was trained on simulated distortions without experimentally acquired paired distorted-undistorted data or additional MRI sequences.

Reproduced under the paper's license (CC BY), from the paper cited above.

Repository

Its files are read in the Code ↔ Paper reader above, with 5 matches between paragraphs and lines of code.

nainye/DeDistortNet

License: Apache-2.0
State: the link answers, verified on 28 September 2026
Evidence: files inventoried
Commit: 553406232f87111a737a1582f96b1978f10f64e4, 10 June 2026
Languages: Jupyter (3), Shell (1), Python (1)
Size: 16 files, 5 scripts
Software Heritage: not archived
Found in: “Code availability”
Holds: README, license file, environment (scripts/requirements.txt), 3 notebooks
Not found: CITATION.cff, tests, continuous integration, documentation
Tools: NumPy (4 files), SimpleITK (4 files), Matplotlib (3 files), SciPy (2 files), pandas (1 file), pydicom (1 file), PyTorch (1 file), scikit-image (1 file), Hugging Face Transformers (1 file)
Availability: 1 check, the latest on 28 September 2026: the link answers
  • 28 September 2026: the link answers
7 files
At the source: github.com/nainye/DeDistortNet

Code availability

The derived slice-level distortion severity labels and model implementation codes developed for this study are publicly available in a dedicated GitHub repository at: https://github.com/nainye/DeDistortNet.

Reproduced under the paper's license (CC BY), from the paper cited above.

Tracing map

Proposed by the machine: these links were found in the paper and verified at the source, without human review. The map will receive a Zenodo DOI once one of the paper's authors has validated it with their ORCID.

What the map holds:

  • 1 repository of the authors' code, each at its verified commit, with its license and how the link was found in the paper;
  • 5 scripts, each with its path and the digest of its content;
  • 5 matches between paragraphs of the paper and lines of the code (method lexical-v1);
  • neither the text of the paper nor the code itself.

Its JSON (tracing-map.json) is deposited on Zenodo with its DOI once the map is validated.

Data

No dataset and no data link were found in the paper.

Data availability

The dataset analyzed during this study is publicly available in The Cancer Imaging Archive (TCIA) at https://www.cancerimagingarchive.net/collection/prostatex/.

Reproduced under the paper's license (CC BY), from the paper cited above.

Versions

The history of this record: each version stored by the harvester or made by a correction of its authors or of the maintainers of its code, and what changed in its facts. The texts of the paper (its abstract, its availability statements) are not part of it; versions that changed only those are not listed.

Version 1, 28 September 2026: the first record

Recorded: type, language, journal, volume, issue, pages, dates, 5 authors, 5 keywords, 8 MeSH terms, 1 funder, 31 references.

Cite

This paper

Na, I., Miao, Q., Kim, J., Sung, K., & Park, H. (2026). Dual-conditioned diffusion model with anatomical guidance for geometric distortion correction in prostate MRI. European radiology experimental, 10(1), 66. https://doi.org/10.1186/s41747-026-00735-w

BibTeX

@article{na2026dual,
author = {Na, Inye and Miao, Qi and Kim, Jonghun and Sung, Kyunghyun and Park, Hyunjin},
title = {{Dual-conditioned diffusion model with anatomical guidance for geometric distortion correction in prostate MRI}},
journal = {European radiology experimental},
year = {2026},
month = may,
volume = {10},
number = {1},
pages = {66},
publisher = {Springer},
issn = {2509-9280},
doi = {10.1186/s41747-026-00735-w},
url = {https://doi.org/10.1186/s41747-026-00735-w},
pmid = {42126718},
pmcid = {PMC13172242}
}

RIS

TY - JOUR
AU - Na, Inye
AU - Miao, Qi
AU - Kim, Jonghun
AU - Sung, Kyunghyun
AU - Park, Hyunjin
TI - Dual-conditioned diffusion model with anatomical guidance for geometric distortion correction in prostate MRI
T2 - European radiology experimental
J2 - Eur Radiol Exp
PY - 2026
DA - 2026/05/13
VL - 10
IS - 1
SP - 66
SN - 2509-9280
PB - Springer
DO - 10.1186/s41747-026-00735-w
UR - https://doi.org/10.1186/s41747-026-00735-w
LA - en
ER -

CSL-JSON

{
"id": "10.1186/s41747-026-00735-w",
"type": "article-journal",
"title": "Dual-conditioned diffusion model with anatomical guidance for geometric distortion correction in prostate MRI",
"container-title": "European radiology experimental",
"author": [
{
"family": "Na",
"given": "Inye"
},
{
"family": "Miao",
"given": "Qi"
},
{
"family": "Kim",
"given": "Jonghun"
},
{
"family": "Sung",
"given": "Kyunghyun"
},
{
"family": "Park",
"given": "Hyunjin"
}
],
"container-title-short": "Eur Radiol Exp",
"volume": "10",
"issue": "1",
"page": "66",
"DOI": "10.1186/s41747-026-00735-w",
"PMID": "42126718",
"PMCID": "PMC13172242",
"ISSN": "2509-9280",
"publisher": "Springer",
"URL": "https://doi.org/10.1186/s41747-026-00735-w",
"language": "en",
"issued": {
"date-parts": [
[
2026,
5,
13
]
]
}
}

The tracing map gets a citation of its own once an author has validated it and it has a DOI.

Similar papers

The papers with a page that share the most with this one: the tools found in their code, their categories, datasets, cited references and authors, the rarest counting most.

[1] doi:10.1111/joa.70203 [code]
Two-step workflow integrating automatic registration and manual refinement for the accurate alignment of serial histological sections in 3D reconstruction.
Journal: Journal of anatomy
In common: pydicom, Hugging Face Transformers, SimpleITK, 6 other tools
[2] doi:10.1162/imag.a.1326 [code]
RAVEN: Robust, generalizable, multi-resolution structural MRI upsampling using autoencoders.
Journal: Imaging neuroscience (Cambridge, Mass.)
In common: Hugging Face Transformers, SimpleITK, scikit-image, 5 other tools, structural MRI / diffusion, 1 reference
[3] doi:10.1371/journal.pcbi.1014555 [code]
Body surface potential driven personalisation of electrophysiological digital twins in hypertrophic cardiomyopathy.
Journal: PLoS computational biology
In common: pydicom, SimpleITK, scikit-image, 5 other tools, structural MRI / diffusion
[4] doi:10.1080/07853890.2026.2685416 [code]
Pulmonary and cerebral damage in COVID-19 survivors: is there any association?
Journal: Annals of medicine
In common: pydicom, SimpleITK, scikit-image, 5 other tools, structural MRI / diffusion
[5] doi:10.3389/frai.2026.1771088 [code]
Few-shot deployment of pretrained MRI transformers in brain imaging tasks.
Journal: Frontiers in artificial intelligence
In common: pydicom, SimpleITK, scikit-image, 5 other tools, structural MRI / diffusion
[6] doi:10.21037/qims-2026-0792 [code]
An nnU-Net-based framework with adaptive feature representation for 3D brain tumor segmentation.
Journal: Quantitative imaging in medicine and surgery
In common: SimpleITK, scikit-image, PyTorch, 4 other tools, structural MRI / diffusion, 2 references
[7] doi:10.1038/s42003-026-10957-8 [code]
Brain defence by the extracellular matrix protein Cochlin.
Journal: Communications biology
In common: pydicom, Hugging Face Transformers, scikit-image, 5 other tools
[8] doi:10.1038/s41598-026-48496-1 [code]
A unified FLAIR hyperintensity segmentation model for various CNS tumor types and acquisition time points.
Journal: Scientific reports
In common: SimpleITK, scikit-image, pandas, 3 other tools, structural MRI / diffusion, 2 references
[9] doi:10.1002/hipo.70124 [code]
Association Between Anterior Hippocampal Gyrification and Episodic Memory Performance in Neurotypical Young Adults.
Journal: Hippocampus
In common: SimpleITK, scikit-image, PyTorch, 4 other tools, structural MRI / diffusion, 1 reference
[10] doi:10.1186/s12880-026-02335-x [code]
Automatic lateral ventricle and choroid plexus segmentation method in infant brain MR images.
Journal: BMC medical imaging
In common: SimpleITK, scikit-image, PyTorch, 4 other tools, structural MRI / diffusion, 1 reference

Contribute

The authors of this paper can claim it, correct its record and validate its tracing map, and the maintainers of its code (its owner, or a public member of its organization) correct what it says of their repository; anyone signed in can ask for its removal. Every request goes to OSCR's own machine, which answers it; your account page follows them.

Sign in with ORCID to claim this paper as one of its authors, correct its record or validate its tracing map: when the paper's metadata lists your ORCID iD, you are recognized at once. Maintainers of its code: sign in with GitHub, then claim the repository on your account page.

Request its removal

To ask OSCR to remove this record, the copies of its authors' scripts or its tracing map, use the removal request page: signed in, you say who you are, what to remove and why, then review and confirm the request. Published rules decide every request (how).

Discussion, reproductions, activity

Discussion: questions and error reports about this paper and its code, from signed-in readers and its authors. It opens with sign-in.

Reproductions: reports from readers who ran the authors' code: what they reproduced, with which environment, commit and data. It opens with sign-in.

Activity: what happens around this paper: new versions of its record, its map's validation, discussions and reproductions. It opens with sign-in.