Repository: bkong999/COVNet
Branch: master
Commit: 932c1cf7e4a0
Files: 11
Total size: 35.1 KB
Directory structure:
gitextract_17ft28df/
├── .gitignore
├── LICENSE
├── README.md
├── config.py
├── data/
│ ├── train.csv
│ └── val.csv
├── dataset.py
├── main.py
├── model.py
├── test.py
└── utils.py
================================================
FILE CONTENTS
================================================
================================================
FILE: .gitignore
================================================
#cmake related files
CMakeCache.txt
CMakeFiles
CMakeScripts
Makefile
cmake_install.cmake
install_manifest.txt
CTestTestfile.cmake
#ignore all bin dirs
bin/
.DS_Store
*.vtp
*.pathlinesinfo
*.pyc
*.pyo
*~
*.bak
*.swp
*.o
*.nii.gz
*.mha
*.raw
*.mhd
*.a
*.orig # git merge conflict file
# model file
*.pth
*.p
### following copied from https://gist.github.com/octocat/9257657
# Compiled source #
###################
*.com
*.class
*.dll
*.exe
*.o
*.so
# Packages #
############
# it's better to unpack these files and commit the raw source
# git has its own built in compression methods
*.7z
*.dmg
*.gz
*.iso
*.jar
*.rar
*.tar
*.zip
# Logs and databases #
######################
*.log
*.sql
*.sqlite
# OS generated files #
######################
.DS_Store
.DS_Store?
._*
.Spotlight-V100
.Trashes
ehthumbs.db
Thumbs.db
# Ignore data and experiments dirs, but keep folder organizations#
######################
experiments
================================================
FILE: LICENSE
================================================
“Commons Clause” License Condition v1.0
The Software is provided to you by the Licensor under the
License, as defined below, subject to the following condition.
Without limiting other conditions in the License, the grant
of rights under the License will not include, and the License
does not grant to you, the right to Sell the Software.
For purposes of the foregoing, “Sell” means practicing any or
all of the rights granted to you under the License to provide
to third parties, for a fee or other consideration (including
without limitation fees for hosting or consulting/ support
services related to the Software), a product or service whose
value derives, entirely or substantially, from the functionality
of the Software. Any license notice or attribution required by
the License must also include this Commons Clause License
Condition notice.
Software: COVNet
License: Apache 2.0
Licensor: CuraCloud Corporation
---------------------------------------------------------------------
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.
================================================
FILE: README.md
================================================
# COVID-19 Detection Neural Network (COVNet)
This is a PyTorch implementation of the paper "[Artificial Intelligence Distinguishes COVID-19 from Community Acquired Pneumonia on Chest CT](https://pubs.rsna.org/doi/10.1148/radiol.2020200905)". It supports training, validation and testing for COVNet.
<img src="assets/overview.png" width="600">
### Updates & Notices
- 2020-03-30: Thanks for the interest in our work. Unfortunately, we do not own the data, and we have to get permission from our collaborators before we share the data and model. We will update later.
## Citation
If you find this code is useful for your research, please consider citing:
```
@article{li2020artificial,
title={Artificial Intelligence Distinguishes COVID-19 from Community Acquired Pneumonia on Chest CT},
author={Li, Lin and Qin, Lixin and Xu, Zeguo and Yin, Youbing and Wang, Xin and Kong, Bin and Bai, Junjie and Lu, Yi and Fang, Zhenghan and Song, Qi and Cao, Kunlin and others},
journal={Radiology},
year={2020}
}
```
## Setup
### Prerequisites
- Anaconda 3.7
- PyTorch 1.4
- SimpleITK
- batchgenerators
- tensorboardX
### Prepare data
Preprocess the data according to the [Appendix E1 section](https://pubs.rsna.org/doi/suppl/10.1148/radiol.2020200905/suppl_file/ry_200905_supp_in%20press.pdf) of the paper and organize them as the following. A example of train.csv and val.csv are also provided.
```
data
├── caseid1
| ├── masked_ct.nii
| └── mask.nii.gz
├── caseid2
| ├── masked_ct.nii
| └── mask.nii.gz
├── caseid3
| ├── masked_ct.nii
| └── mask.nii.gz
├── caseid4
| ├── masked_ct.nii
| └── mask.nii.gz
├── train.csv
└── val.csv
```
## COVNet
<img src="assets/demo.png" width="600">
### Training
Training a COVNet with default arguments. Model checkpoints and tensorboard logs are written out to a unique directory created by default within `experiments/models` and `experiments/logs` respectively after starting training.
```
python main.py
```
### Validation and Testing
You can run validation and testing on the checkpointed best model by:
```
python test.py
```
================================================
FILE: config.py
================================================
import argparse
def parse_arguments():
"""Argument Parser for the commandline argments
:returns: command line arguments
"""
##########################################################################
# Training setting #
##########################################################################
parser = argparse.ArgumentParser()
parser.add_argument('--n_classes', help="Number of classes", type=int,
default=3)
parser.add_argument('--lr', type=float, default=1e-5)
parser.add_argument('--epochs', type=int, default=50)
parser.add_argument('--lr_scheduler', type=str,
default='plateau', choices=['plateau', 'step'])
parser.add_argument('--gamma', type=float,
help='LR Multiplicative factor if lr_scheduler is step',
default=0.1)
parser.add_argument('--patience', type=int, default=9)
parser.add_argument('--log-every', type=int, default=100)
parser.add_argument('--save-model', type=bool, default=True)
args = parser.parse_args()
return args
================================================
FILE: data/train.csv
================================================
case,label
caseid1,0
caseid2,1
caseid3,2
================================================
FILE: data/val.csv
================================================
case,label
caseid4,0
================================================
FILE: dataset.py
================================================
import os
import pandas as pd
import numpy as np
import torch
import torch.utils.data as data
import SimpleITK as sitk
from batchgenerators.transforms import noise_transforms
from batchgenerators.transforms import spatial_transforms
class NCovDataset(data.Dataset):
def __init__(self, root_dir, stage='train'):
super().__init__()
self.root_dir = root_dir
self.stage = stage
assert stage in ['train', 'val', 'test']
if stage == 'train':
split_file = 'train.csv'
elif stage == 'val':
split_file = 'val.csv'
elif stage == 'test':
# We just assume validation set is the same as test set
split_file = 'val.csv'
df = pd.read_csv(os.path.join(root_dir, split_file),
converters={'case': str, 'label': int})
df = df.set_index('case')
self.case_ids = list(df.index)
self.labels = df['label'].values.tolist()
def __len__(self):
return len(self.case_ids)
def __getitem__(self, index):
fn = os.path.join(self.root_dir, self.case_ids[index], 'masked_ct.nii')
image = sitk.ReadImage(fn)
array = sitk.GetArrayFromImage(image)
mask_fn = os.path.join(self.root_dir, self.case_ids[index],
'mask.nii.gz')
mask_image = sitk.ReadImage(mask_fn)
mask = sitk.GetArrayFromImage(mask_image)
array, mask = array[None, ...], mask[None, ...]
if self.stage == 'train':
# Default randomly mirroring the second and third axes
array, mask = spatial_transforms.augment_mirroring(
array, sample_seg=mask, axes=(1, 2))
array, mask = array[0], mask[0]
######################################################
# Preprocessing for both train and validation data #
######################################################
min_value, max_value = -1250, 250
np.clip(array, min_value, max_value, out=array)
array = (array - min_value) / (max_value - min_value)
# data should be a numpy array with shape [x, y, z] or [c, x, y, z]
# seg should be a numpy array with shape [x, y, z]
full_channel = np.stack([array, array, array])
if self.stage == 'train':
full_channel, mask = self.do_augmentation(full_channel, mask)
else:
mask = mask[None, ...]
# remove the noise in the non-lung regions
mask = mask[0]
full_channel[0][mask == 0] = 0
full_channel[1][mask == 0] = 0
full_channel[2][mask == 0] = 0
label = self.labels[index]
full_channel = torch.FloatTensor(full_channel).permute((1, 0, 2, 3))
return full_channel, label, self.case_ids[index]
def do_augmentation(self, array, mask):
"""Augmentation for the training data.
:array: A numpy array of size [c, x, y, z]
:returns: augmented image and the corresponding mask
"""
# normalize image to range [0, 1], then apply this transform
patch_size = np.asarray(array.shape)[1:]
augmented = noise_transforms.augment_gaussian_noise(
array, noise_variance=(0, .015))
# need to become [bs, c, x, y, z] before augment_spatial
augmented = augmented[None, ...]
mask = mask[None, None, ...]
r_range = (0, (3 / 360.) * 2 * np.pi)
cval = 0.
augmented, mask = spatial_transforms.augment_spatial(
augmented, seg=mask, patch_size=patch_size,
do_elastic_deform=True, alpha=(0., 100.), sigma=(8., 13.),
do_rotation=True, angle_x=r_range, angle_y=r_range, angle_z=r_range,
do_scale=True, scale=(.9, 1.1),
border_mode_data='constant', border_cval_data=cval,
order_data=3,
p_el_per_sample=0.5,
p_scale_per_sample=.5,
p_rot_per_sample=.5,
random_crop=False
)
mask = mask[0]
return augmented[0], mask
def make_weights_for_balanced_classes(self):
"""Making sampling weights for the data samples
:returns: sampling weigghts for dealing with class imbalance problem
"""
n_samples = len(self.labels)
unique, cnts = np.unique(self.labels, return_counts=True)
cnt_dict = dict(zip(unique, cnts))
weights = []
for label in self.labels:
weights.append(n_samples / float(cnt_dict[label]))
return weights
================================================
FILE: main.py
================================================
import os
import numpy as np
import time
from tensorboardX import SummaryWriter
import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
import dataset
import model
import config
import utils
def get_lr(optimizer):
for param_group in optimizer.param_groups:
return param_group['lr']
def train_model(model, train_loader, epoch, num_epochs, optimizer, writer,
current_lr, log_every=100):
n_classes = model.n_classes
metric = torch.nn.CrossEntropyLoss()
y_probs = np.zeros((0, n_classes), np.float)
losses, y_trues = [], []
model.train()
for m in model.modules():
if isinstance(m, nn.BatchNorm2d):
m.train()
m.weight.requires_grad = False
m.bias.requires_grad = False
for i, (image, label, case_id) in enumerate(train_loader):
optimizer.zero_grad()
if torch.cuda.is_available():
image = image.cuda()
label = label.cuda()
prediction = model.forward(image.float())
loss = metric(prediction, label.long())
loss.backward()
optimizer.step()
loss_value = loss.item()
losses.append(loss_value)
y_prob = F.softmax(prediction, dim=1)
y_probs = np.concatenate([y_probs, y_prob.detach().cpu().numpy()])
y_trues.append(label.item())
metric_collects = utils.calc_multi_cls_measures(y_probs, y_trues)
n_iter = epoch * len(train_loader) + i
writer.add_scalar('Train/Loss', loss_value, n_iter)
if (i % log_every == 0) & (i > 0):
utils.print_progress(epoch + 1, num_epochs, i, len(train_loader),
np.mean(losses), current_lr, metric_collects)
train_loss_epoch = np.round(np.mean(losses), 4)
return train_loss_epoch, metric_collects
def evaluate_model(model, val_loader, epoch, num_epochs, writer, current_lr,
log_every=20):
n_classes = model.n_classes
metric = torch.nn.CrossEntropyLoss()
model.eval()
for m in model.modules():
if isinstance(m, nn.BatchNorm2d):
m.train()
m.weight.requires_grad = False
m.bias.requires_grad = False
y_probs = np.zeros((0, n_classes), np.float)
losses, y_trues = [], []
for i, (image, label, case_id) in enumerate(val_loader):
if torch.cuda.is_available():
image = image.cuda()
label = label.cuda()
prediction = model.forward(image.float())
loss = metric(prediction, label.long())
loss_value = loss.item()
losses.append(loss_value)
y_prob = F.softmax(prediction, dim=1)
y_probs = np.concatenate([y_probs, y_prob.detach().cpu().numpy()])
y_trues.append(label.item())
metric_collects = utils.calc_multi_cls_measures(y_probs, y_trues)
n_iter = epoch * len(val_loader) + i
writer.add_scalar('Val/Loss', loss_value, n_iter)
if (i % log_every == 0) & (i > 0):
prefix = '*Val|'
utils.print_progress(epoch + 1, num_epochs, i, len(val_loader),
np.mean(losses), current_lr, metric_collects,
prefix=prefix)
val_loss_epoch = np.round(np.mean(losses), 4)
return val_loss_epoch, metric_collects
def main(args):
"""Main function for the training pipeline
:args: commandlien arguments
:returns: None
"""
##########################################################################
# Basic settings #
##########################################################################
exp_dir = 'experiments'
log_dir = os.path.join(exp_dir, 'logs')
model_dir = os.path.join(exp_dir, 'models')
os.makedirs(model_dir, exist_ok=True)
##########################################################################
# Define all the necessary variables for model training and evaluation #
##########################################################################
writer = SummaryWriter(log_dir)
train_dataset = dataset.NCovDataset('data/', stage='train')
weights = train_dataset.make_weights_for_balanced_classes()
weights = torch.DoubleTensor(weights)
sampler = torch.utils.data.sampler.WeightedRandomSampler(
weights, len(train_dataset.case_ids))
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=1, num_workers=20,
drop_last=False, sampler=sampler)
val_dataset = dataset.NCovDataset('data/', stage='val')
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=1, shuffle=False, num_workers=11,
drop_last=False)
cov_net = model.COVNet(n_classes=args.n_classes)
if torch.cuda.is_available():
cov_net = cov_net.cuda()
optimizer = optim.Adam(cov_net.parameters(), lr=args.lr, weight_decay=0.1)
if args.lr_scheduler == "plateau":
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, patience=3, factor=.3, threshold=1e-4, verbose=True)
elif args.lr_scheduler == "step":
scheduler = torch.optim.lr_scheduler.StepLR(
optimizer, step_size=3, gamma=args.gamma)
best_val_loss = float('inf')
best_val_accu = float(0)
iteration_change_loss = 0
t_start_training = time.time()
##########################################################################
# Main training loop #
##########################################################################
for epoch in range(args.epochs):
current_lr = get_lr(optimizer)
t_start = time.time()
############################################################
# The actual training and validation step for each epoch #
############################################################
train_loss, train_metric = train_model(
cov_net, train_loader, epoch, args.epochs, optimizer, writer,
current_lr, args.log_every)
with torch.no_grad():
val_loss, val_metric = evaluate_model(
cov_net, val_loader, epoch, args.epochs, writer, current_lr)
##############################
# Adjust the learning rate #
##############################
if args.lr_scheduler == 'plateau':
scheduler.step(val_loss)
elif args.lr_scheduler == 'step':
scheduler.step()
t_end = time.time()
delta = t_end - t_start
utils.print_epoch_progress(train_loss, val_loss, delta, train_metric,
val_metric)
iteration_change_loss += 1
print('-' * 30)
train_acc, val_acc = train_metric['accuracy'], val_metric['accuracy']
file_name = ('train_acc_{}_val_acc_{}_epoch_{}.pth'.
format(train_acc, val_acc, epoch))
torch.save(cov_net, os.path.join(model_dir, file_name))
if val_acc > best_val_accu:
best_val_accu = val_acc
if bool(args.save_model):
torch.save(cov_net, os.path.join(model_dir, 'best.pth'))
if val_loss < best_val_loss:
best_val_loss = val_loss
iteration_change_loss = 0
if iteration_change_loss == args.patience:
print(('Early stopping after {0} iterations without the decrease ' +
'of the val loss').format(iteration_change_loss))
break
t_end_training = time.time()
print('training took {}s'.
format(t_end_training - t_start_training))
if __name__ == "__main__":
args = config.parse_arguments()
main(args)
================================================
FILE: model.py
================================================
import torch
import torch.nn as nn
from torchvision import models
class COVNet(nn.Module):
def __init__(self, n_classes):
super().__init__()
model = models.resnet50(pretrained=True)
layer_list = list(model.children())[:-2]
self.pretrained_model = nn.Sequential(*layer_list)
self.pooling_layer = nn.AdaptiveAvgPool2d(1)
self.classifer = nn.Linear(2048, n_classes)
self.n_classes = n_classes
def forward(self, x):
x = torch.squeeze(x, dim=0)
features = self.pretrained_model(x)
pooled_features = self.pooling_layer(features)
pooled_features = pooled_features.view(pooled_features.size(0), -1)
flattened_features = torch.max(pooled_features, 0, keepdim=True)[0]
output = self.classifer(flattened_features)
return output
================================================
FILE: test.py
================================================
import os
from tqdm import tqdm
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch
import model
import config
import dataset
import utils
def evaluate_model(model, val_loader):
metric = torch.nn.CrossEntropyLoss()
model.eval()
for m in model.modules():
if isinstance(m, nn.BatchNorm2d):
m.train()
m.weight.requires_grad = False
m.bias.requires_grad = False
y_probs = np.zeros((0, 3), np.float)
losses, y_trues = [], []
for i, (image, label, case_id) in enumerate(tqdm(val_loader)):
if torch.cuda.is_available():
image = image.cuda()
label = label.cuda()
prediction = model.forward(image.float())
loss = metric(prediction, label.long())
loss_value = loss.item()
losses.append(loss_value)
y_prob = F.softmax(prediction, dim=1).detach().cpu().numpy()
y_probs = np.concatenate([y_probs, y_prob])
y_trues.append(label.item())
metric_collects = utils.calc_multi_cls_measures(y_probs, y_trues)
val_loss_epoch = np.mean(losses)
return val_loss_epoch, metric_collects
def main(args):
"""Main function for the testing pipeline
:args: commandline arguments
:returns: None
"""
##########################################################################
# Basic settings #
##########################################################################
exp_dir = 'experiments'
model_dir = os.path.join(exp_dir, 'models')
model_file = os.path.join(model_dir, 'best.pth')
val_dataset = dataset.NCovDataset('data/', stage='val')
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=1, shuffle=False, num_workers=11,
drop_last=False)
cov_net = model.COVNet(n_classes=args.n_classes)
if torch.cuda.is_available():
cov_net.cuda()
state = torch.load(model_file)
cov_net.load_state_dict(state.state_dict())
with torch.no_grad():
val_loss, metric_collects = evaluate_model(cov_net, val_loader)
prefix = '******Evaluate******'
utils.print_progress(mean_loss=val_loss, metric_collects=metric_collects,
prefix=prefix)
if __name__ == "__main__":
args = config.parse_arguments()
main(args)
================================================
FILE: utils.py
================================================
import numpy as np
import pandas as pd
from sklearn.metrics import accuracy_score, precision_score, recall_score
def calc_multi_cls_measures(probs, label):
"""Calculate multi-class classification measures (Accuracy, precision,
Recall, AUC.
:probs: NxC numpy array storing probabilities for each case
:label: ground truth label
:returns: a dictionary of accuracy, precision and recall
"""
n_classes = probs.shape[1]
preds = np.argmax(probs, axis=1)
accuracy = accuracy_score(label, preds)
precisions = precision_score(label, preds, average=None,
labels=range(n_classes), zero_division=0.)
recalls = recall_score(label, preds, average=None, labels=range(n_classes),
zero_division=0.)
metric_collects = {'accuracy': accuracy, 'precisions': precisions,
'recalls': recalls}
return metric_collects
def print_progress(epoch=None, n_epoch=None, n_iter=None, iters_one_batch=None,
mean_loss=None, cur_lr=None, metric_collects=None,
prefix=None):
"""Print the training progress.
:epoch: epoch number
:n_epoch: total number of epochs
:n_iter: current iteration number
:mean_loss: mean loss of current batch
:iters_one_batch: number of iterations per batch
:cur_lr: current learning rate
:metric_collects: dictionary returned by function calc_multi_cls_measures
:returns: None
"""
accuracy = metric_collects['accuracy']
precisions = metric_collects['precisions']
recalls = metric_collects['recalls']
n_classes = len(precisions)
log_str = ''
if epoch is not None:
log_str += 'Ep: {0}/{1}|'.format(epoch, n_epoch)
if n_iter is not None:
log_str += 'It: {0}/{1}|'.format(n_iter, iters_one_batch)
if mean_loss is not None:
log_str += 'Loss: {0:.4f}|'.format(mean_loss)
log_str += 'Acc: {:.4f}|'.format(accuracy)
templ = 'Pr: ' + ', '.join(['{:.4f}'] * (n_classes-1)) + '|'
log_str += templ.format(*(precisions[1:].tolist()))
templ = 'Re: ' + ', '.join(['{:.4f}'] * (n_classes-1)) + '|'
log_str += templ.format(*(recalls[1:].tolist()))
if cur_lr is not None:
log_str += 'lr: {0}'.format(cur_lr)
log_str = log_str if prefix is None else prefix + log_str
print(log_str)
def print_epoch_progress(train_loss, val_loss, time_duration, train_metric,
val_metric):
"""Print all the information after each epoch.
:train_loss: average training loss
:val_loss: average validation loss
:time_duration: time duration for current epoch
:train_metric_collects: a performance dictionary for training
:val_metric_collects: a performance dictionary for validation
:returns: None
"""
train_acc, val_acc = train_metric['accuracy'], val_metric['accuracy']
train_prec, val_prec = train_metric['precisions'], val_metric['precisions']
train_recalls, val_recalls = train_metric['recalls'], val_metric['recalls']
log_str = 'Train/Val| Loss: {:.4f}/{:.4f}|'.format(train_loss, val_loss)
log_str += 'Acc: {:.4f}/{:.4f}|'.format(train_acc, val_acc)
n_classes = len(train_prec)
templ = 'Pr: ' + ', '.join(['{:.4f}'] * (n_classes-1)) + '/'
log_str += templ.format(*(train_prec[1:].tolist()))
templ = ', '.join(['{:.4f}'] * (n_classes-1)) + '|'
log_str += templ.format(*(val_prec[1:].tolist()))
templ = 'Re: ' + ', '.join(['{:.4f}'] * (n_classes - 1)) + '/'
log_str += templ.format(*(train_recalls[1:].tolist()))
templ = ', '.join(['{:.4f}'] * (n_classes - 1)) + '|'
log_str += templ.format(*(val_recalls[1:].tolist()))
log_str += 'T(s) {:.2f}'.format(time_duration)
print(log_str)
gitextract_17ft28df/ ├── .gitignore ├── LICENSE ├── README.md ├── config.py ├── data/ │ ├── train.csv │ └── val.csv ├── dataset.py ├── main.py ├── model.py ├── test.py └── utils.py
SYMBOL INDEX (19 symbols across 6 files)
FILE: config.py
function parse_arguments (line 4) | def parse_arguments():
FILE: dataset.py
class NCovDataset (line 12) | class NCovDataset(data.Dataset):
method __init__ (line 13) | def __init__(self, root_dir, stage='train'):
method __len__ (line 33) | def __len__(self):
method __getitem__ (line 36) | def __getitem__(self, index):
method do_augmentation (line 79) | def do_augmentation(self, array, mask):
method make_weights_for_balanced_classes (line 112) | def make_weights_for_balanced_classes(self):
FILE: main.py
function get_lr (line 17) | def get_lr(optimizer):
function train_model (line 22) | def train_model(model, train_loader, epoch, num_epochs, optimizer, writer,
function evaluate_model (line 66) | def evaluate_model(model, val_loader, epoch, num_epochs, writer, current...
function main (line 111) | def main(args):
FILE: model.py
class COVNet (line 6) | class COVNet(nn.Module):
method __init__ (line 7) | def __init__(self, n_classes):
method forward (line 17) | def forward(self, x):
FILE: test.py
function evaluate_model (line 15) | def evaluate_model(model, val_loader):
function main (line 47) | def main(args):
FILE: utils.py
function calc_multi_cls_measures (line 6) | def calc_multi_cls_measures(probs, label):
function print_progress (line 28) | def print_progress(epoch=None, n_epoch=None, n_iter=None, iters_one_batc...
function print_epoch_progress (line 71) | def print_epoch_progress(train_loss, val_loss, time_duration, train_metric,
Condensed preview — 11 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (38K chars).
[
{
"path": ".gitignore",
"chars": 920,
"preview": "#cmake related files\nCMakeCache.txt\nCMakeFiles\nCMakeScripts\nMakefile\ncmake_install.cmake\ninstall_manifest.txt\nCTestTestf"
},
{
"path": "LICENSE",
"chars": 12356,
"preview": "“Commons Clause” License Condition v1.0\n\nThe Software is provided to you by the Licensor under the\nLicense, as defined b"
},
{
"path": "README.md",
"chars": 2087,
"preview": "# COVID-19 Detection Neural Network (COVNet)\nThis is a PyTorch implementation of the paper \"[Artificial Intelligence Dis"
},
{
"path": "config.py",
"chars": 1164,
"preview": "import argparse\n\n\ndef parse_arguments():\n \"\"\"Argument Parser for the commandline argments\n :returns: command line "
},
{
"path": "data/train.csv",
"chars": 41,
"preview": "case,label\ncaseid1,0\ncaseid2,1\ncaseid3,2\n"
},
{
"path": "data/val.csv",
"chars": 21,
"preview": "case,label\ncaseid4,0\n"
},
{
"path": "dataset.py",
"chars": 4550,
"preview": "import os\nimport pandas as pd\nimport numpy as np\n\nimport torch\nimport torch.utils.data as data\nimport SimpleITK as sitk\n"
},
{
"path": "main.py",
"chars": 7796,
"preview": "import os\nimport numpy as np\nimport time\n\nfrom tensorboardX import SummaryWriter\nimport torch\nimport torch.optim as opti"
},
{
"path": "model.py",
"chars": 841,
"preview": "import torch\nimport torch.nn as nn\nfrom torchvision import models\n\n\nclass COVNet(nn.Module):\n def __init__(self, n_cl"
},
{
"path": "test.py",
"chars": 2389,
"preview": "import os\nfrom tqdm import tqdm\nimport numpy as np\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\n"
},
{
"path": "utils.py",
"chars": 3785,
"preview": "import numpy as np\nimport pandas as pd\nfrom sklearn.metrics import accuracy_score, precision_score, recall_score\n\n\ndef c"
}
]
About this extraction
This page contains the full source code of the bkong999/COVNet GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 11 files (35.1 KB), approximately 8.2k tokens, and a symbol index with 19 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.