def inference_module(_data_path, _save_path, _pth_path):
model = Network(channel=32, n_class=1)
model.load_state_dict(torch.load(_pth_path))
model.cuda()
model.eval()
os.makedirs(_save_path, exist_ok=True)
# FIXME
image_root = '{}/'.format(_data_path)
# gt_root = '{}/mask/'.format(data_path)
test_loader = test_dataset(image_root, image_root, 352)
for i in range(test_loader.size):
image, name = test_loader.load_data()
#gt = np.asarray(gt, np.float32)
#gt /= (gt.max() + 1e-8)
image = image.cuda()
lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2, lateral_edge = model(
image)
res = lateral_map_2 # final segmentation
#res = F.upsample(res, size=gt.shape, mode='bilinear', align_corners=False)
res = res.sigmoid().data.cpu().numpy().squeeze()
res = (res - res.min()) / (res.max() - res.min() + 1e-8)
misc.imsave(_save_path + '/' + name, res)
def inference():
parser = argparse.ArgumentParser()
parser.add_argument('--testsize', type=int, default=352, help='testing size')
parser.add_argument('--data_path', type=str, default='./Dataset/TestingSet/LungSegmentation/',
help='Path to test data')
parser.add_argument('--pth_path', type=str, default='./snapshots/save_weights/Semi-Inf-Net/Inf-Net-100.pth',
help='Path to weights file. If `semi-sup`, edit it to `Semi-Inf-Net/Semi-Inf-Net-100.pth`')
parser.add_argument('--save_path', type=str, default='./Results/LungSegmentation/Inf-Net/',
help='Path to save the predictions. if `semi-sup`, edit it to `Semi-Inf-Net`')
parser.add_argument('--backbone', type=str, default='Res2Net50',
help='change different backbone, choice: VGGNet16, ResNet50, Res2Net50')
opt = parser.parse_args()
print("#" * 20, "\nStart Testing (Inf-Net)\n{}\nThis code is written for 'Inf-Net: Automatic COVID-19 Lung "
"Infection Segmentation from CT Scans', 2020, arXiv.\n"
"----\nPlease cite the paper if you use this code and dataset. "
"And any questions feel free to contact me "
"via E-mail ([email protected])\n----\n".format(opt), "#" * 20)
if opt.backbone == 'Res2Net50':
print('Backbone loading: Res2Net50')
from Code.model_lung_infection.InfNet_Res2Net import Inf_Net as Network
elif opt.backbone == 'ResNet50':
print('Backbone loading: ResNet50')
from Code.model_lung_infection.InfNet_ResNet import Inf_Net as Network
elif opt.backbone == 'VGGNet16':
print('Backbone loading: VGGNet16')
from Code.model_lung_infection.InfNet_VGGNet import Inf_Net as Network
else:
raise ValueError('Invalid backbone parameters: {}'.format(opt.backbone))
model = Network()
# model = torch.nn.DataParallel(model, device_ids=[0, 1]) # uncomment it if you have multiply GPUs.
model.load_state_dict(torch.load(opt.pth_path, map_location={'cuda:1':'cuda:0'}))
model.cuda()
model.eval()
image_root = '{}/Imgs/'.format(opt.data_path)
# gt_root = '{}/GT/'.format(opt.data_path)
test_loader = test_dataset(image_root, opt.testsize)
os.makedirs(opt.save_path, exist_ok=True)
for i in range(test_loader.size):
image, name = test_loader.load_data()
image = image.cuda()
lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2, lateral_edge = model(image)
res = lateral_map_2
#res = F.upsample(res, size=(ori_size[1],ori_size[0]), mode='bilinear', align_corners=False)
res = res.sigmoid().data.cpu().numpy().squeeze()
res = (res - res.min()) / (res.max() - res.min() + 1e-8)
imsave(opt.save_path + name, res)
print('Test Done!')
def train_module(_opt):
#def train_module(_train_path, _train_save, _resume_snapshot,_batchsize):
#parser = argparse.ArgumentParser()
#parser.add_argument('--epoch', type=int, default=10, help='epoch number')
#parser.add_argument('--lr', type=float, default=3e-4, help='learning rate')
#parser.add_argument('--batchsize', type=int, default=_batchsize, help='training batch size')
#parser.add_argument('--trainsize', type=int, default=352, help='training dataset size')
#parser.add_argument('--clip', type=float, default=0.5, help='gradient clipping margin')
#parser.add_argument('--decay_rate', type=float, default=0.1, help='decay rate of learning rate')
#parser.add_argument('--decay_epoch', type=int, default=50, help='every n epochs decay learning rate')
#parser.add_argument('--train_path', type=str, default=_train_path)
#parser.add_argument('--train_save', type=str, default=_train_save)
#parser.add_argument('--resume_snapshot', type=str, default=_resume_snapshot)
#opt = parser.parse_args()
opt = _opt
# ---- build models ----
torch.cuda.set_device(0)
model = Network(channel=32, n_class=1).cuda()
model.load_state_dict(torch.load(opt.resume_snapshot))
params = model.parameters()
optimizer = torch.optim.Adam(params, opt.lr)
image_root = '{}/Imgs/'.format(opt.train_path)
gt_root = '{}/GT/'.format(opt.train_path)
edge_root = '{}/Edge/'.format(opt.train_path)
train_loader = get_loader(image_root,
gt_root,
edge_root,
batchsize=opt.batchsize,
trainsize=opt.trainsize)
total_step = len(train_loader)
print("#" * 20, "Start Training", "#" * 20)
for epoch in range(1, opt.epoch):
adjust_lr(optimizer, opt.lr, epoch, opt.decay_rate, opt.decay_epoch)
trainer(train_loader=train_loader,
model=model,
optimizer=optimizer,
epoch=epoch,
opt=opt,
total_step=total_step)
def inference_module(_data_path, _save_path, _pth_path):
model = Network(channel=32, n_class=1)
model.load_state_dict(torch.load(_pth_path))
model.cuda()
model.eval()
os.makedirs(_save_path, exist_ok=True)
image_root = '{}/'.format(_data_path)
# gt_root = '{}/GT/'.format(_data_path)
test_loader = test_dataset(image_root, 352)
# initially test_loader = test_dataset(image_root, image_root, 352)
# but test_dataset class from Code.utils.dataloader_LungInf have folloving:
# def __init__(self, image_root, testsize) ...
# because self doesn't have to be in input, one image_root is need to be removed.
for i in range(test_loader.size):
# initially image, gt, name = test_loader.load_data()
# but load_data() function of class test_dataset have following: return image, name
# this error have makes no possible to find shape from gt, so we find shape from image.
# this is made truth sh = image.shape[2:4]
image, name = test_loader.load_data()
image = image.cuda()
lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2, lateral_edge = model(
image)
sh = image.shape[2:4]
res = lateral_map_2 # final segmentation
res = F.upsample(res, size=sh, mode='bilinear', align_corners=False)
res = res.sigmoid().data.cpu().numpy().squeeze()
res = (res - res.min()) / (res.max() - res.min() + 1e-8)
# GT map is binary image, so we need threshold. 0.5 was chosen and because some resulting images need
# lower threshold, and some need higher, 0.5 good enough.
res = np.where(res < np.max(res) / 2, True, False)
imsave(_save_path + '/' + name, res, cmap='binary')
def inference():
parser = argparse.ArgumentParser()
parser.add_argument('--testsize', type=int, default=352, help='testing size')
parser.add_argument('--data_path', type=str, default='./Dataset/TestingSet/LungSegmentation/',
help='Path to test data')
parser.add_argument('--pth_path', type=str, default='./Snapshots/save_weights/Inf-Net/Inf-Net-100.pth',
help='Path to weights file. If `semi-sup`, edit it to `Semi-Inf-Net/Semi-Inf-Net-100.pth`')
parser.add_argument('--save_path', type=str, default='./Results/LungSegmentation/Inf-Net/',
help='Path to save the predictions. if `semi-sup`, edit it to `Semi-Inf-Net`')
opt = parser.parse_args()
print("#" * 20, "\nStart Testing (Inf-Net)\n{}\n".format(opt), "#" * 20)
model = Network()
model.load_state_dict(torch.load(opt.pth_path, map_location={'cuda:1':'cuda:0'}))
model.cuda()
model.eval()
image_root = '{}/Imgs/'.format(opt.data_path)
test_loader = test_dataset(image_root, opt.testsize)
os.makedirs(opt.save_path, exist_ok=True)
for i in range(test_loader.size):
image, name = test_loader.load_data()
image = image.cuda()
lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2, lateral_edge = model(image)
res = lateral_map_2
res = res.sigmoid().data.cpu().numpy().squeeze()
res = (res - res.min()) / (res.max() - res.min() + 1e-8)
res = np.where(res < np.max(res) / 2, True, False)
res = np.invert(res)
imsave(opt.save_path + name, res, cmap = 'gray')
print('Test Done!')
for batch_index, batch_samples in enumerate(test_loader):
data, target = batch_samples['img'].to(device), batch_samples['label'].to(device)
output = model(data)
test_loss += criteria(output, target.long())
score = F.softmax(output, dim=1)
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.long().view_as(pred)).sum().item()
targetcpu = target.long().cpu().numpy()
predlist = np.append(predlist, pred.cpu().numpy())
scorelist = np.append(scorelist, score.cpu().numpy()[:, 1])
targetlist = np.append(targetlist, targetcpu)
return targetlist, scorelist, predlist
seg_model = Network()
seg_model.load_state_dict(torch.load('pretrained/Semi-Inf-Net-100.pth', map_location={'cuda:1':'cuda:0'}))
cls_model = Densenet.densenet169()
pretrained_net = torch.load('/home/gcy/projects/COVID-CT-master/baseline methods/Self-Trans/Self-Trans.pt')
cls_model.load_state_dict(pretrained_net)
class wh_net(nn.Module):
def __init__(self, seg_net, cls_net):
super().__init__()
self.seg_net = seg_net
self.cls_net = cls_net
self.con_c = torch.nn.Conv2d(in_channels=4, out_channels=3, kernel_size=1, stride=1, padding=0)
def forward(self, x):
lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2, lateral_edge = self.seg_net(x)
def inference():
parser = argparse.ArgumentParser()
parser.add_argument('--testsize',
type=int,
default=352,
help='testing size')
parser.add_argument(
'--data_path',
type=str,
default='./Dataset/TrainingSet/LungSegmentation/Montgomery',
help='Path to test data')
parser.add_argument(
'--pth_path',
type=str,
default='./snapshots/save_weights/Inf-Net/Inf-Net-100.pth',
help=
'Path to weights file. If `semi-sup`, edit it to `Semi-Inf-Net/Semi-Inf-Net-100.pth`'
)
parser.add_argument(
'--save_path',
type=str,
default='./Results/LungSegmentation/IOUandDICE',
help=
'Path to save the predictions. if `semi-sup`, edit it to `Semi-Inf-Net`'
)
opt = parser.parse_args()
print("#" * 20, "\nStart Testing (Inf-Net)\n{}\n".format(opt), "#" * 20)
model = Network()
model.load_state_dict(
torch.load(opt.pth_path, map_location={'cuda:1': 'cuda:0'}))
model.cuda()
model.eval()
image_root = '{}/Imgs/'.format(opt.data_path)
gt_root = '{}/GT/'.format(opt.data_path)
test_loader = dice_test_dataset(image_root, gt_root, opt.testsize)
os.makedirs(opt.save_path, exist_ok=True)
iouCoef = []
diceCoef = []
metrics = []
for i in range(test_loader.size):
image, gt, name = test_loader.load_data()
image = image.cuda()
lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2, lateral_edge = model(
image)
res = lateral_map_2
res = res.sigmoid().data.cpu().numpy().squeeze()
res = (res - res.min()) / (res.max() - res.min() + 1e-8)
# change Tensor image to numpy image
gt = gt.numpy()
gt = gt.reshape(gt.shape[1:])
gt = gt.astype(bool)
gt = gt[0, :, :]
res = np.where(res < np.max(res) / 2, True, False)
res = np.invert(res)
imsave(opt.save_path + name, res, cmap='gray')
iouC = iou(res, gt)
diceC = dice(res, gt)
iouCoef.append(iouC)
diceCoef.append(diceC)
result = {'file': str(name), 'iouC': iouC, 'dice': diceC}
metrics.append(result)
print("IOU mean = ", np.mean(iouCoef), "\nDice maen = ", np.mean(diceCoef))
print("IOU max = ", np.max(iouCoef), "\nDice max = ", np.max(diceCoef))
print("IOU min = ", np.min(iouCoef), "\nDice min = ", np.min(diceCoef))
output_file = open(opt.save_path + 'metrics.txt', 'w', encoding='utf-8')
output_file.write("Name, IOU, Dice\n")
for item in metrics:
for d in item.items():
output_file.write(str(d) + " ")
output_file.write("\n")
output_file.write("\nIOU mean = " + str(np.mean(iouCoef)))
output_file.write("\nDice mean = " + str(np.mean(diceCoef)))
output_file.write("\nIOU max = " + str(np.max(iouCoef)))
output_file.write("\nDice max = " + str(np.max(diceCoef)))
output_file.write("\nIOU min = " + str(np.min(iouCoef)))
output_file.write("\nDice min = " + str(np.min(diceCoef)))
def inference():
parser = argparse.ArgumentParser()
parser.add_argument('--testsize',
type=int,
default=352,
help='testing size')
parser.add_argument('--data_path',
type=str,
default='/Users/vignavramesh/Documents/CT_Scans/',
help='Path to test data')
parser.add_argument(
'--pth_path',
type=str,
default='./Snapshots/save_weights/Inf-Net/Inf-Net-100.pth',
help=
'Path to weights file. If `semi-sup`, edit it to `Semi-Inf-Net/Semi-Inf-Net-100.pth`'
)
parser.add_argument(
'--save_path',
type=str,
default='/Users/vignavramesh/Documents/CT_Masks/',
help=
'Path to save the predictions. if `semi-sup`, edit it to `Semi-Inf-Net`'
)
opt = parser.parse_args()
print(
"#" * 20,
"\nStart Testing (Inf-Net)\n{}\nThis code is written for 'Inf-Net: Automatic COVID-19 Lung "
"Infection Segmentation from CT Scans', 2020, arXiv.\n"
"----\nPlease cite the paper if you use this code and dataset. "
"And any questions feel free to contact me "
"via E-mail ([email protected])\n----\n".format(opt), "#" * 20)
model = Network()
model.load_state_dict(
torch.load(opt.pth_path, map_location=torch.device('cpu')))
model.eval()
count = 0
list = glob.glob(opt.data_path + '*')
list = sort_list(list)
for image_root in list: #subs
count = image_root[(image_root.rindex('Volume') + 6):]
image_root += '/'
test_loader = test_dataset(image_root, opt.testsize)
os.makedirs(opt.save_path, exist_ok=True)
for i in range(test_loader.size):
image, name = test_loader.load_data()
lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2, lateral_edge = model(
image)
res = lateral_map_2
res = res.sigmoid().data.cpu().numpy().squeeze()
res = (res - res.min()) / (res.max() - res.min() + 1e-8)
string = opt.save_path + 'Volume' + str(count)
if not os.path.exists(string):
os.makedirs(string)
imageio.imwrite(string + '/' + name, res)
print('Test Done!')
def inference(image_root):
parser = argparse.ArgumentParser()
parser.add_argument('--testsize',
type=int,
default=352,
help='testing size')
parser.add_argument('--data_path',
type=str,
default='./Dataset/TestingSet/LungInfection-Test/',
help='Path to test data')
parser.add_argument(
'--pth_path',
type=str,
default='./Snapshots/save_weights/Inf-Net/Inf-Net-100.pth',
help=
'Path to weights file. If `semi-sup`, edit it to `Inf-Net/Inf-Net-100.pth`'
)
parser.add_argument(
'--save_path',
type=str,
default='./Results/Lung infection segmentation/Inf-Net/Labels/',
help='Path to save the predictions. if `semi-sup`, edit it to `Inf-Net`'
)
parser.add_argument(
'--save_path_append',
type=str,
default='./Results/Lung infection segmentation/Inf-Net/Append_result/',
help='Path to save the predictions. if `semi-sup`, edit it to `Inf-Net`'
)
parser.add_argument(
'--save_path_mask',
type=str,
default='./Results/Lung infection segmentation/Inf-Net/Mask/',
help='Path to save the predictions. if `semi-sup`, edit it to `Inf-Net`'
)
opt = parser.parse_args()
model = Network()
# model = torch.nn.DataParallel(model, device_ids=[0, 1]) # uncomment it if you have multiply GPUs.
model.load_state_dict(
torch.load(opt.pth_path, map_location={'cuda:1': 'cuda:0'}))
model.cuda()
model.eval()
# gt_root = '{}/GT/'.format(opt.data_path)
test_loader = test_dataset(image_root, opt.testsize)
os.makedirs(opt.save_path, exist_ok=True)
image, name = test_loader.load_data()
image = image.cuda()
lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2, lateral_edge = model(
image)
res = lateral_map_2
# res = F.upsample(res, size=(ori_size[1],ori_size[0]), mode='bilinear', align_corners=False)
res = res.sigmoid().data.cpu().numpy().squeeze()
res = (res - res.min()) / (res.max() - res.min() + 1e-8)
print(name)
misc.imsave(opt.save_path + name, res)
icon = show_label(res, opt.save_path_mask + name)
img = Image.open(image_root)
img = img.convert("RGBA")
img_w, img_h = img.size
icon = icon.resize((img_w, img_h), Image.ANTIALIAS)
layer = Image.new("RGBA", img.size, (0, 0, 0, 0))
layer.paste(icon, (0, 0))
out = Image.composite(layer, img, layer)
out.save(opt.save_path_append + name)
print('Test Done!')
def inference():
parser = argparse.ArgumentParser()
parser.add_argument('--testsize',
type=int,
default=352,
help='testing size')
parser.add_argument('--data_path',
type=str,
default='./Dataset/TestingSet/LungInfection-Test/',
help='Path to test data')
parser.add_argument(
'--pth_path',
type=str,
default='./Snapshots/save_weights/Semi-Inf-Net/Semi-Inf-Net-100.pth',
help=
'Path to weights fileif `semi-sup`, edit it to `Semi-Inf-Net/Semi-Inf-Net-100.pth`'
)
parser.add_argument(
'--save_path',
type=str,
default='./Results/Lung infection segmentation/Semi-Inf-Net/',
help=
'Path to save the predictions. if `semi-sup`, edit it to `Semi-Inf-Net`'
)
opt = parser.parse_args()
print(
"#" * 20,
"\nStart Testing (Inf-Net)\n{}\nThis code is written for 'Inf-Net: Automatic COVID-19 Lung "
"Infection Segmentation from CT Scans', 2020, arXiv.\n"
"----\nPlease cite the paper if you use this code and dataset. "
"And any questions feel free to contact me "
"via E-mail ([email protected])\n----\n".format(opt.backbone,
opt), "#" * 20)
model = Network()
# model = torch.nn.DataParallel(model, device_ids=[0, 1]) # uncomment it if you have multiply GPUs.
model.load_state_dict(torch.load(opt.pth_path))
model.cuda()
model.eval()
image_root = '{}/Imgs/'.format(opt.data_path)
gt_root = '{}/GT/'.format(opt.data_path)
test_loader = test_dataset(image_root, gt_root, opt.testsize)
os.makedirs(opt.save_path, exist_ok=True)
for i in range(test_loader.size):
image, gt, name = test_loader.load_data()
gt = np.asarray(gt, np.float32)
gt /= (gt.max() + 1e-8)
image = image.cuda()
lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2, lateral_edge = model(
image)
res = lateral_map_2
res = F.upsample(res,
size=gt.shape,
mode='bilinear',
align_corners=False)
res = res.sigmoid().data.cpu().numpy().squeeze()
res = (res - res.min()) / (res.max() - res.min() + 1e-8)
misc.imsave(opt.save_path + name, res)
print('Test Done!')
edge_root = '{}/Edge/'.format(opt.train_path)
models = []
fold_error = []
# ---- start !! -----
print("#" * 20, "\nStart Training (Inf-Net-{})\n{}\nThis code is written for 'Inf-Net: Automatic COVID-19 Lung "
"Infection Segmentation from CT Scans', 2020, arXiv.\n"
"----\nPlease cite the paper if you use this code and dataset. "
"And any questions feel free to contact me "
"via E-mail ([email protected])\n----\n".format(opt.backbone, opt), "#" * 20)
for k in range(opt.k_size):
model = Network(channel=opt.net_channel, n_class=opt.n_classes).cuda()
# ---- load pre-trained weights (mode=Semi-Inf-Net) ----
# - See Sec.2.3 of `README.md` to learn how to generate your own img/pseudo-label from scratch.
if opt.is_semi and opt.backbone == 'Res2Net50':
print('Loading weights from weights file trained on pseudo label')
model.load_state_dict(torch.load('./snapshots/save_weights/Inf-Net_Pseudo/Inf-Net-100.pth'))
else:
print('Not loading weights from weights file')
# ---- calculate FLOPs and Params ----
if opt.is_thop:
from Code.utils.utils import CalParams
x = torch.randn(1, 3, opt.trainsize, opt.trainsize).cuda()
CalParams(model, x)
