yaoppeng / u-net_v2 Goto Github PK

Python 98.29% Shell 0.61% Jupyter Notebook 1.10%

u-net_v2's Introduction

Pytorch implementation of U-Net v2: RETHINKING THE SKIP CONNECTIONS OF U-NET FOR MEDICAL IMAGE SEGMENTATION

nnUNet is the GOAT! Thanks to Fabian et al. for making pure U-Net great again. Less is more.

Please make sure you have installed all the packages with the correct versions as shown in requirements.txt. Most of the issues are caused by incompatible package versions.

The pretrained PVT model: google drive

1 ISIC segmentation

Download the dataset from google drive

set the nnUNet_raw, nnUNet_preprocessed and nnUNet_results environment variable using the following command:

export nnUNet_raw=/path/to/input_raw_dir
export nnUNet_preprocessed=/path/to/preprocessed_dir
export nnUNet_results=/path/to/result_save_dir

run the training and testing using the following command:

python /path/to/U-Net_v2/run/run_training.py dataset_id 2d 0 --no-debug -tr ISICTrainer --c

The nnUNet preprocessed data can be downloaded from ISIC 2017 and ISIC 2018

2. Polyp segmentation

Download the training dataset from google drive and testing dataset from google drive

run the training and testing using the following command:

python /path/to/U-Net_v2/PolypSeg/Train.py

3. On your own data

I only used the 4× downsampled results on my dataset. You may need to modify the code:

f1, f2, f3, f4, f5, f6 = self.encoder(x)

...
f61 = self.sdi_6([f1, f2, f3, f4, f5, f6], f6)
f51 = self.sdi_5([f1, f2, f3, f4, f5, f6], f5)
f41 = self.sdi_4([f1, f2, f3, f4, f5, f6], f4)
f31 = self.sdi_3([f1, f2, f3, f4, f5, f6], f3)
f21 = self.sdi_2([f1, f2, f3, f4, f5, f6], f2)
f11 = self.sdi_1([f1, f2, f3, f4, f5, f6], f1)

and delete the following code:

for i, o in enumerate(seg_outs):
    seg_outs[i] = F.interpolate(o, scale_factor=4, mode='bilinear')

By doing this, you are using all the resolution results rather than the 4× downsampled ones.

The following code snippet shows how to use U-Net v2 in training and testing.

For training:

from unet_v2.UNet_v2 import *

n_classes=2
pretrained_path="/path/to/pretrained/pvt"
model = UNetV2(n_classes=n_classes, deep_supervision=True, pretrained_path=pretrained_path)

x = torch.rand((2, 3, 256, 256))

ys = model(x)  # ys is a list because of deep supervision

Now you can use ys and label to compute the loss and do back-propagation.

In the testing phase:

model.eval()
model.deep_supervision = False

x = torch.rand((2, 3, 256, 256))
y = model(x)  # y is a tensor since the deep supervision is turned off in the testing phase
print(y.shape)  # (2, n_classes, 256, 256)

pred = torch.argmax(y, dim=1)

for convience, the U-Net v2 model file has been copied to ./unet_v2/UNet_v2.py

4. Citation

@article{peng2023u,
  title={U-Net v2: Rethinking the Skip Connections of U-Net for Medical Image Segmentation},
  author={Peng, Yaopeng and Sonka, Milan and Chen, Danny Z},
  journal={arXiv preprint arXiv:2311.17791},
  year={2023}
}

u-net_v2's People

Contributors

Stargazers

Watchers

Forkers

cqray1990 jamesben6688 dl-cnn wstchh xiamenwcy eliza-and-black marenan gityesm eltociear romitavia cwinjet pipixiaipip git-sachin1066 strategist922 yanxin12 anywayany markcvv githublqss

u-net_v2's Issues

作者是否考虑过使用swin-transfomer作为骨干？

作者您好，您新提出的这个想法非常有趣，我在阅读过程中有一个问题，就是你试验过使用swintransformer作为backbone或者用swintransformer块堆叠为图像金字塔结构进行实验吗？效果如何？

python /path/to/U-Net_v2/run/run_training.py dataset_id 2d 0 --no-debug -tr ISICTrainer --c
What is datasetid here ?
For example if I am using ISIC 2017. And is this command should be run in script or can be directly run in jupyter notebook?

训练超参数epochs的选择

作者您好，我看到您在论文中中写道"The maximum number of training epochs is set to
300."，但是github中 PolypSeg/Train.py 中周期数设置为100，请问：论文里息肉分割的评估结果是基于300个还是100个epoch跑的？

我复现Unet-V2的代码，按照Train.py的默认参数跑，（epoch设的100），结果精度都略低于原文的结果，不知是否是周期的原因？

dataset              dsc    iou    mae
-----------------  -----  -----  -----
CVC-300            0.875  0.805  0.011
CVC-ClinicDB       0.937  0.890  0.007
Kvasir             0.923  0.871  0.021
CVC-ColonDB        0.808  0.729  0.025
ETIS-LaribPolypDB  0.757  0.684  0.016
mean               0.860  0.796  0.016

MRI segmentation

I want to use this code to split the brain image, I want to enter a 2D slice and a 3D image, which part of the code should I change?

how to solve the error 'ValueError: too many values to unpack (expected 2)'

Traceback (most recent call last):
File "D:\Unet V2\PolypSeg\Train.py", line 341, in
train(train_loader, model, optimizer, epoch, opt.test_path)
File "D:\Unet V2\PolypSeg\Train.py", line 145, in train
P1, P2= model(images)
ValueError: too many values to unpack (expected 2)

[Question] Training on custom dataset

Dear Author,
Thank you for interesting repo.

I have a question about training on custom dataset.
I am using binary image for segmentation. (1 class)
And I am using dice loss (1-Dice) as the loss function.
But when I train the model with the code below, loss returns a negative number from 1 epoch.

model = UNetV2(n_classes=1, deep_supervision=True, pretrained_path='./pretrained/pvt_v2_b2.pth')
model.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
for epoch in range(num_epochs):
    model.train()  
    loss_batch, val_loss_batch = [], []
    
    # train
    for inputs, labels in tqdm.tqdm(train_dataloader, desc=f'Epoch {epoch + 1}/{num_epochs}'):
        inputs, labels = inputs.to(device).float(), labels.to(device).float()
        # opt zerograd -> model output -> loss -> backward -> opt step
        optimizer.zero_grad()
        outputs = model(inputs)[::-1][-1] # shape
        outputs = torch.sigmoid(outputs) 
        loss = Dice_loss(outputs, labels)
        loss.backward()
        optimizer.step()
        loss_batch.append(loss.detach())
        
    loss_batch = (torch.stack(loss_batch)).mean()
    loss_history.append(loss_batch.cpu())

    print(f"Epoch [{epoch + 1}/{num_epochs}], Training Loss: {loss_batch}")
    #val

labels.shape: (8,1,256,256) #8 is batch size
outputs[::-1][-1].shape:  (8,1,256,256)
outputs[::-1][-2].shape:  (8,1,128,128)

Like original U-Net, I used the model output as model(inputs)[::-1][-1] to match the shape to the labels.
Could you please check if there is a problem in the code?

Use of Optimizer

Hi,

Great job on your work! I've noticed a few discrepancies between the paper and the implementation in your code. Specifically, the optimizer mentioned in the paper is Adam, but in your code, it appears to be SGD (as seen in line 102 of polypTrainer.py). Could you please clarify which optimizer is the correct one used in your research?

Thank you in advance for your response.

Best regards,
Yiwen

my own data set input size

Hello, if I use my own data set, change the input size to 112, and then an error will appear.“y = self.deconv2(f41) + f31
RuntimeError: The size of tensor a (8) must match the size of tensor b (7) at non-singleton dimension 3”，can i solve this problem，Thank you very much.

Small object segmentation result may not good, how to enhance?

Thank you for your work, But I found that Unet_V2 is not better than Unet full resolution baseline in small object segmentation, for example in vessel segmentation task.
How to modify the encoder to output full resolution instead of 4X downsampled, could you give a example, so we can compare more result on different task?

Thank you again!

U-Net_v2是否能用于做实例分割的任务？

我迫切需要使用U-Net_v2（或者U-Net系列）做实例分割的任务，请问您有U-Net_v2做实例分割的示例吗？谢谢！期待您的回复！

U-Net_v2 如何实现高分辨率,多类分割

请问 U-Net_v2 如何实现高分辨率(1280x960）,多类分割？

关于息肉分割结果的一些疑问

您好，我在我的论文中计划和UNetV2对比，我的实验部分有息肉分割实验，但是我的实验设置并不是按照PraNet的设置，而是每个数据集单独训练和测试，我几乎按照论文提供的超参数进行训练，然而得到的分割结果却不太理想，最好Dice也仅是0.5左右，这让我很困惑

模型问题

您好！我有几个问题想提问~
1、为什么您的跳跃连接是用add操作，而不是用cat呢？

2、我在ISIC2018数据集上测试，使用pvt encoder的unet（姑且称为pvt-Unet）结果似乎不如unet，但是pvt-Unet结合SDI模块（也就是您文章中提出的Unet_v2）的结果远远好于unet这是什么原因呢？

3、我在ISIC2018数据集上测试，使用Unet结合SDI模块的结果似乎不如unet，但是pvt-Unet结合SDI模块（也就是您文章中提出的Unet_v2）的结果远远好于unet这是什么原因呢？

不知您是否做过这两个消融实验？能否向您要份相关代码？

下面是我的结果，实验设置是adam（0.5，0.999），lr=0.001，lr每10回合衰减为原来的0.1倍

希望能得到您的指导和回复~

这是我用来实验的pvt-Unet代码
`
class transEncoderUNet(nn.Module):
"""
use SpatialAtt + ChannelAtt
"""
def init(self, channel=32, n_classes=1, deep_supervision=True, pretrained_path=None):
super(transEncoderUNet,self).init()
self.deep_supervision = deep_supervision

    self.encoder = Encoder(pretrained_path)
    self.down4 = down(512, 512)

    self.up1 = up(1024, 320)
    self.up2 = up(640, 128)
    self.up3 = up(256, 64)
    self.up4 = up(128, 64)
    self.outc = outconv(64, 1)


def forward(self, x):
    seg_outs = []
    f1, f2, f3, f4 = self.encoder(x)

    f5 = self.down4(f4)
    #
    y = self.up1(f5, f4)
    seg_outs.append(y)
    y = self.up2(y, f3)
    seg_outs.append(y)
    y = self.up3(y, f2)
    seg_outs.append(y)
    y = self.up4(y, f1)
    y = self.outc(y)
    seg_outs.append(y)

    for i, o in enumerate(seg_outs):
        seg_outs[i] = F.interpolate(o, scale_factor=4, mode='bilinear')

    if self.deep_supervision:
        return seg_outs[::-1]
    else:
        return seg_outs[-1]

这是我用来实验的Unet+SDI代码

`
class SDIUnet(nn.Module):
def init(self, n_channels=3, n_classes=1):
super(SDIUnet, self).init()

    self.inc = inconv(n_channels, 64)
    self.down1 = down(64, 128)
    self.down2 = down(128, 256)
    self.down3 = down(256, 512)
    self.down4 = down(512, 512)

    self.ca_1 = ChannelAttention(64)
    self.sa_1 = SpatialAttention()

    self.ca_2 = ChannelAttention(128)
    self.sa_2 = SpatialAttention()

    self.ca_3 = ChannelAttention(256)
    self.sa_3 = SpatialAttention()

    self.ca_4 = ChannelAttention(512)
    self.sa_4 = SpatialAttention()

    self.Translayer_1 = BasicConv2d(64, 32, 1)
    self.Translayer_2 = BasicConv2d(128, 32, 1)
    self.Translayer_3 = BasicConv2d(256, 32, 1)
    self.Translayer_4 = BasicConv2d(512, 32, 1)

    self.sdi_1 = SDI(32)
    self.sdi_2 = SDI(32)
    self.sdi_3 = SDI(32)
    self.sdi_4 = SDI(32)

    self.up1 = up(544, 32)
    self.up2 = up(64, 32)
    self.up3 = up(64, 32)
    self.up4 = up(64, 32)
    self.outc = outconv(32, n_classes)

def forward(self, x):
    x = x.float()
    x1 = self.inc(x)
    x2 = self.down1(x1)
    x3 = self.down2(x2)
    x4 = self.down3(x3)
    x5 = self.down4(x4)

    f1 = self.ca_1(x1) * x1
    f1 = self.sa_1(f1) * f1
    f1 = self.Translayer_1(f1)

    f2 = self.ca_2(x2) * x2
    f2 = self.sa_2(f2) * f2
    f2 = self.Translayer_2(f2)

    f3 = self.ca_3(x3) * x3
    f3 = self.sa_3(f3) * f3
    f3 = self.Translayer_3(f3)

    f4 = self.ca_4(x4) * x4
    f4 = self.sa_4(f4) * f4
    f4 = self.Translayer_4(f4)

    f41 = self.sdi_4([f1, f2, f3, f4], f4)
    f31 = self.sdi_3([f1, f2, f3, f4], f3)
    f21 = self.sdi_2([f1, f2, f3, f4], f2)
    f11 = self.sdi_1([f1, f2, f3, f4], f1)

    x = self.up1(x5, f41)
    x = self.up2(x, f31)
    x = self.up3(x, f21)
    x = self.up4(x, f11)

    x = self.outc(x)

    return x

Question on #11 issue

Hello, Sir. I have some question on a closed issue #11
I wonder if there is difference in two downsamping ways that "Conv with stride=2" and "MaxPool"? And "avoid too many downsamplings" is equal to "make Unet shallower" ?

Using the same Conv2d layer instance may cause problems

I notice that there is a subtle problem in the initialization part of SDI layer.

Using the code self.convs = nn.ModuleList([nn.Conv2d(channel, channel, kernel_size=3, stride=1, padding=1)] * 4) in Python results in creating a list with four references to the same nn.Conv2d instance because of the list multiplication ([conv] * 4).

This means that instead of having four independent nn.Conv2d instances, all four entries in the module list actually point to the same object and share the same set of weights and biases.

class SDI(nn.Module):
    def __init__(self, channel):
        super().__init__()

        self.convs = nn.ModuleList(
            [nn.Conv2d(channel, channel, kernel_size=3, stride=1, padding=1)] * 4)

And the same problem can be also found in self.seg_outs = nn.ModuleList([nn.Conv2d(channel, n_classes, 1, 1)] * 4).

The correct approach is to use a list comprehension to create separate nn.Conv2d instances:

self.convs = nn.ModuleList(
    [nn.Conv2d(channel, channel, kernel_size=3, stride=1, padding=1) for _ in range(4)]
)

It is recommended that you write like this to support more backbones

It is recommended that you write the UNetV2 class like this to support more backbones, timm==0.9.12

import timm
class UNetV2(nn.Module):
    """
    use SpatialAtt + ChannelAtt
    """
    def __init__(self, channel=32, n_classes=1, deep_supervision=True, backbone ='pvt_v2_b2',pretrained=False):
        super().__init__()
        self.deep_supervision = deep_supervision

        self.encoder = timm.create_model(backbone,pretrained=pretrained,features_only=True,out_indices=(0,1,2,3))
        
        channel1,channel2,channel3,channel4  = self.encoder.feature_info.channels()
        

        self.ca_1 = ChannelAttention(channel1)
        self.sa_1 = SpatialAttention()

        self.ca_2 = ChannelAttention(channel2)
        self.sa_2 = SpatialAttention()

        self.ca_3 = ChannelAttention(channel3)
        self.sa_3 = SpatialAttention()

        self.ca_4 = ChannelAttention(channel4)
        self.sa_4 = SpatialAttention()

        self.Translayer_1 = BasicConv2d(channel1, channel, 1)
        self.Translayer_2 = BasicConv2d(channel2, channel, 1)
        self.Translayer_3 = BasicConv2d(channel3, channel, 1)
        self.Translayer_4 = BasicConv2d(channel4, channel, 1)

        self.sdi_1 = SDI(channel)
        self.sdi_2 = SDI(channel)
        self.sdi_3 = SDI(channel)
        self.sdi_4 = SDI(channel)

        self.seg_outs = nn.ModuleList([
            nn.Conv2d(channel, n_classes, 1, 1)] * 4)

        self.deconv2 = nn.ConvTranspose2d(channel, channel, kernel_size=4, stride=2, padding=1,
                                          bias=False)
        self.deconv3 = nn.ConvTranspose2d(channel, channel, kernel_size=4, stride=2,
                                          padding=1, bias=False)
        self.deconv4 = nn.ConvTranspose2d(channel, channel, kernel_size=4, stride=2,
                                          padding=1, bias=False)
        self.deconv5 = nn.ConvTranspose2d(channel, channel, kernel_size=4, stride=2,
                                          padding=1, bias=False)

模型推理时，图片尺寸的疑问

项目中SDI模块的forward函数，传入了xs和anchor，在for循环中，仅仅考虑了x元素的最后一个维度（w），也就是说默认H=W，但实际情况中很容易输入的H与W并不相等，因此直接如果输入没有resize可能就会引起报错。我的疑问是，如果将模型用于推理，此时H与W不相等，将无法使用。所以是先对图片进行填充，再对图片进行区域裁剪吗？

2.顺便指出一下，network中的演示示例resunet似乎仍有bug，按

`class SDI(nn.Module):
def init(self, channel):
super().init()

    self.convs = nn.ModuleList(
        [nn.Conv2d(channel, channel, kernel_size=3, stride=1, padding=1) for _ in range(4)])

def forward(self, xs, anchor):
    ans = torch.ones_like(anchor)
    target_size = anchor.shape[-1]

    for i, x in enumerate(xs):
        if x.shape[-1] > target_size:
            x = F.adaptive_avg_pool2d(x, (target_size, target_size))
        elif x.shape[-1] < target_size:
            x = F.interpolate(x, size=(target_size, target_size),
                                  mode='bilinear', align_corners=True)

        ans = ans * self.convs[i](x)

RuntimeError: Given groups=1, weight of size [32, 64, 3, 3], expected input[2, 96, 112, 112] to have 64 channels, but got 96 channels instead

    return ans


`RuntimeError: Given groups=1, weight of size [32, 64, 3, 3], expected input[2, 96, 112, 112] to have 64 channels, but got 96 channels instead`

About Training

Download the training dataset from google drive and testing dataset from google drive
run the training using the following command:python /path/to/U-Net_v2/PolypSeg/Train.py. An error occurred in the following statement: P1, P2= model(images)
ValueError: too many values to unpack (expected 2)

关于ISIC数据集的训练问题

作者您好，我想请教一下：
先按照您给的步骤设置了环境变量，并下载了Google Drive的ISIC2017的raw data以及 preprocessed data，但是训练的时候dsc和miou始终都是0%。

==========num_iterations_per_epoch: 250===========
wandb: Network error (ReadTimeout), entering retry loop.
2024-04-10 16:10:45.853054: finished training epoch 3
2024-04-10 16:10:47.647442: Using splits from existing split file: /media/dell/D/cjt/cjt/unetv2/nnunetv2/data/preprocessed_data/Dataset122_ISIC2017/splits_final.json
2024-04-10 16:10:47.664358: The split file contains 1 splits.
2024-04-10 16:10:47.665259: Desired fold for training: 0
2024-04-10 16:10:47.665558: This split has 1500 training and 650 validation cases.
start computing score....
2024-04-10 16:26:39.131107: dsc: 0.00%
2024-04-10 16:26:39.134617: miou: 0.00%
2024-04-10 16:26:39.135679: acc: 83.25%, sen: 0.00%, spe: 100.00%
2024-04-10 16:26:39.137912: current best miou: 0.0 at epoch: 0, (0, 0.0, 0.0)
2024-04-10 16:26:39.138560: current best dsc: 0.0 at epoch: 0, (0, 0.0, 0.0)
2024-04-10 16:26:39.139042: finished real validation
/media/dell/D/cjt/cjt/unetv2/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py:1115: RuntimeWarning: Mean of empty slice

我对_internal_maybe_mirror_and_predict函数做了一些修改，因为用深度监督的时候，会输出一个tuple，我取了tuple中的第二个作为输出。如下所示，另外我看raw_data的每张mask都是全黑的，应该是作者预处理把255像素值调到了1吧，以上会对训练有影响吗？

def _internal_maybe_mirror_and_predict(self, x: torch.Tensor) -> torch.Tensor:
    mirror_axes = self.allowed_mirroring_axes if self.use_mirroring else None
    x = x.to(torch.float16)
    # prediction = self.network(x)
    prediction = self.network(x)[1]
    # print("x.shape", x.shape)

    if mirror_axes is not None:
        # check for invalid numbers in mirror_axes
        # x should be 5d for 3d images and 4d for 2d. so the max value of mirror_axes cannot exceed len(x.shape) - 3
        assert max(mirror_axes) <= len(x.shape) - 3, 'mirror_axes does not match the dimension of the input!'

        num_predictons = 2 ** len(mirror_axes)
        if 0 in mirror_axes:
            # prediction += torch.flip(self.network(torch.flip(x, (2,))), (2,))
            #预测只取最后一个
            prediction += torch.flip(self.network(torch.flip(x, (2,)))[1], (2,))
        if 1 in mirror_axes:
            prediction += torch.flip(self.network(torch.flip(x, (3,)))[1], (3,))
        if 2 in mirror_axes:
            prediction += torch.flip(self.network(torch.flip(x, (4,)))[1], (4,))
        if 0 in mirror_axes and 1 in mirror_axes:
            prediction += torch.flip(self.network(torch.flip(x, (2, 3)))[1], (2, 3))
        if 0 in mirror_axes and 2 in mirror_axes:
            prediction += torch.flip(self.network(torch.flip(x, (2, 4)))[1], (2, 4))
        if 1 in mirror_axes and 2 in mirror_axes:
            prediction += torch.flip(self.network(torch.flip(x, (3, 4)))[1], (3, 4))
        if 0 in mirror_axes and 1 in mirror_axes and 2 in mirror_axes:
            prediction += torch.flip(self.network(torch.flip(x, (2, 3, 4)))[1], (2, 3, 4))
        prediction /= num_predictons
        # if 1 in mirror_axes:
        #     prediction += torch.flip(self.network(torch.flip(x, (3,))), (3,))
        # if 2 in mirror_axes:
        #     prediction += torch.flip(self.network(torch.flip(x, (4,))), (4,))
        # if 0 in mirror_axes and 1 in mirror_axes:
        #     prediction += torch.flip(self.network(torch.flip(x, (2, 3))), (2, 3))
        # if 0 in mirror_axes and 2 in mirror_axes:
        #     prediction += torch.flip(self.network(torch.flip(x, (2, 4))), (2, 4))
        # if 1 in mirror_axes and 2 in mirror_axes:
        #     prediction += torch.flip(self.network(torch.flip(x, (3, 4))), (3, 4))
        # if 0 in mirror_axes and 1 in mirror_axes and 2 in mirror_axes:
        #     prediction += torch.flip(self.network(torch.flip(x, (2, 3, 4))), (2, 3, 4))
        # prediction /= num_predictons
    return prediction