关于GCU、沐曦GPGPU、MLU、0卡V100资源4月7日恢复上架的公告>>> 关于共建具身智能开源数据集的倡议>>> 关于云脑任务中统一路径访问方式的公告>>> 关于将启智集群GPU资源迁移至智算集群的公告>>>
OpenMedIA
/
WeaklyLesionLocalisation
Not watched
4
15
Fork 20
Code
Releases 0 Wiki Activity
Issues 0 Pull Requests 0 Datasets Cloudbrain
Pull Requests
/
Detail
Labels Milestones
New Pull Request

#7 master

Closed
shirley111222 wants to merge 6 commits from shirley111222:master into master
Conversation 0 Commits 6 Files Changed 10
10 changed files with 82 additions and 1012 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. +0
    -1
      .gitignore
  2. +0
    -208
      LICENSE
  3. +77
    -0
      README2.md
  4. BIN
      imgs/model.png
  5. +5
    -0
      requiresment22.txt
  6. +0
    -1
      u_d/__init__.py
  7. +0
    -273
      u_d/__pycache__/base.py
  8. +0
    -128
      u_d/__pycache__/update_u_d.py
  9. +0
    -273
      u_d/base.py
  10. +0
    -128
      u_d/update_u_d.py

+ 0
- 1
.gitignore View File

@@ -1 +0,0 @@
/tags

+ 0
- 208
LICENSE View File

@@ -1,208 +0,0 @@
Apache License

Version 2.0, January 2004

http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION,
AND DISTRIBUTION

1. Definitions.


"License" shall mean the terms and conditions for use, reproduction, and distribution
as defined by Sections 1 through 9 of this document.


"Licensor" shall mean the copyright owner or entity authorized by the copyright
owner that is granting the License.


"Legal Entity" shall mean the union of the acting entity and all other entities
that control, are controlled by, or are under common control with that entity.
For the purposes of this definition, "control" means (i) the power, direct
or indirect, to cause the direction or management of such entity, whether
by contract or otherwise, or (ii) ownership of fifty percent (50%) or more
of the outstanding shares, or (iii) beneficial ownership of such entity.


"You" (or "Your") shall mean an individual or Legal Entity exercising permissions
granted by this License.


"Source" form shall mean the preferred form for making modifications, including
but not limited to software source code, documentation source, and configuration
files.


"Object" form shall mean any form resulting from mechanical transformation
or translation of a Source form, including but not limited to compiled object
code, generated documentation, and conversions to other media types.


"Work" shall mean the work of authorship, whether in Source or Object form,
made available under the License, as indicated by a copyright notice that
is included in or attached to the work (an example is provided in the Appendix
below).


"Derivative Works" shall mean any work, whether in Source or Object form,
that is based on (or derived from) the Work and for which the editorial revisions,
annotations, elaborations, or other modifications represent, as a whole, an
original work of authorship. For the purposes of this License, Derivative
Works shall not include works that remain separable from, or merely link (or
bind by name) to the interfaces of, the Work and Derivative Works thereof.


"Contribution" shall mean any work of authorship, including the original version
of the Work and any modifications or additions to that Work or Derivative
Works thereof, that is intentionally submitted to Licensor for inclusion in
the Work by the copyright owner or by an individual or Legal Entity authorized
to submit on behalf of the copyright owner. For the purposes of this definition,
"submitted" means any form of electronic, verbal, or written communication
sent to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems, and
issue tracking systems that are managed by, or on behalf of, the Licensor
for the purpose of discussing and improving the Work, but excluding communication
that is conspicuously marked or otherwise designated in writing by the copyright
owner as "Not a Contribution."


"Contributor" shall mean Licensor and any individual or Legal Entity on behalf
of whom a Contribution has been received by Licensor and subsequently incorporated
within the Work.

2. Grant of Copyright License. Subject to the terms and conditions of this
License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive,
no-charge, royalty-free, irrevocable copyright license to reproduce, prepare
Derivative Works of, publicly display, publicly perform, sublicense, and distribute
the Work and such Derivative Works in Source or Object form.

3. Grant of Patent License. Subject to the terms and conditions of this License,
each Contributor hereby grants to You a perpetual, worldwide, non-exclusive,
no-charge, royalty-free, irrevocable (except as stated in this section) patent
license to make, have made, use, offer to sell, sell, import, and otherwise
transfer the Work, where such license applies only to those patent claims
licensable by such Contributor that are necessarily infringed by their Contribution(s)
alone or by combination of their Contribution(s) with the Work to which such
Contribution(s) was submitted. If You institute patent litigation against
any entity (including a cross-claim or counterclaim in a lawsuit) alleging
that the Work or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses granted to You
under this License for that Work shall terminate as of the date such litigation
is filed.

4. Redistribution. You may reproduce and distribute copies of the Work or
Derivative Works thereof in any medium, with or without modifications, and
in Source or Object form, provided that You meet the following conditions:

(a) You must give any other recipients of the Work or Derivative Works a copy
of this License; and

(b) You must cause any modified files to carry prominent notices stating that
You changed the files; and

(c) You must retain, in the Source form of any Derivative Works that You distribute,
all copyright, patent, trademark, and attribution notices from the Source
form of the Work, excluding those notices that do not pertain to any part
of the Derivative Works; and

(d) If the Work includes a "NOTICE" text file as part of its distribution,
then any Derivative Works that You distribute must include a readable copy
of the attribution notices contained within such NOTICE file, excluding those
notices that do not pertain to any part of the Derivative Works, in at least
one of the following places: within a NOTICE text file distributed as part
of the Derivative Works; within the Source form or documentation, if provided
along with the Derivative Works; or, within a display generated by the Derivative
Works, if and wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and do not modify the
License. You may add Your own attribution notices within Derivative Works
that You distribute, alongside or as an addendum to the NOTICE text from the
Work, provided that such additional attribution notices cannot be construed
as modifying the License.

You may add Your own copyright statement to Your modifications and may provide
additional or different license terms and conditions for use, reproduction,
or distribution of Your modifications, or for any such Derivative Works as
a whole, provided Your use, reproduction, and distribution of the Work otherwise
complies with the conditions stated in this License.

5. Submission of Contributions. Unless You explicitly state otherwise, any
Contribution intentionally submitted for inclusion in the Work by You to the
Licensor shall be under the terms and conditions of this License, without
any additional terms or conditions. Notwithstanding the above, nothing herein
shall supersede or modify the terms of any separate license agreement you
may have executed with Licensor regarding such Contributions.

6. Trademarks. This License does not grant permission to use the trade names,
trademarks, service marks, or product names of the Licensor, except as required
for reasonable and customary use in describing the origin of the Work and
reproducing the content of the NOTICE file.

7. Disclaimer of Warranty. Unless required by applicable law or agreed to
in writing, Licensor provides the Work (and each Contributor provides its
Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied, including, without limitation, any warranties
or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR
A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness
of using or redistributing the Work and assume any risks associated with Your
exercise of permissions under this License.

8. Limitation of Liability. In no event and under no legal theory, whether
in tort (including negligence), contract, or otherwise, unless required by
applicable law (such as deliberate and grossly negligent acts) or agreed to
in writing, shall any Contributor be liable to You for damages, including
any direct, indirect, special, incidental, or consequential damages of any
character arising as a result of this License or out of the use or inability
to use the Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all other commercial
damages or losses), even if such Contributor has been advised of the possibility
of such damages.

9. Accepting Warranty or Additional Liability. While redistributing the Work
or Derivative Works thereof, You may choose to offer, and charge a fee for,
acceptance of support, warranty, indemnity, or other liability obligations
and/or rights consistent with this License. However, in accepting such obligations,
You may act only on Your own behalf and on Your sole responsibility, not on
behalf of any other Contributor, and only if You agree to indemnify, defend,
and hold each Contributor harmless for any liability incurred by, or claims
asserted against, such Contributor by reason of your accepting any such warranty
or additional liability. END OF TERMS AND CONDITIONS

APPENDIX: How to apply the Apache License to your work.

To apply the Apache License to your work, attach the following boilerplate
notice, with the fields enclosed by brackets "[]" replaced with your own identifying
information. (Don't include the brackets!) The text should be enclosed in
the appropriate comment syntax for the file format. We also recommend that
a file or class name and description of purpose be included on the same "printed
page" as the copyright notice for easier identification within third-party
archives.

Copyright [yyyy] [name of copyright owner]

Licensed under the Apache License, Version 2.0 (the "License");

you may not use this file except in compliance with the License.

You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software

distributed under the License is distributed on an "AS IS" BASIS,

WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

See the License for the specific language governing permissions and

limitations under the License.

+ 77
- 0
README2.md View File

@@ -0,0 +1,77 @@
#
<h1 align="left"> An Official PyTorch Implementation of “Towards Unbiased COVID-19 Lesion Localisation and Segmentation via Weakly Supervised Learning (ISBI-2021)” <a href="https://arxiv.org/abs/2103.00780"><img src="https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg" ></a></h1>

by Yang Yang, Jiancong Chen, Ruixuan Wang*, Ting Ma, Lingwei Wang, Jie Chen, Wei-Shi Zheng, Tong Zhang*.



## Framwork

![Model](./imgs/model.png)



```
@inproceedings{yang2021towards,
title={Towards Unbiased COVID-19 Lesion Localisation and Segmentation via Weakly Supervised Learning},
author={Yang, Yang and Chen, Jiancong and Wang, Ruixuan and Ma, Ting and Wang, Lingwei and Chen, Jie and Zheng, Wei-Shi and Zhang, Tong},
booktitle={IEEE International Symposium on Biomedical Imaging},
pages={1966--1970},
year={2021},
organization={IEEE}
}
```
## Abstract
Despite tremendous efforts, it is very challenging to generate a robust model to assist in the accurate quantification assessment of COVID-19 on chest CT images. Due to the nature of blurred boundaries, the supervised segmentation methods usually suffer from annotation biases. To support unbiased lesion localisation and to minimise the labelling costs, we propose a data-driven framework supervised by only image level labels. The framework can explicitly separate potential lesions from original images, with the help of an generative adversarial network and a lesion-specific decoder. Experiments on two COVID-19 datasets demonstrates the effectiveness of the proposed framework and its superior performance to several existing methods.

## Installation

* Python 3.8+

* Pytorch 1.10+

* Torchvision 0.11.0

* GPU, such as a 1080Ti

You can install require modules with following command.

```bash
pip install -r requirements.txt
```



## Data

Data is stored in ./data, you can directory download it (~400M).



## Training and Evaluation

You can train with following code, which save logs and models to /exp/exp-1.

```bash
python train.py -p ./exp/exp-1 -l 10 -s 0.01 -a 0.01 -g 200 -t 0 -i 15 -b 4 -e 120 --debug
```

## Citation

If you find this project useful for your research, please cit our work use the following BibTeX entry.

@INPROCEEDINGS{9433806,
author={Yang, Yang and Chen, Jiancong and Wang, Ruixuan and Ma, Ting and Wang, Lingwei and Chen, Jie and Zheng, Wei-Shi and Zhang, Tong},
booktitle={2021 IEEE 18th International Symposium on Biomedical Imaging (ISBI)},
title={Towards Unbiased Covid-19 Lesion Localisation And Segmentation Via Weakly Supervised Learning},
year={2021},
volume={},
number={},
pages={1966-1970},
doi={10.1109/ISBI48211.2021.9433806}}


## Contact

If you have any questions about this paper, welcome to email to [zhangt02@pcl.ac.cn](mailto:zhangt02@pcl.ac.cn)


BIN
imgs/model.png View File

Before After
Width: 2342  |  Height: 996  |  Size: 337 KiB

+ 5
- 0
requiresment22.txt View File

@@ -0,0 +1,5 @@
tensorboardX
torch
torchvision
matplotlib
opencv-python

+ 0
- 1
u_d/__init__.py View File

@@ -1 +0,0 @@
from u_d import update_u_d

+ 0
- 273
u_d/__pycache__/base.py View File

@@ -1,273 +0,0 @@
import os
import time
import numpy as np
import matplotlib.pyplot as plt
import torch
from torch.nn import DataParallel
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader
from torchvision.transforms import transforms
from networks.discriminator import get_discriminator
from networks.resnet import resnet18
from networks.unet import UNet
# from utils.Logger import Logger
from utils.read_data import ConcatDataset
from utils.util import add_prefix, weight_to_cpu, rgb2gray, write_list, copy, write, AddGaussianNoise
from tensorboardX import SummaryWriter
plt.switch_backend('agg')
class base(object):
def __init__(self, args):
"""
"""
self.debug = args.debug
self.prefix = args.prefix
self.pretrain_unet_path = args.pretrain_unet_path
self.is_pretrained_unet = args.is_pretrained_unet
self.use_gpu = torch.cuda.is_available()
self.epoch_interval = 1 if self.debug else 5
self.power = args.power
self.data = args.data
self.batch_size = args.batch_size
self.epsi = args.epsi
self.gan_type = args.gan_type
self.u_depth = args.u_depth
self.d_depth = args.d_depth
self.dowmsampling = args.dowmsampling
self.lr = args.lr
self.beta1 = args.beta1
self.eta = args.eta
self.interval = args.interval
self.epochs = args.epochs
self.local = args.local
# self.logger = Logger(add_prefix(self.prefix, 'tensorboard'))
self.mean, self.std = 0.5, 0.5
self.dataloader = self.get_dataloader()
self.auto_encoder = self.get_unet()
self.d = get_discriminator(self.gan_type, self.d_depth, self.dowmsampling)
self.log_lst = []
self.tb = SummaryWriter(log_dir=os.path.join(self.prefix, 'tb'))
if self.use_gpu:
self.auto_encoder = DataParallel(self.auto_encoder, device_ids=[0]).cuda()
self.d = DataParallel(self.d, device_ids=[0]).cuda()
else:
raise RuntimeError('there is no gpu available.')
self.save_init_paras()
self.get_optimizer()
self.save_hyperparameters(args)
def save_hyperparameters(self, args):
write(vars(args), add_prefix(self.prefix, 'para.txt'))
print('save hyperparameters successfully.')
def train(self, epoch):
pass
def validate(self, epoch):
self.d.eval()
self.auto_encoder.eval()
real_data_score = []
fake_data_score = []
for i, (lesion_data, _, lesion_names, _, real_data, _, normal_names, _) in enumerate(self.dataloader):
if i > 2:
break
if self.use_gpu:
lesion_data, real_data = lesion_data.cuda(), real_data.cuda()
phase = 'lesion_data'
prefix_path = '%s/epoch_%d/%s' % (self.prefix, epoch, phase)
lesion_output = self.d(self.auto_encoder(lesion_data)[0])
fake_data_score += list(lesion_output.squeeze().cpu().data.numpy().flatten())
for idx in range(self.batch_size):
single_image = lesion_data[idx:(idx + 1), :, :, :]
single_name = lesion_names[idx]
self.save_single_image(prefix_path, single_name, single_image)
if self.debug:
if idx > 10:
break
phase = 'normal_data'
prefix_path = '%s/epoch_%d/%s' % (self.prefix, epoch, phase)
normal_output = self.d(real_data)
real_data_score += list(normal_output.squeeze().cpu().data.numpy().flatten())
for idx in range(self.batch_size):
single_image = real_data[idx:(idx + 1), :, :, :]
single_name = normal_names[idx]
self.save_single_image(prefix_path, single_name, single_image)
if self.debug:
if idx > 10:
break
prefix_path = '%s/epoch_%d' % (self.prefix, epoch)
self.plot_hist('%s/score_distribution.png' % prefix_path, real_data_score, fake_data_score)
# if not self.debug:
# torch.save(self.auto_encoder.state_dict(), add_prefix(prefix_path, 'g.pkl'))
# torch.save(self.d.state_dict(), add_prefix(prefix_path, 'd.pkl'))
# print('save model parameters successfully when epoch=%d' % epoch)
torch.save(self.auto_encoder.state_dict(), add_prefix(prefix_path, 'g.pkl'))
torch.save(self.d.state_dict(), add_prefix(prefix_path, 'd.pkl'))
print('save model parameters successfully when epoch=%d' % epoch)
def main(self):
print('training start!')
start_time = time.time()
for epoch in range(1, self.epochs + 1):
# self.u_lr_scheduler.step()
# self.c_lr_scheduler.step()
# self.d_lr_scheduler.step()
self.train(epoch)
self.u_lr_scheduler.step()
self.d_lr_scheduler.step()
if epoch % self.epoch_interval == 0:
self.validate(epoch)
self.validate(self.epochs)
total_ptime = time.time() - start_time
print('Training complete in {:.0f}m {:.0f}s'.format(
total_ptime // 60, total_ptime % 60))
def save_init_paras(self):
if not os.path.exists(self.prefix):
os.makedirs(self.prefix)
torch.save(self.auto_encoder.state_dict(), add_prefix(self.prefix, 'init_g_para.pkl'))
torch.save(self.d.state_dict(), add_prefix(self.prefix, 'init_d_para.pkl'))
print('save initial model parameters successfully')
def get_optimizer(self):
self.u_optimizer = torch.optim.Adam(self.auto_encoder.parameters(), lr=self.lr, betas=(self.beta1, 0.9))
self.d_optimizer = torch.optim.Adam(self.d.parameters(), lr=self.lr, betas=(self.beta1, 0.9))
self.u_lr_scheduler = lr_scheduler.ExponentialLR(self.u_optimizer, gamma=self.epsi)
self.d_lr_scheduler = lr_scheduler.ExponentialLR(self.d_optimizer, gamma=self.epsi)
def restore(self, x):
x = torch.squeeze(x)
x = x.data.cpu()
for t, m, s in zip([x], [self.mean], [self.std]):
t.mul_(s).add_(m)
# transform Tensor to numpy
x = x.numpy()
# x = np.transpose(x, (1, 2, 0))
x = np.clip(x * 255, 0, 255).astype(np.uint8)
return x
def get_unet(self):
unet = UNet(1, depth=self.u_depth, in_channels=1)
print(unet)
print('load uent with depth %d and downsampling will be performed for %d times!!' % (
self.u_depth, self.u_depth - 1))
if self.is_pretrained_unet:
unet.load_state_dict(weight_to_cpu(self.pretrain_unet_path))
print('load pretrained unet!')
return unet
def get_dataloader(self):
if self.local:
raise ValueError('')
else:
print('load data from data center.')
if self.data is not None:
print('loading train data')
else:
raise ValueError("the parameter data must be in ['./data/gan']")
transform = transforms.Compose([
# AddGaussianNoise(mean=0, variance=0.01, amplitude=1.0),
transforms.ToTensor(),
transforms.Normalize(self.mean, self.std)
])
dataset = ConcatDataset(data_dir=self.data,
transform=transform,
alpha=self.power
)
data_loader = DataLoader(dataset,
batch_size=self.batch_size,
shuffle=True,
num_workers=2,
drop_last=True,
pin_memory=True if self.use_gpu else False)
return data_loader
def save_single_image(self, saved_path, name, inputs):
"""
save unet output as a form of image
"""
if not os.path.exists(saved_path):
os.makedirs(saved_path)
output, gdp = self.auto_encoder(inputs)
left = self.restore(inputs)
right = self.restore(output)
lesion = self.restore(gdp)
diff = np.where(left > right, left - right, right - left).clip(0, 255).astype(np.uint8)
plt.figure(num='unet result', figsize=(8, 8))
plt.subplot(2, 3, 1)
plt.title('source image')
plt.imshow(left, cmap='gray')
plt.axis('off')
plt.subplot(2, 3, 2)
plt.title('unet output')
plt.imshow(right, cmap='gray')
plt.axis('off')
plt.subplot(2, 3, 3)
plt.title('GDP')
plt.imshow(lesion, cmap='gray')
plt.axis('off')
plt.subplot(2, 3, 4)
plt.imshow(diff, cmap='jet')
plt.colorbar(orientation='horizontal')
plt.title('difference in heatmap')
plt.axis('off')
plt.subplot(2, 3, 5)
plt.imshow(diff.clip(0, 32), cmap='jet')
plt.colorbar(orientation='horizontal')
plt.axis('off')
plt.tight_layout()
plt.savefig(add_prefix(saved_path, name))
plt.close()
def plot_hist(self, path, real_data, fake_data):
bins = np.linspace(min(min(real_data), min(fake_data)), max(max(real_data), max(fake_data)), 60)
plt.hist(real_data, bins=bins, alpha=0.3, label='real_score', edgecolor='k')
plt.hist(fake_data, bins=bins, alpha=0.3, label='fake_score', edgecolor='k')
plt.legend(loc='upper right')
plt.savefig(path)
plt.close()
def save_log(self):
write_list(self.log_lst, add_prefix(self.prefix, 'log.txt'))
print('save running log successfully')
def save_running_script(self, script_path):
"""
save the main running script to get differences between scripts
"""
copy(script_path, add_prefix(self.prefix, script_path.split('/')[-1]))
def get_lr(self):
lr = []
for param_group in self.d_optimizer.param_groups:
lr += [param_group['lr']]
return lr[0]

+ 0
- 128
u_d/__pycache__/update_u_d.py View File

@@ -1,128 +0,0 @@
import torch
from torch import nn
from torch.autograd import grad
from torch.optim import lr_scheduler
from u_d.base import base
from utils.tv_loss import TVLoss
class update_u_d(base):
def __init__(self, args):
base.__init__(self, args)
self.alpha = args.alpha
self.sigma = args.sigma
self.lmbda = args.lmbda
self.gamma = args.gamma
self.theta = args.theta
self.pretrained_steps = args.pretrained_steps
self.l1_criterion = nn.L1Loss(reduce=False).cuda()
self.tv_loss_criterion = TVLoss().cuda()
def train(self, epoch):
l_n = len(self.dataloader)
l_d_real_loss = 0.0
l_d_fake_loss = 0.0
l_d_loss = 0.0
l_d_loss_ = 0.0
l_u_loss_ = 0.0
l_u_d_loss = 0.0
self.d.train()
self.auto_encoder.train()
for idx, data in enumerate(self.dataloader, 1):
step = (epoch - 1) * (len(self.dataloader.dataset) // self.batch_size) + idx
lesion_data, lesion_labels, lesion_name, lesion_gradient, real_data, normal_labels, normal_name, normal_gradient = data
if self.use_gpu:
lesion_data, lesion_labels, normal_gradient = lesion_data.cuda(), lesion_labels.cuda(), normal_gradient.unsqueeze(
1).cuda()
real_data, normal_labels, lesion_gradient = real_data.cuda(), normal_labels.cuda(), lesion_gradient.unsqueeze(
1).cuda()
# training network: update d
self.d_optimizer.zero_grad()
# fake_data[0] is fake image without lesion, fake_data[1] is fake lesion
fake_data, fake_lesion = self.auto_encoder(lesion_data)
real_dis_output = self.d(real_data)
fake_dis_output = self.d(fake_data.detach())
theta = torch.rand((real_data.size(0), 1, 1, 1))
if self.use_gpu:
theta = theta.cuda()
x_hat = theta * real_data.data + (1 - theta) * fake_data.data
x_hat.requires_grad = True
pred_hat = self.d(x_hat)
if self.use_gpu:
gradients = grad(outputs=pred_hat, inputs=x_hat, grad_outputs=torch.ones(pred_hat.size()).cuda(),
create_graph=True, retain_graph=True, only_inputs=True)[0]
else:
gradients = grad(outputs=pred_hat, inputs=x_hat, grad_outputs=torch.ones(pred_hat.size()),
create_graph=True, retain_graph=True, only_inputs=True)[0]
gradient_penalty = self.eta * ((gradients.view(gradients.size()[0], -1).norm(2, 1) - 1) ** 2).mean()
d_real_loss = -torch.mean(real_dis_output)
d_fake_loss = torch.mean(fake_dis_output)
# don't update D too much by set weight
d_loss = self.sigma * (d_real_loss + d_fake_loss + gradient_penalty)
d_loss.backward()
self.d_optimizer.step()
# track loss
l_d_real_loss += d_real_loss.item()
l_d_fake_loss += d_fake_loss.item()
l_d_loss += d_loss.item()
# update u
if step > self.pretrained_steps:
self.u_optimizer.zero_grad()
dis_output = self.d(fake_data)
d_loss_ = -torch.mean(dis_output)
real_data_, real_lesion = self.auto_encoder(real_data)
normal_l1_loss = (normal_gradient * self.l1_criterion((real_data_ + real_lesion), real_data)).mean()
lesion_l1_loss = (lesion_gradient * self.l1_criterion((fake_data + fake_lesion), lesion_data)).mean()
normal_l1_loss_ = (normal_gradient * self.l1_criterion(real_data_, real_data)).mean()
# normal_l1_loss = (self.l1_criterion((real_data_ + real_lesion), real_data)).mean()
# lesion_l1_loss = (self.l1_criterion((fake_data + fake_lesion), lesion_data)).mean()
# normal_l1_loss_ = (self.l1_criterion(real_data_, real_data)).mean()
u_loss_ = normal_l1_loss + normal_l1_loss_ + lesion_l1_loss
u_d_loss = self.alpha * d_loss_ + self.gamma * u_loss_
u_d_loss.backward()
self.u_optimizer.step()
# track loss
l_d_loss_ += d_loss_.item()
l_u_loss_ += u_loss_.item()
l_u_d_loss += u_d_loss.item()
w_distance = d_real_loss.item() + d_fake_loss.item()
if idx % self.interval == 0:
log_1 = '[{}/{}] {}={}*({}(d_real_loss)+{}(d_fake_loss)+{}*{}(gradient_penalty)), '.format(epoch,
self.epochs,
d_loss,
self.sigma,
d_real_loss,
d_fake_loss,
self.eta,
gradient_penalty)
log_2 = 'w_distance:{} ,'.format(w_distance)
log_3 = '{}={}({}(normal_l1_loss)+{}(normal_l1_loss_)+{}(lesion_l1_loss))+{}*{}(d_loss_)'.format(
u_d_loss,
self.gamma,
normal_l1_loss,
normal_l1_loss_,
lesion_l1_loss,
self.alpha, d_loss_)
log = log_1 + log_2 + log_3
print(log)
self.log_lst.append(log)
self.tb.add_scalar('D_loss/d_real_loss', l_d_real_loss, epoch)
self.tb.add_scalar('D_loss/d_fake_loss', l_d_fake_loss, epoch)
self.tb.add_scalar('D_loss/d_loss', l_d_loss, epoch)
self.tb.add_scalar('G_loss/d_loss_', l_d_loss_, epoch)
self.tb.add_scalar('G_loss/u_loss_', l_u_loss_, epoch)
self.tb.add_scalar('G_loss/u_d_loss', l_u_d_loss, epoch)

+ 0
- 273
u_d/base.py View File

@@ -1,273 +0,0 @@
import os
import time
import numpy as np
import matplotlib.pyplot as plt
import torch
from torch.nn import DataParallel
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader
from torchvision.transforms import transforms
from networks.discriminator import get_discriminator
from networks.resnet import resnet18
from networks.unet import UNet
# from utils.Logger import Logger
from utils.read_data import ConcatDataset
from utils.util import add_prefix, weight_to_cpu, rgb2gray, write_list, copy, write, AddGaussianNoise
from tensorboardX import SummaryWriter
plt.switch_backend('agg')
class base(object):
def __init__(self, args):
"""
"""
self.debug = args.debug
self.prefix = args.prefix
self.pretrain_unet_path = args.pretrain_unet_path
self.is_pretrained_unet = args.is_pretrained_unet
self.use_gpu = torch.cuda.is_available()
self.epoch_interval = 1 if self.debug else 5
self.power = args.power
self.data = args.data
self.batch_size = args.batch_size
self.epsi = args.epsi
self.gan_type = args.gan_type
self.u_depth = args.u_depth
self.d_depth = args.d_depth
self.dowmsampling = args.dowmsampling
self.lr = args.lr
self.beta1 = args.beta1
self.eta = args.eta
self.interval = args.interval
self.epochs = args.epochs
self.local = args.local
# self.logger = Logger(add_prefix(self.prefix, 'tensorboard'))
self.mean, self.std = 0.5, 0.5
self.dataloader = self.get_dataloader()
self.auto_encoder = self.get_unet()
self.d = get_discriminator(self.gan_type, self.d_depth, self.dowmsampling)
self.log_lst = []
self.tb = SummaryWriter(log_dir=os.path.join(self.prefix, 'tb'))
if self.use_gpu:
self.auto_encoder = DataParallel(self.auto_encoder, device_ids=[0]).cuda()
self.d = DataParallel(self.d, device_ids=[0]).cuda()
else:
raise RuntimeError('there is no gpu available.')
self.save_init_paras()
self.get_optimizer()
self.save_hyperparameters(args)
def save_hyperparameters(self, args):
write(vars(args), add_prefix(self.prefix, 'para.txt'))
print('save hyperparameters successfully.')
def train(self, epoch):
pass
def validate(self, epoch):
self.d.eval()
self.auto_encoder.eval()
real_data_score = []
fake_data_score = []
for i, (lesion_data, _, lesion_names, _, real_data, _, normal_names, _) in enumerate(self.dataloader):
if i > 2:
break
if self.use_gpu:
lesion_data, real_data = lesion_data.cuda(), real_data.cuda()
phase = 'lesion_data'
prefix_path = '%s/epoch_%d/%s' % (self.prefix, epoch, phase)
lesion_output = self.d(self.auto_encoder(lesion_data)[0])
fake_data_score += list(lesion_output.squeeze().cpu().data.numpy().flatten())
for idx in range(self.batch_size):
single_image = lesion_data[idx:(idx + 1), :, :, :]
single_name = lesion_names[idx]
self.save_single_image(prefix_path, single_name, single_image)
if self.debug:
if idx > 10:
break
phase = 'normal_data'
prefix_path = '%s/epoch_%d/%s' % (self.prefix, epoch, phase)
normal_output = self.d(real_data)
real_data_score += list(normal_output.squeeze().cpu().data.numpy().flatten())
for idx in range(self.batch_size):
single_image = real_data[idx:(idx + 1), :, :, :]
single_name = normal_names[idx]
self.save_single_image(prefix_path, single_name, single_image)
if self.debug:
if idx > 10:
break
prefix_path = '%s/epoch_%d' % (self.prefix, epoch)
self.plot_hist('%s/score_distribution.png' % prefix_path, real_data_score, fake_data_score)
# if not self.debug:
# torch.save(self.auto_encoder.state_dict(), add_prefix(prefix_path, 'g.pkl'))
# torch.save(self.d.state_dict(), add_prefix(prefix_path, 'd.pkl'))
# print('save model parameters successfully when epoch=%d' % epoch)
torch.save(self.auto_encoder.state_dict(), add_prefix(prefix_path, 'g.pkl'))
torch.save(self.d.state_dict(), add_prefix(prefix_path, 'd.pkl'))
print('save model parameters successfully when epoch=%d' % epoch)
def main(self):
print('training start!')
start_time = time.time()
for epoch in range(1, self.epochs + 1):
# self.u_lr_scheduler.step()
# self.c_lr_scheduler.step()
# self.d_lr_scheduler.step()
self.train(epoch)
self.u_lr_scheduler.step()
self.d_lr_scheduler.step()
if epoch % self.epoch_interval == 0:
self.validate(epoch)
self.validate(self.epochs)
total_ptime = time.time() - start_time
print('Training complete in {:.0f}m {:.0f}s'.format(
total_ptime // 60, total_ptime % 60))
def save_init_paras(self):
if not os.path.exists(self.prefix):
os.makedirs(self.prefix)
torch.save(self.auto_encoder.state_dict(), add_prefix(self.prefix, 'init_g_para.pkl'))
torch.save(self.d.state_dict(), add_prefix(self.prefix, 'init_d_para.pkl'))
print('save initial model parameters successfully')
def get_optimizer(self):
self.u_optimizer = torch.optim.Adam(self.auto_encoder.parameters(), lr=self.lr, betas=(self.beta1, 0.9))
self.d_optimizer = torch.optim.Adam(self.d.parameters(), lr=self.lr, betas=(self.beta1, 0.9))
self.u_lr_scheduler = lr_scheduler.ExponentialLR(self.u_optimizer, gamma=self.epsi)
self.d_lr_scheduler = lr_scheduler.ExponentialLR(self.d_optimizer, gamma=self.epsi)
def restore(self, x):
x = torch.squeeze(x)
x = x.data.cpu()
for t, m, s in zip([x], [self.mean], [self.std]):
t.mul_(s).add_(m)
# transform Tensor to numpy
x = x.numpy()
# x = np.transpose(x, (1, 2, 0))
x = np.clip(x * 255, 0, 255).astype(np.uint8)
return x
def get_unet(self):
unet = UNet(1, depth=self.u_depth, in_channels=1)
print(unet)
print('load uent with depth %d and downsampling will be performed for %d times!!' % (
self.u_depth, self.u_depth - 1))
if self.is_pretrained_unet:
unet.load_state_dict(weight_to_cpu(self.pretrain_unet_path))
print('load pretrained unet!')
return unet
def get_dataloader(self):
if self.local:
raise ValueError('')
else:
print('load data from data center.')
if self.data is not None:
print('loading train data')
else:
raise ValueError("the parameter data must be in ['./data/gan']")
transform = transforms.Compose([
# AddGaussianNoise(mean=0, variance=0.01, amplitude=1.0),
transforms.ToTensor(),
transforms.Normalize(self.mean, self.std)
])
dataset = ConcatDataset(data_dir=self.data,
transform=transform,
alpha=self.power
)
data_loader = DataLoader(dataset,
batch_size=self.batch_size,
shuffle=True,
num_workers=2,
drop_last=True,
pin_memory=True if self.use_gpu else False)
return data_loader
def save_single_image(self, saved_path, name, inputs):
"""
save unet output as a form of image
"""
if not os.path.exists(saved_path):
os.makedirs(saved_path)
output, gdp = self.auto_encoder(inputs)
left = self.restore(inputs)
right = self.restore(output)
lesion = self.restore(gdp)
diff = np.where(left > right, left - right, right - left).clip(0, 255).astype(np.uint8)
plt.figure(num='unet result', figsize=(8, 8))
plt.subplot(2, 3, 1)
plt.title('source image')
plt.imshow(left, cmap='gray')
plt.axis('off')
plt.subplot(2, 3, 2)
plt.title('unet output')
plt.imshow(right, cmap='gray')
plt.axis('off')
plt.subplot(2, 3, 3)
plt.title('GDP')
plt.imshow(lesion, cmap='gray')
plt.axis('off')
plt.subplot(2, 3, 4)
plt.imshow(diff, cmap='jet')
plt.colorbar(orientation='horizontal')
plt.title('difference in heatmap')
plt.axis('off')
plt.subplot(2, 3, 5)
plt.imshow(diff.clip(0, 32), cmap='jet')
plt.colorbar(orientation='horizontal')
plt.axis('off')
plt.tight_layout()
plt.savefig(add_prefix(saved_path, name))
plt.close()
def plot_hist(self, path, real_data, fake_data):
bins = np.linspace(min(min(real_data), min(fake_data)), max(max(real_data), max(fake_data)), 60)
plt.hist(real_data, bins=bins, alpha=0.3, label='real_score', edgecolor='k')
plt.hist(fake_data, bins=bins, alpha=0.3, label='fake_score', edgecolor='k')
plt.legend(loc='upper right')
plt.savefig(path)
plt.close()
def save_log(self):
write_list(self.log_lst, add_prefix(self.prefix, 'log.txt'))
print('save running log successfully')
def save_running_script(self, script_path):
"""
save the main running script to get differences between scripts
"""
copy(script_path, add_prefix(self.prefix, script_path.split('/')[-1]))
def get_lr(self):
lr = []
for param_group in self.d_optimizer.param_groups:
lr += [param_group['lr']]
return lr[0]

+ 0
- 128
u_d/update_u_d.py View File

@@ -1,128 +0,0 @@
import torch
from torch import nn
from torch.autograd import grad
from torch.optim import lr_scheduler
from u_d.base import base
from utils.tv_loss import TVLoss
class update_u_d(base):
def __init__(self, args):
base.__init__(self, args)
self.alpha = args.alpha
self.sigma = args.sigma
self.lmbda = args.lmbda
self.gamma = args.gamma
self.theta = args.theta
self.pretrained_steps = args.pretrained_steps
self.l1_criterion = nn.L1Loss(reduce=False).cuda()
self.tv_loss_criterion = TVLoss().cuda()
def train(self, epoch):
l_n = len(self.dataloader)
l_d_real_loss = 0.0
l_d_fake_loss = 0.0
l_d_loss = 0.0
l_d_loss_ = 0.0
l_u_loss_ = 0.0
l_u_d_loss = 0.0
self.d.train()
self.auto_encoder.train()
for idx, data in enumerate(self.dataloader, 1):
step = (epoch - 1) * (len(self.dataloader.dataset) // self.batch_size) + idx
lesion_data, lesion_labels, lesion_name, lesion_gradient, real_data, normal_labels, normal_name, normal_gradient = data
if self.use_gpu:
lesion_data, lesion_labels, normal_gradient = lesion_data.cuda(), lesion_labels.cuda(), normal_gradient.unsqueeze(
1).cuda()
real_data, normal_labels, lesion_gradient = real_data.cuda(), normal_labels.cuda(), lesion_gradient.unsqueeze(
1).cuda()
# training network: update d
self.d_optimizer.zero_grad()
# fake_data[0] is fake image without lesion, fake_data[1] is fake lesion
fake_data, fake_lesion = self.auto_encoder(lesion_data)
real_dis_output = self.d(real_data)
fake_dis_output = self.d(fake_data.detach())
theta = torch.rand((real_data.size(0), 1, 1, 1))
if self.use_gpu:
theta = theta.cuda()
x_hat = theta * real_data.data + (1 - theta) * fake_data.data
x_hat.requires_grad = True
pred_hat = self.d(x_hat)
if self.use_gpu:
gradients = grad(outputs=pred_hat, inputs=x_hat, grad_outputs=torch.ones(pred_hat.size()).cuda(),
create_graph=True, retain_graph=True, only_inputs=True)[0]
else:
gradients = grad(outputs=pred_hat, inputs=x_hat, grad_outputs=torch.ones(pred_hat.size()),
create_graph=True, retain_graph=True, only_inputs=True)[0]
gradient_penalty = self.eta * ((gradients.view(gradients.size()[0], -1).norm(2, 1) - 1) ** 2).mean()
d_real_loss = -torch.mean(real_dis_output)
d_fake_loss = torch.mean(fake_dis_output)
# don't update D too much by set weight
d_loss = self.sigma * (d_real_loss + d_fake_loss + gradient_penalty)
d_loss.backward()
self.d_optimizer.step()
# track loss
l_d_real_loss += d_real_loss.item()
l_d_fake_loss += d_fake_loss.item()
l_d_loss += d_loss.item()
# update u
if step > self.pretrained_steps:
self.u_optimizer.zero_grad()
dis_output = self.d(fake_data)
d_loss_ = -torch.mean(dis_output)
real_data_, real_lesion = self.auto_encoder(real_data)
normal_l1_loss = (normal_gradient * self.l1_criterion((real_data_ + real_lesion), real_data)).mean()
lesion_l1_loss = (lesion_gradient * self.l1_criterion((fake_data + fake_lesion), lesion_data)).mean()
normal_l1_loss_ = (normal_gradient * self.l1_criterion(real_data_, real_data)).mean()
# normal_l1_loss = (self.l1_criterion((real_data_ + real_lesion), real_data)).mean()
# lesion_l1_loss = (self.l1_criterion((fake_data + fake_lesion), lesion_data)).mean()
# normal_l1_loss_ = (self.l1_criterion(real_data_, real_data)).mean()
u_loss_ = normal_l1_loss + normal_l1_loss_ + lesion_l1_loss
u_d_loss = self.alpha * d_loss_ + self.gamma * u_loss_
u_d_loss.backward()
self.u_optimizer.step()
# track loss
l_d_loss_ += d_loss_.item()
l_u_loss_ += u_loss_.item()
l_u_d_loss += u_d_loss.item()
w_distance = d_real_loss.item() + d_fake_loss.item()
if idx % self.interval == 0:
log_1 = '[{}/{}] {}={}*({}(d_real_loss)+{}(d_fake_loss)+{}*{}(gradient_penalty)), '.format(epoch,
self.epochs,
d_loss,
self.sigma,
d_real_loss,
d_fake_loss,
self.eta,
gradient_penalty)
log_2 = 'w_distance:{} ,'.format(w_distance)
log_3 = '{}={}({}(normal_l1_loss)+{}(normal_l1_loss_)+{}(lesion_l1_loss))+{}*{}(d_loss_)'.format(
u_d_loss,
self.gamma,
normal_l1_loss,
normal_l1_loss_,
lesion_l1_loss,
self.alpha, d_loss_)
log = log_1 + log_2 + log_3
print(log)
self.log_lst.append(log)
self.tb.add_scalar('D_loss/d_real_loss', l_d_real_loss, epoch)
self.tb.add_scalar('D_loss/d_fake_loss', l_d_fake_loss, epoch)
self.tb.add_scalar('D_loss/d_loss', l_d_loss, epoch)
self.tb.add_scalar('G_loss/d_loss_', l_d_loss_, epoch)
self.tb.add_scalar('G_loss/u_loss_', l_u_loss_, epoch)
self.tb.add_scalar('G_loss/u_d_loss', l_u_d_loss, epoch)

Write Preview
Loading…
Cancel
Save
Please read the following content carefully:

Dear OpenI User

Thank you for your continuous support to the Openl Qizhi Community AI Collaboration Platform. In order to protect your usage rights and ensure network security, we updated the Openl Qizhi Community AI Collaboration Platform Usage Agreement in January 2024. The updated agreement specifies that users are prohibited from using intranet penetration tools. After you click "Agree and continue", you can continue to use our services. Thank you for your cooperation and understanding.

For more agreement content, please refer to the《Openl Qizhi Community AI Collaboration Platform Usage Agreement》

Agree and continue
Disagree, exit