ControlNets扭转医学图像生成的游戏规定 逾越边界

摘要

文章引见了如何经常使用ControlNets来控制Latent Diffusion Models生成医学图像的环节。首先 探讨了分散模型的开展和其在生成环节中的控制应战，而后引见了ControlNets的概念和长处。接着，文章具体解释了如何训练Latent Diffusion Model和ControlNet，以及如何经常使用ControlNet启动采样和评价。最后，文章展现了经常使用ControlNets生成的医学图像，并提供了性能评价结果。文章的目的是展现这些模型在将脑图像转换为各种对比度的才干，并激励读者在自己的名目中经常使用这些工具。

开篇

本文将引见如何训练一个ControlNet，经常使用户能够准确地控制Latent Diffusion Model（如Stable Diffusion！）的生成环节。咱们的目的是展现模型将脑图像转换为各种对比度的才干。为了成功这一目的，咱们将应用最近推出的MONAI开源裁减，即MONAI Generative Models！

名目代码可以在这个公共仓库中找到：

引言

近年来，文本到图像的分散模型取得了清楚的提高，人们可以依据放开畛域的文本形容生成高真切的图像。这些生成的图像具备丰盛的细节，明晰的轮廓，连接的结构和无心义的场景。但是，虽然分散模型取得了严重的成就，但在生成环节的准确控制方面仍存在应战。即使是内容丰盛的文本形容，也或许不可准确地捕捉到用户的想法。

正如Lvmin Zhang和Maneesh Agrawala在他们的开创性论文“Adding Conditional Control to Text-to-Image Diffusion Models”（2023）中所提出的那样，ControlNets的引入清楚提高了分散模型的可控性和定制性。这些神经网络充任轻量级的适配器，可以准确地控制和定制，同时保管分散模型的原始生成才干。经过对这些适配器启动微调，同时坚持原始分散模型不被修正，可以有效地增强文本到图像模型的多样性。

【编者：裁减模型准许咱们在坚持原始模型不变的同时，对模型的行为启动微调。这象征着咱们可以应用原始模型的生成才干，同时经过ControlNets来调整生成环节，以满足特定的需求或改良性能。】

ControlNet的共同之处在于它处置了空间分歧性的疑问。与以往的方法不同，ControlNet准许对结构空间、结构和几何方面启动明白的控制，同时保管从文本题目中取得的语义控制。原始钻研引入了各种模型，使得可以基于边缘、姿态、语义掩码和深度图启动条件生成，为计算机视觉畛域激动人心的提高铺平了路线。

【编者：空间分歧性是指在图像生成或处置环节中，生成的图像应该在空间上坚持分歧性，即相邻的像素或区域之间应该有正当的相关和延续性。例如，在生成一团体脸的图像时，眼睛、鼻子和嘴巴的相对位置应该是分歧的，不能随机散布。在传统的生成模型中，坚持空间分歧性或许是一个应战，由于模型或许会在尝试生成复杂的图像时发生不分歧的结果。

例如，你正在经常使用一个文本到图像的模型来生成一张猫的图像，来看看如何经常使用边缘、姿态、语义掩码和深度图这四个条件。

1. 边缘：创立一个边缘图来示意猫的轮廓。包括猫的头部、身材、尾巴等关键局部的边缘信息。

2.姿态：创立姿态图来示意猫的姿态。例如，一只正在腾跃的猫，可以在姿态图中示意出这个腾跃的举措。

3.语义掩码：创立一个语义掩码来示意猫的各个局部。例如，在语义掩码中标出猫的眼睛、耳朵、鼻子等局部。

4.深度图：创立深度图来示意猫的三维外形。例如，示意出猫的头部比尾巴更凑近观察者。

经过这四个步骤，就可以指点模型生成一只合乎咱们需求的猫的图像。】

在医学成像畛域，经常会遇到图像转换的运行场景，因此ControlNet的经常使用就十分有价值。在这些运行场景中，有一个场景须要将图像在不同的畛域之间启动转换，例如将计算机断层扫描（CT）转换为磁共振成像（MRI），或许将图像在不同的对比度之间启动转换，例如从T1加权到T2加权的MRI图像。在这篇文章中，咱们将关注一个特定的案例：经常使用从FLAIR图像失掉的脑图像的2D切片来生成相应的T1加权图像。咱们的目的是展现MONAI裁减（MONAI Generative Models）以及ControlNets如何启动医学数据训练，并生成评价模型。经过深化钻研这个例子，咱们能提供关于这些技术在医学成像畛域的最佳通常。

Latent Diffusion Model训练

为了从FLAIR图像生成T1加权（T1w）图像，首先须要训练一个能够生成T1w图像的分散模型。在咱们的例子中，咱们经常使用从英国动物银行数据集（依据这个数据协定可用）中提取的脑MRI图像的2D切片。而后 经常使用你最青睐的方法（例如，ANTs或UniRes）将原始的3D脑部图像注册到MNI空间后，咱们从脑部的核心局部提取五个2D切片。之所以选用这个区域，由于它蕴含各种组织，使得咱们更容易评价图像转换。经常使用这个脚本，咱们最终失掉了大概190,000个切片，空间尺寸为224 × 160像素。接上去，经常使用这个脚本将咱们的图像划分为训练集（约180,000个切片）、验证集（约5,000个切片）和测试集（约5,000个切片）。预备好数据集后，咱们可以开局训练咱们的Latent Diffusion Model了！

为了优化计算资源，潜在分散模型经常使用一个编码器将输入图像x转换为一个低维的潜在空间z，而后可以经过一个解码器启动重构。这种方法使得即使在计算才干有限的状况下也能训练分散模型，同时坚持了它们的原始品质和灵敏性。与咱们在之前的文章中所做的相似（经常使用MONAI生成医学图像），咱们经常使用MONAI Generative models中的KL-regularization模型来创立紧缩模型。经过经常使用这特性能加上L1损失、KL-regularisation，感知损失以及反抗性损失，咱们创立了一个能够以高保真度编码和解码脑图像的自编码器（经常使用这个脚本）。自编码器的重构品质关于Latent Diffusion Model的性能至关关键，由于它定义了生成图像的品质下限。假设自编码器的解码器发生含糊或低品质的图像，咱们的生成模型将不可生成更高品质的图像。

下面这段话用一个例子来解释一下，或许会愈加清楚一些。

假定你是一个艺术家，你的义务是画出一系列的猫的画像。你可以自在地画任何猫，但是你的老板宿愿你画的猫看起来都是"个别的"猫，而不是太奇怪的猫。

这就是你的义务：你须要发明新的猫的画像，同时还要确保这些画像都合乎"个别的"猫的特色。这就像是变分自编码器（VAE）的义务：它须要生成新的数据，同时还要确保这些数据合乎某种预设的散布（也就是"个别的"猫的散布）。

如今，假定你开局画猫。你或许会发现，有些时刻你画的猫看起来太奇怪了，比如它或许有六只眼睛，或许它的尾巴比个别的猫长很多。这时，你的老板或许会提示你，让你画的猫更凑近"个别的"猫。

这就像是KL-regularization的作用：它是一种"处罚"，当你生成的数据偏离预设的散布时，它就会提示你。假设你画的猫太奇怪，你的老板就会提示你。在VAE中，假设生成的数据偏离预设的散布，KL-regularization就会经过参与损失函数的值来提示模型。

经过这种形式，你可以在发明新的猫的画像的同时，还能确保这些画像都合乎"个别的"猫的特色。雷同，VAE也可以在生成新的数据的同时，确保这些数据合乎预设的散布。这就是KL-regularization的关键作用。】

【编者：对三种损失函数启动简明说明：

1. L1损失：L1损失，也称为相对值损失，是预测值和实在值之间差异的相对值的平均。它的公式为 L1 = 1/n Σ|yi - xi|，其中yi是实在值，xi是预测值，n是样本数量。L1损失对意外值不敏感，由于它不会适度处罚预测失误的样本。例如，假设咱们预测一个房价为100万，但实践多少钱为110万，L1损失就是10万。

2. 感知损失：感知损失是一种在图像生成义务中罕用的损失函数，它权衡的是生成图像和实在图像在感知层面的差异。感知损失通常经过比拟图像在某个预训练模型（如VGG网络）的某一层的特色示意来计算。这种方法可以捕捉到图像的初级特性，如纹理和外形，而不只仅是像素级的差异。例如，假设咱们生成的猫的图像和实在的猫的图像在色彩上有纤细的差异，但在外形和纹理上是分歧的，那么感知损失或许就会很小。

3. 反抗性损失：反抗性损失是在生成反抗网络（GAN）中经常使用的一种损失函数。在GAN中，生成器的义务是生成看起来像实在数据的假数据，而判断器的义务是辨别实在数据和假数据。反抗性损失就是用来权衡生成器生成的假数据是否诈骗判断器。例如，假设咱们的生成器生成了一张猫的图像，而判断器简直不可辨别这张图像和实在的猫的图像，那么反抗性损失就会很小。】

经常使用这个 ，咱们可以经过经常使用原始图像和它们的重构之间的多尺度结构相似性指数测量（MS-SSIM）来量化自编码器的保真度。在这个例子中，咱们失掉了一个高性能的MS-SSIM目的，等于0.9876。

【编者：多尺度结构相似性指数测量（Multi-Scale Structural Similarity Index, MS-SSIM）是一种用于权衡两幅图像相似度的目的。它是结构相似性指数（Structural Similarity Index, SSIM）的裁减，思考了图像的多尺度信息。

SSIM是一种比传统的均方误差（Mean Squared Error, MSE）或峰值信噪比（Peak Signal-to-Noise Ratio, PSNR）更合乎人眼视觉感知的图像品质评价目的。它思考了图像的亮度、对比度和结构三个方面的信息，而不只仅是像素级的差异。

MS-SSIM则进一步思考了图像的多尺度信息。它经过在不同的尺度（例如，不同的分辨率）上计算SSIM，而后将这些SSIM值启动加权平均，失掉最终的MS-SSIM值。这样可以更好地捕捉到图像的细节和结构信息。

例如，假定咱们有两幅猫的图像，一幅是原始图像，另一幅是咱们的模型生成的图像。咱们可以在不同的尺度（例如，原始尺度、一半尺度、四分之一尺度等）上计算这两幅图像的SSIM值，而后将这些SSIM值启动加权平均，失掉MS-SSIM值。假设MS-SSIM值凑近1，那么说明生成的图像与原始图像十分相似；假设MS-SSIM值远离1，那么说明生成的图像与原始图像有较大的差异。

MS-SSIM罕用于图像处置和计算机视觉的义务中，例如图像紧缩、图像增强、图像生成等，用于评价处置或生成的图像的品质。文中MS-SSIM被用来量化自编码器的保真度，即自编码重视构的图像与原始图像的相似度。失掉的MS-SSIM目的为0.9876，凑近1，说明自编码器的重构品质十分高，重构的图像与原始图像十分相似。】

在咱们训练了自编码器之后，咱们将在潜在空间z上训练diffusion model（分散模型）。分散模型是一个能够经过在一系列期间步上迭代地去噪来从纯噪声图像生成图像的模型。它通经常常使用一个U-Net架构（具备编码器-解码器格局），其中咱们有编码器的层跳过衔接到解码器局部的层（经过长腾跃衔接），使得特色可重用并稳固训练和收敛。

【编者：这段话形容的是训练分散模型的环节，以及分散模型的基本上班原理。

首先，作者提到在训练了自编码器之后，他们将在潜在空间z上训练分散模型。这象征着他们首先经常使用自编码器将输入图像编码为一个低维的潜在空间z，而后在这个潜在空间上训练分散模型。

分散模型是一种生成模型，它的上班原理是从纯噪声图像开局，而后经过在一系列期间步上迭代地去噪，最终生成目的图像。这个环节就像是将一个含糊的图像逐渐明晰起来，直到生成一个明晰的、与目的图像相似的图像。

分散模型通经常常使用一个U-Net架构。U-Net是一种不凡的卷积神经网络，它有一个编码器局部和一个解码器局部，编码器局部用于将输入图像编码为一个潜在空间，解码器局部用于将潜在空间解码为一个输入图像。U-Net的特点是它有一些腾跃衔接，这些衔接将编码器局部的某些层间接衔接到解码器局部的对应层。这些腾跃衔接可以使得编码器局部的特色被重用在解码器局部，这有助于稳固训练环节并减速模型的收敛。】

在训练环节中，Latent Diffusion Model学习了给定这些提示的条件噪声预测。再次，咱们经常使用MONAI来创立和训练这个网络。在这个 来实例化模型，其中训练和评价在代码的 启动。由于咱们在这个教程中对文本提示不太感兴味，所以咱们对一切的图像经常使用了相反的提示（ “脑部的T1加权图像”）。

咱们的Latent Diffusion Model生成的分解脑图像

再次，咱们可以量化生成模型从而优化其性能，这次咱们评价样本的品质（经常使用Fréchet inception distance (FID)）和模型的多样性（计算一组1,000个样本的一切样本对之间的MS-SSIM）。经常使用这对脚本（

如你在前面的图像和结果中所看到的，咱们如今有一个高分辨率图像的模型，品质十分好。但是，咱们对图像的外观没有任何空间控制。为此，咱们将经常使用一个ControlNet来疏导咱们的Latent Diffusion Model的生成。

ControlNet训练

ControlNet架构包括两个关键组成局部：一个是U-Net模型的编码器（可训练版本），包括两边块，以及一个预训练的“锁定”版本的分散模型。在这里，锁定的正本保管了生成才干，而可训练的正本在特定的图像到图像数据集上启动训练，以学习条件控制。这两个组件经过一个“零卷积”层相互衔接——一个1×1的卷积层，其初始化的权重和偏置设为零。卷积权重逐渐从零过渡到优化的参数，确保在初始训练步骤中，可训练和锁定正本的输入与在没有ControlNet的状况下坚持分歧。换句话说，当ControlNet在任何优化之前运行到某些神经网络块时，它不会对深度神经特色引入任何额外的影响或噪声。

【编者：形容ControlNet架构的设计和上班原理。ControlNet架构包括两个关键局部：一个可训练的U-Net模型的编码器，以及一个预训练的、被"锁定"的分散模型。

用一个例子来协助了解。假定你正在学习画画，你有一个教员（预训练的"锁定"的分散模型）和一个可以修正的画布（可训练的U-Net模型的编码器）。你的教员曾经是一个阅历丰盛的艺术家，他的技艺（生成才干）是固定的，不会扭转。而你的画布是可以修正的，你可以在下面尝试不同的画法，学习如何画画（学习条件控制）。

这两个局部经过一个"零卷积"层相互衔接。这个"零卷积"层就像是一个透明的过滤器，它最后不会扭转任何物品（由于它的权重和偏置都初始化为零），所以你可以看到你的教员的原始画作。而后，随着你的学习提高，这个过滤器会逐渐扭转（卷积权重从零过渡到优化的参数），开局对你的教员的画作启动修正，使其更合乎你的格调。

这个设计确保了在初始训练步骤中，可训练局部和"锁定"局部的输入是分歧的，即你的画作最后会和你的教员的画作一样。而后，随着训练的启动，你的画作会逐渐展现出你自己的格调，但依然坚持着你教员的基本技巧。】

经过整合这两个组件，ControlNet使咱们能够控制Diffusion Model（分散模型）的U-Net中每个级别的行为。

在咱们的例子中，咱们在这个 中实例化了ControlNet，经常使用了以上等效的代码片段。

import torchfrom generative.networks.nets import ControlNet, DiffusionModelUNet# Load pre-trained diffusion modeldiffusion_model = DiffusionModelUNet(spatial_dims=2,in_channels=3,out_channels=3,num_res_blocks=2,num_channels=[256, 512, 768],attention_levels=[False, True, True],with_conditioning=True,cross_attention_dim=1024,num_head_channels=[0, 512, 768],)diffusion_model.load_state_dict(torch.load("diffusion_model.pt"))# Create ControlNetcontrolnet = ControlNet(spatial_dims=2,in_channels=3,num_res_blocks=2,num_channels=[256, 512, 768],attention_levels=[False, True, True],with_conditioning=True,cross_attention_dim=1024,num_head_channels=[0, 512, 768],conditioning_embedding_in_channels=1,conditioning_embedding_num_channels=[64, 128, 128, 256],)# Create trainable copy of the diffusion modelcontrolnet.load_state_dict(diffusion_model.state_dict(), strict=False)# Lock the weighht of the diffusion modelfor p in diffusion_model.parameters():p.requires_grad = False

由于咱们经常使用的是Latent Diffusion Model，这须要ControlNets将基于图像的条件转换为相反的潜在空间，以婚配卷积的大小。为此，咱们经常使用一个与完整模型一同训练的 。在咱们的案例中，咱们有三个下采样级别（相似于智能编码器KL），定义在“conditioning_embedding_num_channels=[64, 128, 128, 256]”。由于咱们的条件图像是一个FLAIR图像，只要一个通道，咱们也须要在“conditioning_embedding_in_channels=1”中指定其输入通道的数量。

初始化咱们的网络后，咱们像训练分散模型一样训练它。在以下的代码片段（ ）中，可以看到首先将条件FLAIR图像传递给可训练的网络，并从其跳过衔接中失掉输入。而后，当计算预测的噪声时，这些值被输入到分散模型中。在外部，分散模型将ControlNets的跳过衔接与自己的衔接相加，而后在馈送解码器局部之前（ ）。以下是训练循环的一局部：

# Training Loop...images = batch["t1w"].to(device)cond = batch["flair"].to(device)...noise = torch.randn_like(latent_representation).to(device)noisy_z = scheduler.add_noise(original_samples=latent_representation, noise=noise, timesteps=timesteps)# Compute trainable partdown_block_res_samples, mid_block_res_sample = controlnet(x=noisy_z, timesteps=timesteps, context=prompt_embeds, controlnet_cond=cond)# Using controlnet outputs to control diffusion model behaviournoise_pred = diffusion_model(x=noisy_z,timesteps=timesteps,context=prompt_embeds,down_block_additional_residuals=down_block_res_samples,mid_block_additional_residual=mid_block_res_sample,)# Then compute diffusion model loss as usual...

ControlNet采样和评价

在训练模型之后，咱们可以对它们启动采样和评价。在这里，经常使用测试集中的FLAIR图像来生成条件化的T1w图像。与咱们的训练相似，采样环节十分凑近于分散模型经常使用的环节，惟一的区别是将条件图像传递给训练过的ControlNet，并在每个采样期间步中经常使用其输入来馈送分散模型。如下图所示，咱们生成的图像具备空中间保真度的原始条件，皮层盘旋遵照相似的外形，图像保管了不同组织之间的边界。

测试集中的原始FLAIR图像作为输入到ControlNet（左），生成的T1加权图像（中），和原始的T1加权图像，也就是预期的输入（右）

在咱们对模型的图像启动采样之后，可以量化ControlNet在将图像在不同对比度之间转换时的性能。由于咱们从测试集中失掉了预期的T1w图像，咱们也可以审核它们的差异，并经常使用平均相对误差（MAE）、峰值信噪比（PSNR）和MS-SSIM计算实在和分解图像之间的距离。在咱们的测试集中，当口头这个脚本时，咱们失掉了PSNR= 26.2458+-1.0092，MAE=0.02632+-0.0036和MSSIM=0.9526+-0.0111。

ControlNet为咱们的分散模型提供了令人难以置信的控制，最近的方法曾经裁减了其方法，联合了不同的训练ControlNets（Multi-ControlNet），在同一模型中处置不同类型的条件（T2I适配器），甚至在模型上设置条件（经常使用像ControlNet 1.1这样的方法）。假设这些方法听起来很幽默，不要犹疑，尝试一下！

总结

在这篇文章中，咱们展现了如何经常使用ControlNets来控制Latent Diffusion Models生成医学图像的环节。咱们的目的是展现这些模型在将脑图像转换为各种对比度的才干。为了成功这一目的，咱们应用了最近推出的MONAI的开源裁减，即MONAI Generative Models！咱们宿愿这篇文章能协助你了解如何经常使用这些工具，并激励你在你自己的名目中经常使用它们。

作者引见

崔皓，社区编辑，资深架构师，领有18年的软件开发和架构阅历，10年散布式架构阅历。

原文题目： Controllable Medical Image Generation with ControlNets ，作者：Walter Hugo Lopez Pinaya