def test(cfg, areaname):
# Setup seeds
torch.manual_seed(cfg.get('seed', 1337))
torch.cuda.manual_seed(cfg.get('seed', 1337))
np.random.seed(cfg.get('seed', 1337))
random.seed(cfg.get('seed', 1337))
# Setup device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Setup Augmentations
augmentations = cfg['training'].get('augmentations', None)
data_aug = get_composed_augmentations(augmentations)
# Setup Dataloader
# data_loader = get_loader(cfg['data']['dataset'])
# data_path = cfg['data']['path']
#
# t_loader = data_loader(
# data_path,
# is_transform=True,
# split=cfg['data']['train_split'],
# img_size=(cfg['data']['img_rows'], cfg['data']['img_cols']),
# augmentations=data_aug)
#
# v_loader = data_loader(
# data_path,
# is_transform=True,
# split=cfg['data']['val_split'],
# img_size=(cfg['data']['img_rows'], cfg['data']['img_cols']),)
#
# n_classes = t_loader.n_classes
# trainloader = data.DataLoader(t_loader,
# batch_size=cfg['training']['batch_size'],
# num_workers=cfg['training']['n_workers'],
# shuffle=True)
#
# valloader = data.DataLoader(v_loader,
# batch_size=cfg['training']['batch_size'],
# num_workers=cfg['training']['n_workers'])
datapath = '/home/chengjjang/Projects/deepres/SatelliteData/{}/'.format(
areaname)
paths = {
'masks': '{}/patch{}_train/gt'.format(datapath, areaname),
'images': '{}/patch{}_train/rgb'.format(datapath, areaname),
'nirs': '{}/patch{}_train/nir'.format(datapath, areaname),
'swirs': '{}/patch{}_train/swir'.format(datapath, areaname),
'vhs': '{}/patch{}_train/vh'.format(datapath, areaname),
'vvs': '{}/patch{}_train/vv'.format(datapath, areaname),
'redes': '{}/patch{}_train/rede'.format(datapath, areaname),
'ndvis': '{}/patch{}_train/ndvi'.format(datapath, areaname),
}
valpaths = {
'masks': '{}/patch{}_val/gt'.format(datapath, areaname),
'images': '{}/patch{}_val/rgb'.format(datapath, areaname),
'nirs': '{}/patch{}_val/nir'.format(datapath, areaname),
'swirs': '{}/patch{}_val/swir'.format(datapath, areaname),
'vhs': '{}/patch{}_val/vh'.format(datapath, areaname),
'vvs': '{}/patch{}_val/vv'.format(datapath, areaname),
'redes': '{}/patch{}_val/rede'.format(datapath, areaname),
'ndvis': '{}/patch{}_val/ndvi'.format(datapath, areaname),
}
n_classes = 3
train_img_paths = [
pth for pth in os.listdir(paths['images'])
if ('_01_' not in pth) and ('_25_' not in pth)
]
val_img_paths = [
pth for pth in os.listdir(valpaths['images'])
if ('_01_' not in pth) and ('_25_' not in pth)
]
ntrain = len(train_img_paths)
nval = len(val_img_paths)
train_idx = [i for i in range(ntrain)]
val_idx = [i for i in range(nval)]
train_idx = [i for i in range(ntrain)]
val_idx = [i for i in range(nval)]
trainds = ImageProvider(MultibandImageType, paths, image_suffix='.png')
valds = ImageProvider(MultibandImageType, valpaths, image_suffix='.png')
print('valds.im_names: {}'.format(valds.im_names))
config_path = 'crop_pspnet_config.json'
with open(config_path, 'r') as f:
mycfg = json.load(f)
train_data_path = '{}/patch{}_train'.format(datapath, areaname)
dataset_path, train_dir = os.path.split(train_data_path)
mycfg['dataset_path'] = dataset_path
config = Config(**mycfg)
config = update_config(config, num_channels=12, nb_epoch=50)
#dataset_train = TrainDataset(trainds, train_idx, config, transforms=augment_flips_color)
dataset_train = TrainDataset(trainds, train_idx, config, 1)
dataset_val = ValDataset(valds, val_idx, config, 1)
trainloader = data.DataLoader(dataset_train,
batch_size=cfg['training']['batch_size'],
num_workers=cfg['training']['n_workers'],
shuffle=True)
valloader = data.DataLoader(dataset_val,
batch_size=cfg['training']['batch_size'],
num_workers=cfg['training']['n_workers'],
shuffle=False)
# Setup Metrics
running_metrics_train = runningScore(n_classes)
running_metrics_val = runningScore(n_classes)
nbackground = 1116403140
ncorn = 44080178
nsoybean = 316698122
print('nbackgraound: {}'.format(nbackground))
print('ncorn: {}'.format(ncorn))
print('nsoybean: {}'.format(nsoybean))
wgts = [1.0, 1.0 * nbackground / ncorn, 1.0 * nbackground / nsoybean]
total_wgts = sum(wgts)
wgt_background = wgts[0] / total_wgts
wgt_corn = wgts[1] / total_wgts
wgt_soybean = wgts[2] / total_wgts
weights = torch.autograd.Variable(
torch.cuda.FloatTensor([wgt_background, wgt_corn, wgt_soybean]))
# Setup Model
model = get_model(cfg['model'], n_classes).to(device)
model = torch.nn.DataParallel(model,
device_ids=range(torch.cuda.device_count()))
# Setup optimizer, lr_scheduler and loss function
optimizer_cls = get_optimizer(cfg)
optimizer_params = {
k: v
for k, v in cfg['training']['optimizer'].items() if k != 'name'
}
optimizer = optimizer_cls(model.parameters(), **optimizer_params)
scheduler = get_scheduler(optimizer, cfg['training']['lr_schedule'])
loss_fn = get_loss_function(cfg)
start_iter = 0
runpath = '/home/chengjjang/arisia/CropPSPNet/runs/pspnet_crop_{}'.format(
areaname)
modelpath = glob.glob('{}/*/*_best_model.pkl'.format(runpath))[0]
print('modelpath: {}'.format(modelpath))
checkpoint = torch.load(modelpath)
model.load_state_dict(checkpoint["model_state"])
val_loss_meter = averageMeter()
time_meter = averageMeter()
best_iou = -100.0
respath = '{}_results_val'.format(areaname)
os.makedirs(respath, exist_ok=True)
model.eval()
with torch.no_grad():
for inputdata in valloader:
imname_val = inputdata['img_name']
images_val = inputdata['img_data']
labels_val = inputdata['seg_label']
images_val = images_val.to(device)
labels_val = labels_val.to(device)
print('imname_val: {}'.format(imname_val))
outputs = model(images_val)
val_loss = loss_fn(input=outputs, target=labels_val)
pred = outputs.data.max(1)[1].cpu().numpy()
gt = labels_val.data.cpu().numpy()
dname = imname_val[0].split('.png')[0]
np.save('{}/pred'.format(respath) + dname + '.npy', pred)
np.save('{}/gt'.format(respath) + dname + '.npy', gt)
np.save('{}/output'.format(respath) + dname + '.npy',
outputs.data.cpu().numpy())
running_metrics_val.update(gt, pred)
val_loss_meter.update(val_loss.item())
#writer.add_scalar('loss/val_loss', val_loss_meter.avg, i+1)
#logger.info("Iter %d Loss: %.4f" % (i + 1, val_loss_meter.avg))
print('Test loss: {}'.format(val_loss_meter.avg))
score, class_iou = running_metrics_val.get_scores()
for k, v in score.items():
print('val_metrics, {}: {}'.format(k, v))
for k, v in class_iou.items():
print('val_metrics, {}: {}'.format(k, v))
val_loss_meter.reset()
running_metrics_val.reset()
def train(cfg, writer, logger):
# Setup seeds
torch.manual_seed(cfg.get('seed', 1337))
torch.cuda.manual_seed(cfg.get('seed', 1337))
np.random.seed(cfg.get('seed', 1337))
random.seed(cfg.get('seed', 1337))
# Setup device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Setup Augmentations
augmentations = cfg['training'].get('augmentations', None)
data_aug = get_composed_augmentations(augmentations)
# Setup Dataloader
# data_loader = get_loader(cfg['data']['dataset'])
# data_path = cfg['data']['path']
#
# t_loader = data_loader(
# data_path,
# is_transform=True,
# split=cfg['data']['train_split'],
# img_size=(cfg['data']['img_rows'], cfg['data']['img_cols']),
# augmentations=data_aug)
#
# v_loader = data_loader(
# data_path,
# is_transform=True,
# split=cfg['data']['val_split'],
# img_size=(cfg['data']['img_rows'], cfg['data']['img_cols']),)
#
# n_classes = t_loader.n_classes
# trainloader = data.DataLoader(t_loader,
# batch_size=cfg['training']['batch_size'],
# num_workers=cfg['training']['n_workers'],
# shuffle=True)
#
# valloader = data.DataLoader(v_loader,
# batch_size=cfg['training']['batch_size'],
# num_workers=cfg['training']['n_workers'])
paths = {
'masks': './satellitedata/patchohio_train/gt/',
'images': './satellitedata/patchohio_train/rgb',
'nirs': './satellitedata/patchohio_train/nir',
'swirs': './satellitedata/patchohio_train/swir',
'vhs': './satellitedata/patchohio_train/vh',
'vvs': './satellitedata/patchohio_train/vv',
'redes': './satellitedata/patchohio_train/rede',
'ndvis': './satellitedata/patchohio_train/ndvi',
}
valpaths = {
'masks': './satellitedata/patchohio_val/gt/',
'images': './satellitedata/patchohio_val/rgb',
'nirs': './satellitedata/patchohio_val/nir',
'swirs': './satellitedata/patchohio_val/swir',
'vhs': './satellitedata/patchohio_val/vh',
'vvs': './satellitedata/patchohio_val/vv',
'redes': './satellitedata/patchohio_val/rede',
'ndvis': './satellitedata/patchohio_val/ndvi',
}
n_classes = 3
train_img_paths = [pth for pth in os.listdir(paths['images']) if ('_01_' not in pth) and ('_25_' not in pth)]
val_img_paths = [pth for pth in os.listdir(valpaths['images']) if ('_01_' not in pth) and ('_25_' not in pth)]
ntrain = len(train_img_paths)
nval = len(val_img_paths)
train_idx = [i for i in range(ntrain)]
val_idx = [i for i in range(nval)]
trainds = ImageProvider(MultibandImageType, paths, image_suffix='.png')
valds = ImageProvider(MultibandImageType, valpaths, image_suffix='.png')
config_path = 'crop_pspnet_config.json'
with open(config_path, 'r') as f:
mycfg = json.load(f)
train_data_path = './satellitedata/'
print('train_data_path: {}'.format(train_data_path))
dataset_path, train_dir = os.path.split(train_data_path)
print('dataset_path: {}'.format(dataset_path) + ', train_dir: {}'.format(train_dir))
mycfg['dataset_path'] = dataset_path
config = Config(**mycfg)
config = update_config(config, num_channels=12, nb_epoch=50)
#dataset_train = TrainDataset(trainds, train_idx, config, transforms=augment_flips_color)
dataset_train = TrainDataset(trainds, train_idx, config, 1)
dataset_val = TrainDataset(valds, val_idx, config, 1)
trainloader = data.DataLoader(dataset_train,
batch_size=cfg['training']['batch_size'],
num_workers=cfg['training']['n_workers'],
shuffle=True)
valloader = data.DataLoader(dataset_val,
batch_size=cfg['training']['batch_size'],
num_workers=cfg['training']['n_workers'],
shuffle=False)
# Setup Metrics
running_metrics_train = runningScore(n_classes)
running_metrics_val = runningScore(n_classes)
k = 0
nbackground = 0
ncorn = 0
#ncotton = 0
#nrice = 0
nsoybean = 0
for indata in trainloader:
k += 1
gt = indata['seg_label'].data.cpu().numpy()
nbackground += (gt == 0).sum()
ncorn += (gt == 1).sum()
#ncotton += (gt == 2).sum()
#nrice += (gt == 3).sum()
nsoybean += (gt == 2).sum()
print('k = {}'.format(k))
print('nbackgraound: {}'.format(nbackground))
print('ncorn: {}'.format(ncorn))
#print('ncotton: {}'.format(ncotton))
#print('nrice: {}'.format(nrice))
print('nsoybean: {}'.format(nsoybean))
wgts = [1.0, 1.0*nbackground/ncorn, 1.0*nbackground/nsoybean]
total_wgts = sum(wgts)
wgt_background = wgts[0]/total_wgts
wgt_corn = wgts[1]/total_wgts
#wgt_cotton = wgts[2]/total_wgts
#wgt_rice = wgts[3]/total_wgts
wgt_soybean = wgts[2]/total_wgts
weights = torch.autograd.Variable(torch.cuda.FloatTensor([wgt_background, wgt_corn, wgt_soybean]))
#weights = torch.autograd.Variable(torch.cuda.FloatTensor([1.0, 1.0, 1.0]))
# Setup Model
model = get_model(cfg['model'], n_classes).to(device)
model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
# Setup optimizer, lr_scheduler and loss function
optimizer_cls = get_optimizer(cfg)
optimizer_params = {k:v for k, v in cfg['training']['optimizer'].items()
if k != 'name'}
optimizer = optimizer_cls(model.parameters(), **optimizer_params)
logger.info("Using optimizer {}".format(optimizer))
scheduler = get_scheduler(optimizer, cfg['training']['lr_schedule'])
loss_fn = get_loss_function(cfg)
logger.info("Using loss {}".format(loss_fn))
start_iter = 0
if cfg['training']['resume'] is not None:
if os.path.isfile(cfg['training']['resume']):
logger.info(
"Loading model and optimizer from checkpoint '{}'".format(cfg['training']['resume'])
)
checkpoint = torch.load(cfg['training']['resume'])
model.load_state_dict(checkpoint["model_state"])
optimizer.load_state_dict(checkpoint["optimizer_state"])
scheduler.load_state_dict(checkpoint["scheduler_state"])
start_iter = checkpoint["epoch"]
logger.info(
"Loaded checkpoint '{}' (iter {})".format(
cfg['training']['resume'], checkpoint["epoch"]
)
)
else:
logger.info("No checkpoint found at '{}'".format(cfg['training']['resume']))
val_loss_meter = averageMeter()
time_meter = averageMeter()
best_iou = -100.0
i = start_iter
flag = True
while i <= cfg['training']['train_iters'] and flag:
for inputdata in trainloader:
i += 1
start_ts = time.time()
scheduler.step()
model.train()
images = inputdata['img_data']
labels = inputdata['seg_label']
#print('images.size: {}'.format(images.size()))
#print('labels.size: {}'.format(labels.size()))
images = images.to(device)
labels = labels.to(device)
optimizer.zero_grad()
outputs = model(images)
#print('outputs.size: {}'.format(outputs[1].size()))
#print('labels.size: {}'.format(labels.size()))
loss = loss_fn(input=outputs[1], target=labels, weight=weights)
loss.backward()
optimizer.step()
time_meter.update(time.time() - start_ts)
if (i + 1) % cfg['training']['print_interval'] == 0:
fmt_str = "Iter [{:d}/{:d}] Loss: {:.4f} Time/Image: {:.4f}"
print_str = fmt_str.format(i + 1,
cfg['training']['train_iters'],
loss.item(),
time_meter.avg / cfg['training']['batch_size'])
print(print_str)
logger.info(print_str)
writer.add_scalar('loss/train_loss', loss.item(), i+1)
time_meter.reset()
if (i + 1) % cfg['training']['val_interval'] == 0 or \
(i + 1) == cfg['training']['train_iters']:
model.eval()
with torch.no_grad():
for inputdata in valloader:
images_val = inputdata['img_data']
labels_val = inputdata['seg_label']
images_val = images_val.to(device)
labels_val = labels_val.to(device)
outputs = model(images_val)
val_loss = loss_fn(input=outputs, target=labels_val)
pred = outputs.data.max(1)[1].cpu().numpy()
gt = labels_val.data.cpu().numpy()
running_metrics_val.update(gt, pred)
val_loss_meter.update(val_loss.item())
writer.add_scalar('loss/val_loss', val_loss_meter.avg, i+1)
logger.info("Iter %d Loss: %.4f" % (i + 1, val_loss_meter.avg))
score, class_iou = running_metrics_val.get_scores()
for k, v in score.items():
print(k, v)
logger.info('{}: {}'.format(k, v))
writer.add_scalar('val_metrics/{}'.format(k), v, i+1)
for k, v in class_iou.items():
logger.info('{}: {}'.format(k, v))
writer.add_scalar('val_metrics/cls_{}'.format(k), v, i+1)
val_loss_meter.reset()
running_metrics_val.reset()
if score["Mean IoU : \t"] >= best_iou:
best_iou = score["Mean IoU : \t"]
state = {
"epoch": i + 1,
"model_state": model.state_dict(),
"optimizer_state": optimizer.state_dict(),
"scheduler_state": scheduler.state_dict(),
"best_iou": best_iou,
}
save_path = os.path.join(writer.file_writer.get_logdir(),
"{}_{}_best_model.pkl".format(
cfg['model']['arch'],
cfg['data']['dataset']))
torch.save(state, save_path)
if (i + 1) == cfg['training']['train_iters']:
flag = False
break
def test(args):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_file_name = os.path.split(args.model_path)[1]
model_name = model_file_name[:model_file_name.find("_")]
print('model_file_name: {}'.format(model_file_name))
print('model_name: {}'.format(model_name))
# Setup image
#print("Read Input Image from : {}".format(args.img_path))
#img = misc.imread(args.img_path)
#data_loader = get_loader(args.dataset)
#data_path = get_data_path(args.dataset)
#loader = data_loader(data_path, is_transform=True, img_norm=args.img_norm)
#n_classes = loader.n_classes
# Dataset and Loader
#list_test = [{'fpath_img': args.test_img}]
#dataset_val = TestDataset(
# list_test, args, max_sample=args.num_val)
# paths = {
# 'masks': '/home/local/KHQ/chengjiang.long/Projects/core3d/data/AOIS/4AOIs/pngdata/refine_gtl',
# 'images': '/home/local/KHQ/chengjiang.long/Projects/core3d/data/AOIS/4AOIs/pngdata/rgb',
# 'ndsms': '/home/local/KHQ/chengjiang.long/Projects/core3d/data/AOIS/4AOIs/pngdata/ndsm',
# 'ndvis': '/home/local/KHQ/chengjiang.long/Projects/core3d/data/AOIS/4AOIs/pngdata/ndvi',
# }
paths = {
'masks': '/data/CORE3D/AOIS/4AOIs/test_tmp/tile_building_fill',
'images': '/data/CORE3D/AOIS/4AOIs/test_tmp/tile_image',
'ndsms': '/data/CORE3D/AOIS/4AOIs/test_tmp/tile_dsm',
'ndvis': '/data/CORE3D/AOIS/4AOIs/test_tmp/NDVI',
}
num_classes = 3
ntest = len(os.listdir(paths['images']))
test_idx = [i for i in range(ntest)]
testds = ImageProvider(MultibandImageType, paths, image_suffix='.png')
config_path = 'lcj_denseunet_1x1080_retrain.json'
with open(config_path, 'r') as f:
cfg = json.load(f)
train_data_path = '/data/CORE3D/AOIS/Dayton_20sqkm/pngdata'
print('train_data_path: {}'.format(train_data_path))
dataset_path, train_dir = os.path.split(train_data_path)
print('dataset_path: {}'.format(dataset_path) +
', train_dir: {}'.format(train_dir))
cfg['dataset_path'] = dataset_path
config = Config(**cfg)
config = update_config(config, num_channels=5, nb_epoch=50)
#dataset_train = TrainDataset(trainds, train_idx, config, transforms=augment_flips_color)
dataset_test = ValDataset(testds, test_idx, config)
loader_test = data.DataLoader(dataset_test,
batch_size=1,
shuffle=False,
num_workers=2,
drop_last=True)
# Setup Model
mymodel_dict = dict()
mymodel_dict['arch'] = model_name
model = get_model(mymodel_dict, num_classes)
state = convert_state_dict(torch.load(args.model_path)["model_state"])
model.load_state_dict(state)
model.eval()
model.to(device)
for i, batch_data in enumerate(loader_test):
#if batch_data['img_name'][0] != 'D4.png':
# continue
# process data
# print('batch_data: {}'.format(batch_data))
#batch_data = batch_data[0]
print('batch_data[img_data].size: {}'.format(
batch_data['img_data'].shape))
segSize = (batch_data['img_data'].shape[2],
batch_data['img_data'].shape[3])
img = batch_data['img_data']
gt = batch_data['seg_label'].data.cpu().numpy()
dsize = img.shape
predicted = np.zeros((dsize[0], num_classes, dsize[2], dsize[3]))
predicted = torch.autograd.variable(torch.FloatTensor(predicted))
with torch.no_grad():
xinsidx, xineidx, xcpsidx, xcpeidx, xoutsidx, xouteidx = index_in_copy_out(
segSize[0])
yinsidx, yineidx, ycpsidx, ycpeidx, youtsidx, youteidx = index_in_copy_out(
segSize[1])
imgdata = torch.autograd.variable(torch.FloatTensor(img))
print('imgdata.size: {}'.format(imgdata.size()))
for i in range(len(xinsidx)):
for j in range(len(yinsidx)):
samples = torch.autograd.Variable(
imgdata[:, :, xinsidx[i]:xineidx[i],
yinsidx[j]:yineidx[j]],
volatile=True).cuda()
samples_img = samples.to(device)
# forward pass
prediction = model(samples_img)
#print('samples.size: {}'.format(samples_img.size()))
#print('prediction.size: {}'.format(prediction.size()))
#prediction = tmp_prediction.data.cpu().numpy()
#print('prediction.size: {}'.format(prediction.shape))
predicted[:, :, xoutsidx[i]:xouteidx[i], youtsidx[j]:
youteidx[j]] = prediction[:, :,
xcpsidx[i]:xcpeidx[i],
ycpsidx[j]:ycpeidx[j]]
#print('predicted: {}'.format(predicted.size()))
running_metrics_test = runningScore(num_classes)
print('gt.shape: {}'.format(gt.shape))
print('predicted.shape: {}'.format(predicted.shape))
pred = predicted.data.max(1)[1].cpu().numpy()
running_metrics_test.update(gt, pred)
print('score: {}'.format(running_metrics_test.get_scores()))
probmaps = np.squeeze(predicted.data.cpu().numpy(), axis=0)
print('preds.shape: {}'.format(probmaps.shape))
pred = np.squeeze(pred, axis=0)
visualize_result((batch_data['img_data'], batch_data['img_name']),
probmaps, pred)
if args.dcrf:
unary = predicted
unary = np.squeeze(unary, 0)
unary = -np.log(unary)
unary = unary.transpose(2, 1, 0)
w, h, c = unary.shape
unary = unary.transpose(2, 0, 1).reshape(num_classes, -1)
unary = np.ascontiguousarray(unary)
resized_img = img
d = dcrf.DenseCRF2D(w, h, loader.n_classes)
d.setUnaryEnergy(unary)
d.addPairwiseBilateral(sxy=5, srgb=3, rgbim=resized_img, compat=1)
q = d.inference(50)
mask = np.argmax(q, axis=0).reshape(w, h).transpose(1, 0)
decoded_crf = loader.decode_segmap(np.array(mask, dtype=np.uint8))
dcrf_path = args.out_path[:-4] + "_drf.png"
misc.imsave(dcrf_path, decoded_crf)
print("Dense CRF Processed Mask Saved at: {}".format(dcrf_path))
#decoded = loader.decode_segmap(pred)
print("Classes found: ", np.unique(pred))