root_path=data_path,
train_transforms=train_transform_fn,
val_transforms=val_transform_fn,
batch_size=batch_size,
num_workers=num_workers,
val_batch_size=batch_size * 2,
limit_train_num_samples=250 if debug else None,
limit_val_num_samples=250 if debug else None,
random_seed=seed)
prepare_batch = prepare_batch_fp32
val_interval = 5
num_classes = 21
model = DeepLabV3(build_resnet18_backbone, num_classes=num_classes)
criterion = nn.CrossEntropyLoss()
lr = 0.007 / 4.0 * batch_size
weight_decay = 5e-4
momentum = 0.9
optimizer = optim.SGD(
[{
'params': model.backbone.parameters()
}, {
'params': chain(model.aspp.parameters(), model.decoder.parameters())
}],
lr=lr,
momentum=momentum,
from train_gf15.measure import SegmentationMetric
from tqdm import tqdm, trange
from train.LovaszSoftmax import lovasz_softmax
batch_size = 8
niter = 100
class_num = 15
learning_rate = 0.0001 * 1
beta1 = 0.5
cuda = True
num_workers = 1
size_h = 256
size_w = 256
flip = 0
band = 3
net = DeepLabV3(band, class_num)
train_path = '../dataset/GF15/train/'
val_path = '../dataset/GF15/val/'
test_path = '../dataset/GF15/test/'
out_file = './checkpoint/' + net.name
save_epoch = 1
test_step = 300
log_step = 1
num_GPU = 1
index = 2000
pre_trained = True
torch.cuda.set_device(0)
try:
import os
os.makedirs(out_file)
def main():
best_metric = 0
train_transform_det = trans.Compose([
trans.Scale(cfg.TRANSFROM_SCALES),
])
val_transform_det = trans.Compose([
trans.Scale(cfg.TRANSFROM_SCALES),
])
train_data = my_dataset.Dataset(cfg.TRAIN_DATA_PATH,
cfg.TRAIN_LABEL_PATH,
cfg.TRAIN_TXT_PATH,
'train',
transform=True,
transform_med=train_transform_det)
val_data = my_dataset.Dataset(cfg.VAL_DATA_PATH,
cfg.VAL_LABEL_PATH,
cfg.VAL_TXT_PATH,
'val',
transform=True,
transform_med=val_transform_det)
train_dataloader = DataLoader(train_data,
batch_size=cfg.BATCH_SIZE,
shuffle=True,
num_workers=1,
pin_memory=True)
val_dataloader = DataLoader(val_data,
batch_size=cfg.BATCH_SIZE,
shuffle=False,
num_workers=1,
pin_memory=True)
model = DeepLabV3(model_id=1, project_dir=cfg.BASE_PATH)
if cfg.RESUME:
checkpoint = torch.load(cfg.TRAINED_LAST_MODEL)
model.load_state_dict(checkpoint['state_dict'])
print('resume success \n')
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model = nn.DataParallel(model)
if torch.cuda.is_available():
model.cuda()
# if torch.cuda.is_available():
# model.cuda()
# params = [{'params': md.parameters()} for md in model.children() if md in [model.classifier]]
optimizer = optim.Adam(model.parameters(),
lr=cfg.INIT_LEARNING_RATE,
weight_decay=cfg.DECAY)
scheduler = ReduceLROnPlateau(optimizer,
mode='min',
factor=0.2,
patience=5,
verbose=True,
threshold=0.0001,
threshold_mode='rel',
cooldown=2,
eps=1e-08)
fl = FocalLoss2d(gamma=cfg.FOCAL_LOSS_GAMMA)
Loss_list = []
Accuracy_list = []
for epoch in range(cfg.EPOCH):
print('epoch {}'.format(epoch + 1))
#training--------------------------
train_loss = 0
train_acc = 0
for batch_idx, train_batch in enumerate(train_dataloader):
model.train()
batch_det_img, batch_y, _, _, _, _, _ = train_batch
batch_det_img, batch_y = Variable(batch_det_img).cuda(), Variable(
batch_y).cuda()
output = model(batch_det_img)
del batch_det_img
loss = calc_loss(output, batch_y)
# train_loss += loss.data[0]
#should change after
# pred = torch.max(out, 1)[0]
# train_correct = (pred == batch_y).sum()
# train_acc += train_correct.data[0]
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (batch_idx) % 5 == 0:
model.eval()
val_loss = 0
for v_batch_idx, val_batch in enumerate(val_dataloader):
v_batch_det_img, v_batch_y, _, _, _, _, _ = val_batch
v_batch_det_img, v_batch_y = Variable(
v_batch_det_img).cuda(), Variable(v_batch_y).cuda()
val_out = model(v_batch_det_img)
del v_batch_det_img
val_loss += float(calc_loss(val_out, v_batch_y))
scheduler.step(val_loss)
del val_out, v_batch_y
print("Train Loss: {:.6f} Val Loss: {:.10f}".format(
loss, val_loss))
if (epoch + 1) % 10 == 0:
torch.save({'state_dict': model.state_dict()},
os.path.join(
cfg.SAVE_MODEL_PATH, cfg.TRAIN_LOSS,
'model_tif_deeplab18_bce_240*240_' +
str(epoch + 1) + '.pth'))
torch.save({'state_dict': model.state_dict()},
os.path.join(cfg.SAVE_MODEL_PATH, cfg.TRAIN_LOSS,
'model_tif_deeplab18_bce_240*240_last.pth'))
def main():
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# initialize parameters
num_steps = args.num_steps
batch_size = args.batch_size
lr = args.lr
save_cp = args.save_cp
img_scale = args.scale
val_percent = args.val / 100
# data input
dataset = BasicDataset(IMG_DIRECTORY, MASK_DIRECTORY, img_scale)
n_val = int(len(dataset) * val_percent)
n_train = len(dataset) - n_val
train, val = random_split(dataset, [n_train, n_val])
tcga_dataset = UnlabeledDataset(TCGA_DIRECTORY)
n_unlabeled = len(tcga_dataset)
# create network
logger = logging.getLogger()
logger.setLevel(logging.INFO)
#logger.addHandler(logging.StreamHandler())
logging.info('Using device %s' % str(device))
logging.info('Network %s' % args.mod)
logging.info('''Starting training:
Num_steps: %.2f
Batch size: %.2f
Learning rate: %.4f_transform
Training size: %.0f
Validation size: %.0f
Unlabeled size: %.0f
Checkpoints: %s
Device: %s
Scale: %.2f
''' % (num_steps, batch_size, lr, n_train, n_val, n_unlabeled,
str(save_cp), str(device.type), img_scale))
if args.mod == 'unet':
net = UNet(n_channels=3, n_classes=NUM_CLASSES)
print('channels = %d , classes = %d' % (net.n_channels, net.n_classes))
elif args.mod == 'modified_unet':
net = modified_UNet(n_channels=3, n_classes=NUM_CLASSES)
print('channels = %d , classes = %d' % (net.n_channels, net.n_classes))
elif args.mod == 'deeplabv3':
net = DeepLabV3(nclass=NUM_CLASSES, pretrained_base=False)
print('channels = 3 , classes = %d' % net.nclass)
elif args.mod == 'deeplabv3plus':
net = DeepLabV3Plus(nclass=NUM_CLASSES, pretrained_base=False)
print('channels = 3 , classes = %d' % net.nclass)
elif args.mod == 'nestedunet':
net = NestedUNet(nclass=NUM_CLASSES, deep_supervision=False)
print('channels = 3 , classes = %d' % net.nlass)
elif args.mod == 'inception3':
net = Inception3(n_classes=4,
inception_blocks=None,
init_weights=True,
bilinear=True)
print('channels = 3 , classes = %d' % net.n_classes)
net.to(device=device)
net.train()
cudnn.benchmark = True
# init D
model_D = FCDiscriminator(num_classes=args.num_classes)
if args.restore_from_D is not None:
model_D.load_state_dict(torch.load(args.restore_from_D))
model_D.train()
model_D.cuda()
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
if args.semi_train is None:
train_loader = DataLoader(train,
batch_size=batch_size,
shuffle=True,
num_workers=8,
pin_memory=True)
val_loader = DataLoader(val,
batch_size=batch_size,
shuffle=False,
num_workers=8,
pin_memory=True)
else:
#read unlabeled data and labeled data
train_loader = DataLoader(train,
batch_size=batch_size,
shuffle=True,
num_workers=4,
pin_memory=True)
val_loader = DataLoader(val,
batch_size=batch_size,
shuffle=False,
num_workers=4,
pin_memory=True)
trainloader_remain = DataLoader(tcga_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=4,
pin_memory=True)
#trainloader_gt = data.DataLoader(train_gt_dataset,
#batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=3, pin_memory=True)
trainloader_remain_iter = enumerate(trainloader_remain)
trainloader_iter = enumerate(train_loader)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
#optimizer = optim.SGD(net.optim_parameters(args),
#lr=args.learning_rate, momentum=args.momentum,weight_decay=args.weight_decay)
optimizer = optim.Adam(net.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
optimizer.zero_grad()
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
10000,
eta_min=1e-6,
last_epoch=-1)
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
#optimizer_D = optim.SGD(model_D.parameters(), lr=args.learning_rate_D, momentum=args.momentum,weight_decay=args.weight_decay)
optimizer_D.zero_grad()
# loss/ bilinear upsampling
bce_loss = BCEWithLogitsLoss2d()
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
'''
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear', align_corners=True)
else:
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
'''
# labels for adversarial training
pred_label = 0
gt_label = 1
for i_iter in range(args.num_steps):
best_acc = 0
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_semi_value = 0
loss_semi_adv_value = 0
optimizer.zero_grad()
#adjust_learning_rate(optimizer, i_iter)
optimizer_D.zero_grad()
adjust_learning_rate_D(optimizer_D, i_iter)
for sub_i in range(args.iter_size):
# train G
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
for param in net.parameters():
param.requires_grad = True
# do semi first
if (args.lambda_semi > 0 or args.lambda_semi_adv > 0
) and i_iter >= args.semi_start_adv:
try:
_, batch = trainloader_remain_iter.__next__()
except:
trainloader_remain_iter = enumerate(trainloader_remain)
_, batch = trainloader_remain_iter.__next__()
# only access to img
images = batch['image']
images = images.type(torch.FloatTensor)
images = Variable(images).cuda()
pred = net(images)
pred_remain = pred.detach()
D_out = interp(model_D(F.softmax(pred, dim=1)))
D_out_sigmoid = torch.sigmoid(
D_out).data.cpu().numpy().squeeze(axis=1)
#D_out_sigmoid = torch.sigmoid(D_out).data.cpu().numpy()
#ignore_mask_remain = np.zeros(D_out_sigmoid.shape).astype(np.bool)
targetr = Variable(torch.ones(D_out.shape))
targetr = Variable(torch.FloatTensor(targetr)).cuda()
loss_semi_adv = args.lambda_semi_adv * bce_loss(D_out, targetr)
loss_semi_adv = loss_semi_adv / args.iter_size
#loss_semi_adv.backward()
#loss_semi_adv_value += loss_semi_adv.data.cpu().numpy()[0]/args.lambda_semi_adv
loss_semi_adv_value += loss_semi_adv.cpu().detach().numpy(
).item() / args.lambda_semi_adv
if args.lambda_semi <= 0 or i_iter < args.semi_start:
loss_semi_adv.backward()
loss_semi_value = 0
else:
# produce ignore mask
semi_ignore_mask = (D_out_sigmoid < args.mask_T)
semi_gt = pred.data.cpu().numpy().argmax(axis=1)
semi_gt[semi_ignore_mask] = 255
semi_ratio = 1.0 - float(
semi_ignore_mask.sum()) / semi_ignore_mask.size
print('semi ratio: {:.4f}'.format(semi_ratio))
if semi_ratio == 0.0:
loss_semi_value += 0
else:
semi_gt = torch.FloatTensor(semi_gt)
loss_semi = args.lambda_semi * loss_calc(pred, semi_gt)
loss_semi = loss_semi / args.iter_size
loss_semi_value += loss_semi.cpu().detach().numpy(
).item() / args.lambda_semi
loss_semi += loss_semi_adv
loss_semi.backward()
else:
loss_semi = None
loss_semi_adv = None
# train with source
try:
_, batch = trainloader_iter.__next__()
except:
trainloader_iter = enumerate(train_loader)
_, batch = trainloader_iter.__next__()
images = batch['image']
labels = batch['mask']
images = images.to(device=device, dtype=torch.float32)
labels = labels.to(device=device, dtype=torch.long)
labels = labels.squeeze(1)
ignore_mask = (labels.cpu().numpy() == 255)
#pred = interp(net(images))
pred = net(images)
criterion = nn.CrossEntropyLoss()
loss_seg = criterion(pred, labels)
#loss_seg = loss_calc(pred, labels)
D_out = interp(model_D(F.softmax(pred, dim=1)))
targetr = Variable(torch.ones(D_out.shape))
targetr = Variable(torch.FloatTensor(targetr)).cuda()
#loss_adv_pred = bce_loss(D_out, targetr)
if i_iter > args.semi_start_adv:
loss_adv_pred = bce_loss(D_out, targetr)
loss = loss_seg + args.lambda_adv_pred * loss_adv_pred
loss_adv_pred_value += loss_adv_pred.cpu().detach().numpy(
).item() / args.iter_size
else:
loss = loss_seg
# proper normalization
loss = loss / args.iter_size
loss.backward()
optimizer.step()
loss_seg_value += loss_seg.cpu().detach().numpy().item(
) / args.iter_size
#loss_adv_pred_value += loss_adv_pred.cpu().detach().numpy().item()/args.iter_size
# train D
# bring back requires_grad
if i_iter > args.semi_start_adv and i_iter % 3 == 0:
for param in net.parameters():
param.requires_grad = False
for param in model_D.parameters():
param.requires_grad = True
# train with pred
pred = pred.detach()
if args.D_remain:
pred = torch.cat((pred, pred_remain), 0)
#ignore_mask = np.concatenate((ignore_mask,ignore_mask_remain), axis = 0)
D_out = interp(model_D(F.softmax(pred, dim=1)))
#targetf = Variable(torch.zeros(D_out.shape))
targetf = 0.1 * np.random.rand(D_out.shape[0], D_out.shape[1],
D_out.shape[2], D_out.shape[3])
targetf = Variable(torch.FloatTensor(targetf)).cuda()
loss_D = bce_loss(D_out, targetf)
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
loss_D_value += loss_D.data.cpu().detach().numpy().item()
# train with gt
# get gt labels
try:
_, batch = trainloader_iter.__next__()
except:
trainloader_iter = enumerate(train_loader)
_, batch = trainloader_iter.__next__()
labels_gt = batch['mask']
D_gt_v = Variable(one_hot(labels_gt)).cuda()
ignore_mask_gt = (labels_gt.numpy() == 255).squeeze(axis=1)
D_out = interp(model_D(D_gt_v))
#targetr = Variable(torch.ones(D_out.shape))
targetr = 0.1 * np.random.rand(D_out.shape[0], D_out.shape[1],
D_out.shape[2],
D_out.shape[3]) + 0.9
targetr = Variable(torch.FloatTensor(targetr)).cuda()
loss_D = bce_loss(D_out, targetr)
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
optimizer_D.step()
loss_D_value += loss_D.cpu().detach().numpy().item()
scheduler.step()
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.3f}, loss_semi = {5:.3f}, loss_semi_adv = {6:.3f}'
.format(i_iter, args.num_steps, loss_seg_value,
loss_adv_pred_value, loss_D_value, loss_semi_value,
loss_semi_adv_value))
'''
if i_iter >= args.num_steps-1:
print 'save model ...'
torch.save(model.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+str(args.num_steps)+'.pth'))
torch.save(model_D.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+str(args.num_steps)+'_D.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter!=0:
print 'taking snapshot ...'
torch.save(model.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+str(i_iter)+'.pth'))
torch.save(model_D.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+str(i_iter)+'_D.pth'))
'''
# save checkpoints
if save_cp and (i_iter % 1000) == 0 and (i_iter != 0):
try:
os.mkdir(DIR_CHECKPOINTS)
logging.info('Created checkpoint directory')
except OSError:
pass
torch.save(net.state_dict(),
DIR_CHECKPOINTS + 'i_iter_%d.pth' % (i_iter + 1))
logging.info('Checkpoint %d saved !' % (i_iter + 1))
if (i_iter % 1000 == 0) and (i_iter != 0):
val_score, accuracy, dice_avr, dice_panck, dice_nuclei, dice_lcell = eval_net(
net, val_loader, device, n_val)
logging.info('Validation cross entropy: {}'.format(val_score))
if accuracy > best_acc:
best_acc = accuracy
result_file = open('result.txt', 'a', encoding='utf-8')
result_file.write('best_acc = ' + str(best_acc) + '\n' +
'iter = ' + str(i_iter) + '\n')
result_file.close
def train(train_loader, val_loader, class_weights, class_encoding):
print("\nTraining...\n")
num_classes = len(class_encoding)
# 初始化ENet
model = DeepLabV3(num_classes)
# model.load_state_dict(torch.load("save/model_13_2_2_2_epoch_580.pth"))
# model.aspp.conv_1x1_4 = torch.nn.Conv2d(256, num_classes, kernel_size=1)
# 检查网络结构是否正确
print(model)
# 交叉熵的损失函数
criterion = nn.CrossEntropyLoss(weight=class_weights)
# Adam as the optimizer
optimizer = optim.Adam(model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay)
# 学习率衰减
# lr_decay_epochs: 学习率衰减期。
# lr_decay: 学习率衰减的乘积因子,默认值:-0.1
lr_updater = lr_scheduler.StepLR(optimizer, args.lr_decay_epochs,
args.lr_decay)
# 评价指标
# if not args.ignore_unlabeled:
# ignore_index = list(class_encoding).index('unlabeled')
# else:
# ignore_index = None
ignore_index = None
metric = IoU(num_classes, ignore_index=ignore_index)
if use_cuda:
print("model使用GPU")
model = model.cuda()
criterion = criterion.cuda()
# Optionally 从checkpoint恢复
if args.resume:
model, optimizer, start_epoch, best_miou = utils.load_checkpoint(
model, optimizer, args.save_dir, args.name)
print("Resuming from model: Start epoch = {0} "
"| Best mean IoU = {1:.4f}".format(start_epoch, best_miou))
else:
start_epoch = 0
best_miou = 0
# 开始 Training
print()
train = Train(model, train_loader, optimizer, criterion, metric, use_cuda)
val = Test(model, val_loader, criterion, metric, use_cuda)
for epoch in range(start_epoch, args.epochs):
print(">>>> [Epoch: {0:d}] Training".format(epoch))
lr_updater.step() # 修改学习率，开始训练
epoch_loss, (iou, miou) = train.run_epoch(args.print_step)
# 打印epoch，loss,mean iou
print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(epoch, epoch_loss, miou))
# 如果当前的epochs结束，打印验证的进行一个验证
if (epoch + 1) % 10 == 0 or epoch + 1 == args.epochs:
print(">>>> [Epoch: {0:d}] Validation".format(epoch))
loss, (iou, miou) = val.run_epoch(args.print_step)
print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(epoch, loss, miou))
# Print per class IoU on last epoch or if best iou
if epoch + 1 == args.epochs or miou > best_miou:
for key, class_iou in zip(class_encoding.keys(), iou):
print("{0}: {1:.4f}".format(key, class_iou))
# Save the model if it's the best thus far
if miou > best_miou:
print("\nBest model thus far. Saving...\n")
best_miou = miou
utils.save_checkpoint(model, optimizer, epoch + 1, best_miou,
args)
return model
# 导入需要的数据
if args.dataset.lower() == 'lane':
from data_provider.Lane_dataset import Lane_dataset as dataset
loaders, w_class, class_encoding = load_dataset(dataset)
train_loader, val_loader, test_loader = loaders
if args.mode.lower() in {'train', 'full'}:
model = train(train_loader, val_loader, w_class, class_encoding)
if args.mode.lower() == 'full':
test(model, test_loader, w_class, class_encoding)
elif args.mode.lower() == 'test':
# 初始化新的 ENet model
num_classes = len(class_encoding)
model = DeepLabV3(num_classes)
# model.load_state_dict(torch.load("save/model_13_2_2_2_epoch_580.pth"))
# model.aspp.conv_1x1_4 = torch.nn.Conv2d(256, num_classes, kernel_size=1)
if use_cuda:
model = model.cuda()
# 初始化优化器
# checkpoint
optimizer = optim.Adam(model.parameters())
# 加载以前存储过的ENet模型
model = utils.load_checkpoint(model, optimizer, args.save_dir,
args.name)[0]
print(model)
test(model, test_loader, w_class, class_encoding)
else:
def prediction( weight):
best_metric = 0
train_transform_det = trans.Compose([
trans.Scale(cfg.TRANSFROM_SCALES),
])
val_transform_det = trans.Compose([
trans.Scale(cfg.TRANSFROM_SCALES),
])
test_transform_det = trans.Compose([
trans.Scale((960,960)),
])
model = DeepLabV3(model_id=1,project_dir=cfg.BASE_PATH)
# model=torch.nn.DataParallel(model)
if torch.cuda.is_available():
model.cuda()
# model.load_state_dict({k.replace('module.', ''): v for k, v in torch.load(weight).items()})
# model.load_state_dict(torch.load(weight))
checkpoint = torch.load(weight)
model.load_state_dict(checkpoint['state_dict'])
# test_data = my_dataset.Dataset(cfg.TEST_DATA_PATH, '',cfg.TEST_TXT_PATH, 'test', transform=True, transform_med=test_transform_det)
test_data = my_dataset.Dataset(cfg.VAL_DATA_PATH, cfg.VAL_LABEL_PATH,cfg.VAL_TXT_PATH, 'val', transform=True, transform_med=test_transform_det)
# test_data = my_dataset.Dataset(cfg.TRAIN_DATA_PATH, cfg.TRAIN_LABEL_PATH,cfg.TRAIN_TXT_PATH, 'train', transform=True, transform_med=test_transform_det)
test_dataloader = DataLoader(test_data, batch_size=cfg.TEST_BATCH_SIZE, shuffle=False, num_workers=8, pin_memory=True)
crop = 0
for batch_idx, val_batch in enumerate(test_dataloader):
model.eval()
#
# batch_x1, batch_x2, _, filename, h, w, green_mask1, green_mask2 = val_batch
batch_det_img, _, filename, h, w,_,green_mask2 = val_batch
# green_mask1 = green_mask1.view(output_w, output_h, -1).data.cpu().numpy()
filename = filename[0].split('/')[-1].replace('image','mask_2017')
if crop:
pass
# outputs = np.zeros((cfg.TEST_BATCH_SIZE,1,960, 960))
#
# while (i + w // rows <= w):
# j = 0
# while (j + h // cols <= h):
# batch_x1_ij = batch_x1[0, :, i:i + w // rows, j:j + h // cols]
# batch_x2_ij = batch_x2[0, :, i:i + w // rows, j:j + h // cols]
# # batch_y_ij = batch_y[batch_idx,: , i:i + w // rows, j:j + h // cols]
# batch_x1_ij = np.expand_dims(batch_x1_ij, axis=0)
# batch_x2_ij = np.expand_dims(batch_x2_ij, axis=0)
# batch_x1_ij, batch_x2_ij = Variable(torch.from_numpy(batch_x1_ij)).cuda(), Variable(
# torch.from_numpy(batch_x2_ij)).cuda()
# with torch.no_grad():
# output = model(batch_x1_ij, batch_x2_ij)
# output_w, output_h = output.shape[-2:]
# output = torch.sigmoid(output).view(-1, output_w, output_h)
#
# output = output.data.cpu().numpy() # .resize([80, 80, 1])
# output = np.where(output > cfg.THRESH, 255, 0)
# outputs[0, :, i:i + w // rows, j:j + h // cols] = output
#
# j += h // cols
# i += w // rows
#
#
# if not os.path.exists('./change'):
# os.mkdir('./change')
# print('./change/{}'.format(filename))
# cv2.imwrite('./change/crop_{}'.format(filename), outputs[0,0,:,:])
else:
batch_det_img = Variable(batch_det_img).cuda()
with torch.no_grad():
outputs = model(batch_det_img)
output_w, output_h = outputs[0].shape[-2:]
# green_mask2 = green_mask2.view(output_w, output_h, -1).data.cpu().numpy()
output = torch.sigmoid(outputs).view(output_w, output_h, -1).data.cpu().numpy()
# print(output.min(),output.max())
output = np.where((output > cfg.THRESH) , 255, 0)
if not os.path.exists('./change'):
os.mkdir('./change')
print('./change/{}'.format(filename))
cv2.imwrite('./change/{}'.format(filename), output)
def train(FLAGS):
# Defining the hyperparameters
device = FLAGS.cuda
batch_size = FLAGS.batch_size
epochs = FLAGS.epochs
lr = FLAGS.learning_rate
print_every = FLAGS.print_every
eval_every = FLAGS.eval_every
save_every = FLAGS.save_every
nc = FLAGS.num_classes
wd = FLAGS.weight_decay
ip = FLAGS.input_path_train
lp = FLAGS.label_path_train
ipv = FLAGS.input_path_val
lpv = FLAGS.label_path_val
print('[INFO]Defined all the hyperparameters successfully!')
# Get the class weights
print('[INFO]Starting to define the class weights...')
pipe = loader(ip, lp, batch_size='all')
class_weights = get_class_weights(pipe, nc)
print('[INFO]Fetched all class weights successfully!')
# Get an instance of the model
deeplabv3 = DeepLabV3(nc)
print('[INFO]Model Instantiated!')
# Move the model to cuda if available
deeplabv3 = deeplabv3.to(device)
# Define the criterion and the optimizer
criterion = nn.CrossEntropyLoss(
weight=torch.FloatTensor(class_weights).to(device))
optimizer = torch.optim.Adam(deeplabv3.parameters(),
lr=lr,
weight_decay=wd)
print('[INFO]Defined the loss function and the optimizer')
# Training Loop starts
print('[INFO]Staring Training...')
print()
train_losses = []
eval_losses = []
# Assuming we are using the CamVid Dataset
bc_train = 367 // batch_size
bc_eval = 101 // batch_size
pipe = loader(ip, lp, batch_size)
eval_pipe = loader(ipv, lpv, batch_size)
epochs = epochs
for e in range(1, epochs + 1):
train_loss = 0
print('-' * 15, 'Epoch %d' % e, '-' * 15)
deeplabv3.train()
for _ in tqdm(range(bc_train)):
X_batch, mask_batch = next(pipe)
#assert (X_batch >= 0. and X_batch <= 1.0).all()
X_batch, mask_batch = X_batch.to(device), mask_batch.to(device)
optimizer.zero_grad()
out = deeplabv3(X_batch.float())
loss = criterion(out, mask_batch.long())
loss.backward()
optimizer.step()
train_loss += loss.item()
print()
train_losses.append(train_loss)
if (e + 1) % print_every == 0:
print('Epoch {}/{}...'.format(e, epochs),
'Loss {:6f}'.format(train_loss))
if e % eval_every == 0:
with torch.no_grad():
deeplabv3.eval()
eval_loss = 0
for _ in tqdm(range(bc_eval)):
inputs, labels = next(eval_pipe)
inputs, labels = inputs.to(device), labels.to(device)
out = deeplabv3(inputs)
loss = criterion(out, labels.long())
eval_loss += loss.item()
print()
print('Loss {:6f}'.format(eval_loss))
eval_losses.append(eval_loss)
if e % save_every == 0:
checkpoint = {'epochs': e, 'state_dict': deeplabv3.state_dict()}
torch.save(checkpoint,
'./ckpt-deeplabv3-{}-{}.pth'.format(e, train_loss))
print('Model saved!')
print('Epoch {}/{}...'.format(e + 1, epochs),
'Total Mean Loss: {:6f}'.format(sum(train_losses) / epochs))
print('[INFO]Training Process complete!')
type=str,
default='./data/Training/',
help='Image range for test')
args = parser.parse_args()
DATA_FOLDER = args.data_folder
SAVE_TEST_IMAGES = False #True
SAVE_TEST_3D = True
TEST_SAVE_DIR = './results/'
print(args.test_img_range)
mrange = tuple(args.test_img_range)
# load the model
model = DeepLabV3(num_classes=len(CLASSES) + 1)
model.load_state_dict(torch.load(args.model_folder))
model.cuda()
model.eval()
def save_images(img,
subject_num,
slice_depth,
multiplier=1.0,
real_img=False,
tag=''):
print("SAVING IMAGES")
#for i in range(img.size()[0]):
new_arr = (img[:, :] * multiplier).cpu().byte()