def main():
opt = TestOptions()
args = opt.initialize()
if not os.path.exists(args.save):
os.makedirs(args.save)
model = CreateModel(args)
model.eval()
model.cuda()
targetloader = CreateTrgDataLoader(args)
for index, batch in enumerate(targetloader):
if index % 100 == 0:
print '%d processd' % index
image, _, name = batch
output = model(Variable(image).cuda())
output = nn.functional.upsample(output, (1024, 2048), mode='bilinear', align_corners=True).cpu().data[0].numpy()
output = output.transpose(1,2,0)
output_nomask = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output_nomask)
output_nomask = Image.fromarray(output_nomask)
name = name[0].split('/')[-1]
output_nomask.save('%s/%s' % (args.save, name))
output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0]))
compute_mIoU(args.gt_dir, args.save, args.devkit_dir, args.restore_from)
def main():
testing_entropy = 0
args = parse_args()
if not os.path.exists(args.save):
os.makedirs(args.save)
args.init_weights = root_base+'/snapshots/' + args.model_name
model = CreateModel(args)
model.eval()
model.cuda()
targetloader = CreateTrgDataLoader(args)
for index, batch in tqdm.tqdm(enumerate(targetloader)):
image, _, name,_ = batch
output = model(Variable(image).cuda())
output = nn.functional.softmax(output, dim=1)
testing_entropy += self_entropy(roll_axis(output.cpu().data[0].numpy()))
output = nn.functional.upsample(output, (1024, 2048), mode='bilinear', align_corners=True).cpu().data[0].numpy()
#output = np.multiply(output,priors)
output = output.transpose(1,2,0)
output_nomask = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output_nomask)
output_nomask = Image.fromarray(output_nomask)
name = name[0].split('/')[-1]
output_nomask.save('%s/%s' % (args.save, name))
output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0]))
mIou_ = compute_mIoU(args.gt_dir, args.save, args.devkit_dir, '')
def main():
opt = TestOptions()
args = opt.initialize()
if not os.path.exists(args.save):
os.makedirs(args.save)
maxiou = 0
root = args.restore_from
for i in range(1, 51):
#args.restore_from = './snapshots/gta2city_vgg/gta5_{0:d}'.format(i*2000)
args.restore_from = root + '_{0:d}'.format(i * 2000)
model = CreateModel(args)
print(args.restore_from)
model.eval()
model.cuda()
targetloader = CreateTrgDataLoader(args)
for index, batch in enumerate(targetloader):
if index % 100 == 0:
print '%d processd' % index
image, _, name = batch
output1, output2 = model(Variable(image).cuda())
output = output1 + output2
output = nn.functional.softmax(output, dim=1)
output = nn.functional.upsample(
output, (1024, 2048), mode='bilinear',
align_corners=True).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output_nomask = np.asarray(np.argmax(output, axis=2),
dtype=np.uint8)
output_col = colorize_mask(output_nomask)
output_nomask = Image.fromarray(output_nomask)
name = name[0].split('/')[-1]
output_nomask.save('%s/%s' % (args.save, name))
output_col.save('%s/%s_color.png' %
(args.save, name.split('.')[0]))
iou = compute_mIoU(args.gt_dir, args.save, args.devkit_dir,
args.restore_from)
#print(save_path)
if iou > maxiou:
maxiou = iou
best_model = args.restore_from
print('The best model is {:s}, the best IOU is {:f}'.format(
best_model, maxiou))
def main():
opt = TestOptions()
args = opt.initialize()
if args.target == 'deepglobe':
size = 612
elif args.target == 'deepglobe224':
size = 224
elif args.target == 'worldview' or args.target == 'worldviewFI' or args.target == 'worldviewGR' or args.target == 'worldviewFIc' or args.target == 'worldviewc':
size = 512
elif args.target == 'sentinel' or args.target == 'sentinelFI' or args.target == 'sentinelGR' or args.target == 'sentinelc':
size = 224
if args.target == 'pleiades' or args.target == 'pleiadesFI' or args.target == 'pleiadesGR':
size = 448
if not os.path.exists(args.save):
os.makedirs(args.save)
model = CreateModel(args)
model.eval()
model.cuda()
targetloader = CreateTrgDataLoader(args)
for index, batch in enumerate(targetloader):
if index % 100 == 0:
print('%d processd' % index)
image, _, name = batch
output = model(Variable(image).cuda())
output = nn.functional.softmax(output, dim=1)
output = nn.functional.upsample(
output, (size, size), mode='bilinear',
align_corners=True).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output_nomask = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output_nomask)
output_nomask = Image.fromarray(output_nomask)
name = name[0].split('/')[-1]
if args.target == 'deepglobe' or args.target == 'deepglobe224':
output_nomask.save('%s/%s.png' %
(args.save, os.path.splitext(name)[0]))
else:
output_nomask.save('%s/%s.png' % (args.save, name))
output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0]))
compute_mIoU(args.gt_dir, args.save, args.target, args.devkit_dir,
args.restore_from)
def main():
opt = TrainOptions()
args = opt.initialize()
_t = {'iter time': Timer()}
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
os.makedirs(os.path.join(args.snapshot_dir, 'logs'))
opt.print_options(args)
targetloader = CreateActTrgDataLoader(args, 'train')
targetloader_iter = iter(targetloader)
testloader = CreateActTrgDataLoader(args, 'test')
model, optimizer = CreateModel(args)
start_iter = 0
if args.restore_from is not None:
start_iter = int(args.restore_from.rsplit('/', 1)[1].rsplit('_')[1])
train_writer = tensorboardX.SummaryWriter(
os.path.join(args.snapshot_dir, "logs"))
bce_loss = torch.nn.BCEWithLogitsLoss()
cudnn.enabled = True
cudnn.benchmark = True
model.train()
model.cuda()
loss = ['loss_trg', 'eval_loss']
_t['iter time'].tic()
best_loss_eval = None
best_step = 0
eval_loss = np.array([0])
for i in range(start_iter, args.num_steps):
model.train()
model.module.adjust_learning_rate(args, optimizer, i)
optimizer.zero_grad()
try:
trg_img, trg_lbl, _, paths = next(targetloader_iter)
except StopIteration:
targetloader_iter = iter(targetloader)
trg_img, trg_lbl, _, paths = next(targetloader_iter)
trg_score, loss_trg = model(trg_img, lbl=trg_lbl)
loss_trg = loss_trg
loss_trg.mean().backward()
optimizer.step()
if (i + 1) % args.save_pred_every == 0:
with torch.no_grad():
model.eval()
eval_loss = 0
for test_img, test_lbl, _, _ in testloader:
test_score, loss_test = model(test_img, lbl=test_lbl)
eval_loss += loss_test.mean().item() * test_img.size(0)
eval_loss /= len(testloader.dataset)
if best_loss_eval == None or eval_loss < best_loss_eval:
best_loss_eval = eval_loss
best_step = i + 1
print('taking snapshot ... eval_loss: {}'.format(eval_loss))
torch.save(
model.module.state_dict(),
os.path.join(args.snapshot_dir,
str(i + 1) + '.pth'))
eval_loss = np.array([eval_loss])
if (i + 1) % args.print_freq == 0:
_t['iter time'].toc(average=False)
print('[it %d][src loss %.4f][lr %.4f][%.2fs]' % \
(i + 1, loss_trg.mean().data, optimizer.param_groups[0]['lr']*10000, _t['iter time'].diff))
if i + 1 > args.num_steps_stop:
print('finish training')
break
_t['iter time'].tic()
for m in loss:
train_writer.add_scalar(m, eval(m).mean(), i + 1)
best_maxSAD = maxSAD
best_avgSAD = avgSAD
torch.save(MODEL.state_dict(), PTH)
print(f'saving(best_maxSAD)... {PTH} ')
elif maxSAD == best_maxSAD:
if avgSAD < best_avgSAD:
best_maxSAD = maxSAD
best_avgSAD = avgSAD
torch.save(MODEL.state_dict(), PTH)
print(f'saving(best_avgSAD)... {PTH} ')
else:
print()
print('Finished Training')
MODEL.eval()
avgSAD = list()
maxSAD = list()
for idx, data in enumerate(data_iter):
X_, Y_ = data
Y_ = Y_.to('cuda')
Y_pred = MODEL(X_)
Y_pred += 0.5
Y_pred = torch.clamp(Y_pred, 0., 1.)
Y_pred = torch.round(255 * Y_pred)
Y_tmp = Y_
Y_tmp += 0.5
Y_tmp = torch.round(255 * Y_tmp)
def main():
opt = TestOptions()
args = opt.initialize()
os.environ["CUDA_VISIBLE_DEVICES"] = args.GPU
if not os.path.exists(args.save):
os.makedirs(args.save)
# 3 models are used because MBT method is used.
args.restore_from = args.restore_opt1
model1 = CreateModel(args)
model1.eval()
model1.cuda()
args.restore_from = args.restore_opt2
model2 = CreateModel(args)
model2.eval()
model2.cuda()
args.restore_from = args.restore_opt3
model3 = CreateModel(args)
model3.eval()
model3.cuda()
targetloader = CreateTrgDataLoader(args)
# change the mean for different dataset other than CS
IMG_MEAN = np.array((104.00698793, 116.66876762, 122.67891434),
dtype=np.float32)
IMG_MEAN = torch.reshape(torch.from_numpy(IMG_MEAN), (1, 3, 1, 1))
mean_img = torch.zeros(1, 1)
# ------------------------------------------------- #
# compute scores and save them
with torch.no_grad():
for index, batch in enumerate(targetloader):
if index % 100 == 0:
print('%d processd' % index)
image, _, name = batch # 1. get image
# create mean image
if mean_img.shape[-1] < 2:
B, C, H, W = image.shape
# 2. get mean image
mean_img = IMG_MEAN.repeat(B, 1, H, W)
image = image.clone() - mean_img # 3, image - mean_img
image = Variable(image).cuda()
# forward
output1 = model1(image)
output1 = nn.functional.softmax(output1, dim=1)
output2 = model2(image)
output2 = nn.functional.softmax(output2, dim=1)
output3 = model3(image)
output3 = nn.functional.softmax(output3, dim=1)
a, b = 0.3333, 0.3333
output = a * output1 + b * output2 + (1.0 - a - b) * output3
output = nn.functional.interpolate(
output, (1024, 2048), mode='bilinear',
align_corners=True).cpu().data[0].numpy()
#output = nn.functional.upsample( output, (1024, 2048), mode='bilinear', align_corners=True).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output_nomask = np.asarray(np.argmax(output, axis=2),
dtype=np.uint8)
output_col = colorize_mask(output_nomask)
output_nomask = Image.fromarray(output_nomask)
name = name[0].split('/')[-1]
output_nomask.save('%s/%s' % (args.save, name))
output_col.save('%s/%s_color.png' %
(args.save, name.split('.')[0]))
# scores computed and saved
# ------------------------------------------------- #
compute_mIoU(args.gt_dir, args.save, args.devkit_dir, args.restore_from)
def main():
opt = TrainOptions()
args = opt.initialize()
_t = {'iter time': Timer()}
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
os.makedirs(os.path.join(args.snapshot_dir, 'logs'))
opt.print_options(args)
sourceloader, targetloader = CreateActSrcDataLoader(
args), CreateActTrgDataLoader(args, 'train')
testloader = CreateActTrgDataLoader(args, 'test')
targetloader_iter, sourceloader_iter = iter(targetloader), iter(
sourceloader)
model, optimizer = CreateModel(args)
model_D, optimizer_D = CreateDiscriminator(args)
start_iter = 0
if args.restore_from is not None:
start_iter = int(args.restore_from.rsplit('/', 1)[1].rsplit('_')[1])
train_writer = tensorboardX.SummaryWriter(
os.path.join(args.snapshot_dir, "logs"))
bce_loss = torch.nn.BCEWithLogitsLoss()
cudnn.enabled = True
cudnn.benchmark = True
model.train()
model.cuda()
model_D.train()
model_D.cuda()
loss = [
'loss_src', 'loss_trg', 'loss_D_trg_fake', 'loss_D_src_real',
'loss_D_trg_real', 'eval_loss'
]
_t['iter time'].tic()
best_loss_eval = None
best_step = 0
eval_loss = np.array([0])
for i in range(start_iter, args.num_steps):
model.module.adjust_learning_rate(args, optimizer, i)
model_D.module.adjust_learning_rate(args, optimizer_D, i)
optimizer.zero_grad()
optimizer_D.zero_grad()
for param in model_D.parameters():
param.requires_grad = False
try:
src_img, src_lbl, _, _ = next(sourceloader_iter)
except StopIteration:
sourceloader_iter = iter(sourceloader)
src_img, src_lbl, _, _ = next(sourceloader_iter)
src_img, src_lbl = Variable(src_img).cuda(), Variable(
src_lbl.long()).cuda()
src_score, loss_src = model(src_img, lbl=src_lbl)
loss_src.mean().backward()
try:
trg_img, trg_lbl, _, _ = next(targetloader_iter)
except StopIteration:
targetloader_iter = iter(targetloader)
trg_img, trg_lbl, _, _ = next(targetloader_iter)
trg_img, trg_lbl = Variable(trg_img).cuda(), Variable(
trg_lbl.long()).cuda()
trg_score, loss_trg = model(trg_img, lbl=trg_lbl)
outD_trg, loss_D_trg_fake = model_D(F.softmax(trg_score, dim=1),
0) # do not apply softmax
loss_trg = args.lambda_adv_target * loss_D_trg_fake + loss_trg
loss_trg.mean().backward()
for param in model_D.parameters():
param.requires_grad = True
src_score, trg_score = src_score.detach(), trg_score.detach()
outD_src, model_D_loss = model_D(F.softmax(src_score, dim=1),
0) # do not apply softmax
loss_D_src_real = model_D_loss / 2
loss_D_src_real.mean().backward()
outD_trg, model_D_loss = model_D(F.softmax(trg_score, dim=1),
1) # do not apply softmax
loss_D_trg_real = model_D_loss / 2
loss_D_trg_real.mean().backward()
optimizer.step()
optimizer_D.step()
for m in loss:
train_writer.add_scalar(m, eval(m).mean(), i + 1)
if (i + 1) % args.save_pred_every == 0:
with torch.no_grad():
model.eval()
eval_loss = 0
for test_img, test_lbl, _, _ in testloader:
test_score, loss_test = model(test_img, lbl=test_lbl)
eval_loss += loss_test.mean().item() * test_img.size(0)
eval_loss /= len(testloader.dataset)
if best_loss_eval == None or eval_loss < best_loss_eval:
best_loss_eval = eval_loss
best_step = i + 1
print('taking snapshot ... eval_loss: {}'.format(eval_loss))
torch.save(
model.module.state_dict(),
os.path.join(args.snapshot_dir,
str(i + 1) + '.pth'))
eval_loss = np.array([eval_loss])
if (i + 1) % args.print_freq == 0:
_t['iter time'].toc(average=False)
print('[it %d][src loss %.4f][lr %.4f][%.2fs]' % \
(i + 1, loss_src.mean().data, optimizer.param_groups[0]['lr']*10000, _t['iter time'].diff))
if i + 1 > args.num_steps_stop:
print('finish training')
break
_t['iter time'].tic()
def main():
opt = TestOptions()
args = opt.initialize()
if not os.path.exists(args.save):
os.makedirs(args.save)
model = CreateModel(args)
model.eval()
model.cuda()
targetloader = CreateTrgDataLoader(args)
id_to_trainid = {2: 0, 1: 128, 0: 255}
for index, batch in enumerate(targetloader):
if index % 10 == 0:
print('%d processd' % index)
image, _, name = batch
_, output, _, _ = model(Variable(image).cuda()) #[1,3,129,129]
#import pdb;pdb.set_trace()
output = nn.functional.softmax(output, dim=1)
output = nn.functional.upsample(
output, (1634, 1634), mode='bilinear',
align_corners=True).cpu().data[0].numpy()
output = output.transpose(1, 2, 0) #(1634,1634,3)
'''
output_crop = output[14:526, 626:1138, 0:2]
np.save("/extracephonline/medai_data2/zhengdzhang/eyes/qikan/cai/output_crop_1.npy", output_crop)
#crop_img = Image.fromarray(output_crop)
#crop_img.save("/extracephonline/medai_data2/zhengdzhang/eyes/qikan/cai/crop_img.png")
'''
output_nomask = np.asarray(np.argmax(output, axis=2),
dtype=np.uint8) #(1644,1634) unique:[0,1,2]
cup_mask = get_bool(output_nomask, 0)
disc_mask = get_bool(output_nomask, 0) + get_bool(output_nomask, 1)
disc_mask = disc_mask.astype(np.uint8)
cup_mask = cup_mask.astype(np.uint8)
for i in range(5):
disc_mask = scipy.signal.medfilt2d(disc_mask, 19)
cup_mask = scipy.signal.medfilt2d(cup_mask, 19)
disc_mask = morphology.binary_erosion(disc_mask,
morphology.diamond(7)).astype(
np.uint8) # return 0,1
cup_mask = morphology.binary_erosion(cup_mask,
morphology.diamond(7)).astype(
np.uint8) # return 0,1
disc_mask = get_largest_fillhole(disc_mask)
cup_mask = get_largest_fillhole(cup_mask)
disc_mask = morphology.binary_dilation(disc_mask,
morphology.diamond(7)).astype(
np.uint8) # return 0,1
cup_mask = morphology.binary_dilation(cup_mask,
morphology.diamond(7)).astype(
np.uint8) # return 0,1
disc_mask = get_largest_fillhole(disc_mask).astype(
np.uint8) # return 0,1
cup_mask = get_largest_fillhole(cup_mask).astype(np.uint8)
output_nomask = disc_mask + cup_mask
output_col = np.ones(output_nomask.shape, dtype=np.float32)
for k, v in id_to_trainid.items():
output_col[output_nomask == k] = v
output_col = Image.fromarray(output_col.astype(np.uint8))
output_nomask = Image.fromarray(output_nomask)
name = name[0].split('.')[0] + '.png'
output_nomask.save('%s/%s' % (args.save, name))
output_col.save('%s/color_%s' % (args.save, name))
disc_dice, cup_dice = calculate_dice(args.gt_dir, args.save,
args.devkit_dir)
print('===> disc_dice:' + str(round(disc_dice, 3)) + '\t' + 'cup_dice:' +
str(round(cup_dice, 3)))
