def rebase_groundtruth(task, fullres, force=False):
"""
Inplace / lazy modification of groundtruth labels
hacky.
"""
# Remap the original three labels to [0, 1, 2]
orig_labels = [2, 6, 65]
mapping = np.full(max(orig_labels) + 1, fill_value=-1)
mapping[orig_labels] = np.arange(len(orig_labels))
datadir = ub.ensuredir((task.workdir, 'data'))
dpath = ub.ensuredir((datadir, 'gt', 'full'))
new_gt_paths = []
for ix in ub.ProgIter(range(len(fullres.paths['gt'])), label='rebase'):
path = fullres.paths['gt'][ix]
name = fullres.dump_im_names[ix]
out_dpath = join(dpath, name)
# Hacky cache
if force or not exists(out_dpath):
in_data = imutil.imread(path)
out_data = mapping[in_data]
imutil.imwrite(out_dpath, out_data)
new_gt_paths.append(out_dpath)
fullres.paths['gt'] = new_gt_paths
return fullres
def _lowres_proc(tag, out_dpath, in_paths, out_names, interpolation):
"""
pip install tifffile
"""
_iter = zip(in_paths, out_names)
prog = ub.ProgIter(_iter,
length=len(in_paths),
label='lowres ' + tag)
for in_path, out_name in prog:
out_path = join(out_dpath, out_name)
if not exists(out_path):
in_data = imutil.imread(in_path)
out_data = cv2.resize(in_data,
target_dsize,
interpolation=interpolation)
imutil.imwrite(out_path, out_data)
yield out_path
def np_loader(fpath, colorspace=None):
im_in = imutil.imread(fpath)
if colorspace is not None:
cv_255 = im_in # Assume we read a (bgr) byte image
cv_01 = cv_255.astype(np.float32) / 255.0
n_channels = imutil.get_num_channels(cv_01)
if n_channels == 1:
input_space = 'GRAY'
elif n_channels == 3:
input_space = 'BGR'
elif n_channels == 4:
input_space = 'BGRA'
else:
raise NotImplementedError()
output_space = colorspace.upper()
if output_space == 'L':
output_space = 'GRAY'
im_out = imutil.convert_colorspace(cv_01,
dst_space=output_space,
src_space=input_space)
else:
im_out = im_in
return im_out
def restitch(task, output_dpath, part_paths, blend=None):
# Group parts by base id
if len(part_paths) == 0:
return []
def stitch_tiles(rc_locs, tiles):
"""
Recombine parts back into an entire image
(TODO: blending / averaging to remove boundry effects)
"""
shapes = [t.shape[0:2] for t in tiles]
n_channels = 1 if len(tiles[0].shape) == 2 else tiles[0].shape[2]
bboxes = np.array([
(r, c, r + h, c + w)
for ((r, c), (h, w)) in zip(rc_locs, shapes)
])
stiched_shape = tuple(bboxes.T[2:4].max(axis=1))
if n_channels > 1:
stiched_shape = stiched_shape + (n_channels,)
stiched = np.zeros(stiched_shape)
for bbox, tile in zip(bboxes, tiles):
r1, c1, r2, c2 = bbox
stiched[r1:r2, c1:c2] = tile
return stiched
def stitch_tiles_ave(rc_locs, tiles, weighted=False):
"""
Recombine parts back into an entire image
Example:
>>> rc_locs = [(0, 0), (0, 5), (0, 10)]
>>> tiles = [np.ones((1, 7, 3)) + i for i in range(len(rc_locs))]
>>> tiles = [np.ones((1, 7)) + i for i in range(len(rc_locs))]
"""
shapes = [t.shape[0:2] for t in tiles]
n_channels = 1 if len(tiles[0].shape) == 2 else tiles[0].shape[2]
bboxes = np.array([
(r, c, r + h, c + w)
for ((r, c), (h, w)) in zip(rc_locs, shapes)
])
stiched_wh = tuple(bboxes.T[2:4].max(axis=1))
stiched_shape = stiched_wh
if n_channels > 1:
stiched_shape = stiched_wh + (n_channels,)
sums = np.zeros(stiched_shape)
nums = np.zeros(stiched_wh)
if weighted:
# assume all shapes are the same
h, w = shapes[0]
weight = np.ones((h, w) )
# Weight borders less than center
# should really use receptive fields for this calculation
# but this should be fine.
weight[:h // 4] = .25
weight[-h // 4:] = .25
weight[:w // 4] = .25
weight[-w // 4:] = .25
weight3c = weight
if n_channels > 1:
weight3c = weight[:, :, None]
else:
weight3c = weight = 1
# Assume we are not in log-space here, so the weighted average
# formula does not need any exponentiation.
for bbox, tile in zip(bboxes, tiles):
r1, c1, r2, c2 = bbox
sums[r1:r2, c1:c2] += (tile * weight3c)
nums[r1:r2, c1:c2] += weight
if len(sums.shape) == 2:
stiched = sums / nums
else:
stiched = sums / nums[:, :, None]
return stiched
def stitch_tiles_vote(rc_locs, tiles):
"""
Recombine parts back into an entire image
(TODO: blending / averaging to remove boundry effects)
"""
shapes = [t.shape[0:2] for t in tiles]
n_channels = 1 if len(tiles[0].shape) == 2 else tiles[0].shape[2]
assert n_channels == 1
bboxes = np.array([
(r, c, r + h, c + w)
for ((r, c), (h, w)) in zip(rc_locs, shapes)
])
stiched_shape = tuple(bboxes.T[2:4].max(axis=1))
n_classes = len(task.classnames)
votes = np.zeros((n_classes,) + stiched_shape)
for bbox, tile in zip(bboxes, tiles):
r1, c1, r2, c2 = bbox
for i in range(n_classes):
votes[i, r1:r2, c1:c2][tile == i] += 1
stiched = votes.argmax(axis=0)
return stiched
def _extract_part_grid(paths):
# hack to use filenames to extract upper left locations of tiles in
# the larger image.
rc_locs = [[int(x) for x in basename(p).split('.')[0].split('_')[-2:]]
for p in paths]
return rc_locs
ext = splitext(part_paths[0])[1]
is_image = ext not in ['.npy', '.h5']
groupid = [basename(p).split('_part')[0] for p in part_paths]
new_paths = []
groups = list(ub.group_items(part_paths, groupid).items())
for tileid, paths in ub.ProgIter(groups, label='restitching'):
# Read all parts belonging to an original group
if is_image:
tiles = [imutil.imread(p) for p in paths]
else:
tiles = [util.read_arr(p) for p in paths]
if tiles[0].shape[0] < 10:
# hack
tiles = [t.transpose(1, 2, 0) for t in tiles]
# try:
# except OSError:
# Find their relative positions and restitch them
rc_locs = _extract_part_grid(paths)
if blend == 'vote':
stiched = stitch_tiles_vote(rc_locs, tiles)
elif blend == 'ave':
stiched = stitch_tiles_ave(rc_locs, tiles, weighted=False)
elif blend == 'avew':
stiched = stitch_tiles_ave(rc_locs, tiles, weighted=True)
elif blend is None:
stiched = stitch_tiles(rc_locs, tiles)
else:
raise KeyError(blend)
# Write them to disk.
if is_image:
fpath = join(output_dpath, tileid + '.png')
imutil.imwrite(fpath, stiched)
else:
fpath = join(output_dpath, tileid + '.npy')
util.write_arr(fpath, stiched)
new_paths.append(fpath)
return new_paths
def create_boundary_groundtruth(task, fullres, force=False):
# Hack: transform task into boundary mode
task._boundary_mode()
NON_BUILDING = task.classname_to_id['non-building']
INNER_INSTANCE = task.classname_to_id['inner_building']
OUTER_INSTANCE = task.classname_to_id['outer_building']
UNKNOWN = task.classname_to_id['uncertain']
def instance_boundary(gti_data, gtl_data):
"""
gti_data = np.array([
[0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0],
[0, 0, 5, 5, 5, 5, 5, 5, 0, 0, 0],
[0, 0, 5, 5, 5, 5, 5, 5, 2, 2, 2],
[0, 0, 5, 5, 5, 5, 5, 5, 2, 2, 2],
[0, 5, 5, 5, 5, 5, 5, 5, 2, 2, 2],
[0, 5, 5, 5, 5, 5, 5, 5, 0, 0, 0],
[0, 5, 5, 5, 5, 5, 5, 5, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 3, 3, 0, 4, 0, 0],
[1, 0, 0, 0, 0, 3, 3, 0, 0, 0, 0],
], dtype=np.uint8)
kernel = np.ones((2, 2))
cv2.erode(im, kernel)
"""
out_data = np.full(gti_data.shape, dtype=np.uint8,
fill_value=NON_BUILDING)
# 5x5 erosion is equivalent to two iterations of a 3x3 erosion
kernel = np.ones((3, 3))
for label, submask, rc_off, rc_sl in instance_submasks(gti_data):
pad = 1
border = cv2.copyMakeBorder(submask, pad, pad, pad, pad,
cv2.BORDER_CONSTANT, value=0)
inner = cv2.erode(border, kernel, iterations=2)[1:-1, 1:-1]
outer = submask - inner
newmask = ((inner * INNER_INSTANCE) | (outer * OUTER_INSTANCE))
out_data[rc_sl] = newmask
# Map old unknown value to the a new one
is_unknown = gtl_data == 65
out_data[is_unknown] = UNKNOWN
return out_data
# Augment the groundtruth so the network must also predict instance
# boundaries
datadir = ub.ensuredir((task.workdir, 'data'))
dpath = ub.ensuredir((datadir, 'gtb', 'full'))
new_gt_paths = []
for ix in ub.ProgIter(range(len(fullres.paths['gti'])), label='boundary'):
path = fullres.paths['gti'][ix]
name = fullres.dump_im_names[ix]
out_dpath = join(dpath, name)
if force or not exists(out_dpath):
gtl_data = imutil.imread(fullres.paths['gt'][ix])
gti_data = imutil.imread(path)
out_data = instance_boundary(gti_data, gtl_data)
imutil.imwrite(out_dpath, out_data)
new_gt_paths.append(out_dpath)
fullres.paths['gt'] = new_gt_paths
return fullres
def dump_task_inference(task, inputs, load_path):
"""
task = get_task('urban_mapper_3d')
(train, test), = task.xval_splits()
inputs = test
eval_dataset = SSegInputsWrapper(inputs, task, colorspace='RGB')
eval_dataset.with_gt = False
eval_dataset.inputs.make_dumpsafe_names()
load_path = 'solver_4214-yxalqwdk_unet_vgg_nttxoagf_a=1,n_ch=5,n_cl=3_epoch_00000236.pt'
"""
eval_dataset = SSegInputsWrapper(inputs, task, colorspace='RGB')
eval_dataset.with_gt = False
eval_dataset.inputs.make_dumpsafe_names()
if True:
# TODO: make model metadata know about this
eval_dataset.center_inputs = eval_dataset._original_urban_mapper_normalizer(
)
eval_dpath = ub.ensuredir(
(task.workdir, 'eval', 'input_' + eval_dataset.input_id))
subdir = list(ub.take(os.path.splitext(load_path)[0].split('/'), [-3, -1]))
# base output dump path on the training id string
test_dump_dpath = ub.ensuredir((eval_dpath, '/'.join(subdir)))
print('test_dump_dpath = {!r}'.format(test_dump_dpath))
datasets = {'eval': eval_dataset}
# TODO n_classes and n_channels should be saved as model metadata
n_classes = datasets['eval'].n_classes
n_channels = datasets['eval'].n_channels
xpu = xpu_device.XPU.from_argv()
print('Loading snapshot onto {}'.format(xpu))
snapshot = torch.load(load_path, map_location=xpu.map_location())
# Infer which model this belongs to
if snapshot['model_class_name'] == 'UNet':
model = models.UNet(in_channels=n_channels, n_classes=n_classes)
elif snapshot['model_class_name'] == 'SegNet':
model = models.SegNet(in_channels=n_channels, n_classes=n_classes)
model = xpu.move(model)
model.load_state_dict(snapshot['model_state_dict'])
print('Preparing to predict {} on {}'.format(model.__class__.__name__,
xpu))
model.train(False)
for ix in ub.ProgIter(range(len(eval_dataset)), label='dumping'):
inputs_ = eval_dataset[ix][None, :]
inputs_ = xpu.move(inputs_)
inputs_ = torch.autograd.Variable(inputs_)
output_tensor = model(inputs_)
log_prob_tensor = torch.nn.functional.log_softmax(output_tensor,
dim=1)[0]
log_probs = log_prob_tensor.data.cpu().numpy()
# Just reload rgb data without trying to go through the reverse
# transform
img = imutil.imread(eval_dataset.inputs.im_paths[ix])
# ut.save_cPkl('crf_testdata.pkl', {
# 'log_probs': log_probs,
# 'img': img,
# })
from clab import filters
posterior = filters.crf_posterior(img, log_probs)
# output = prob_tensor.data.cpu().numpy()[0]
pred = log_probs.argmax(axis=0)
pred_crf = posterior.argmax(axis=0)
fname = eval_dataset.inputs.dump_im_names[ix]
fname = os.path.splitext(fname)[0] + '.png'
# pred = argmax.data.cpu().numpy()[0]
blend_pred = task.colorize(pred, img)
blend_pred_crf = task.colorize(pred_crf, img)
# color_pred = task.colorize(pred)
output_dict = {
'log_probs': log_probs,
'blend_pred': blend_pred,
# 'color_pred': color_pred,
'blend_pred_crf': blend_pred_crf,
'pred_crf': pred_crf,
'pred': pred,
}
if eval_dataset.with_gt:
true = imutil.imread(eval_dataset.inputs.gt_paths[ix])
blend_true = task.colorize(true, img, alpha=.5)
# color_true = task.colorize(true, alpha=.5)
output_dict['true'] = true
output_dict['blend_true'] = blend_true
# output_dict['color_true'] = color_true
for key, img in output_dict.items():
dpath = join(test_dump_dpath, key)
ub.ensuredir(dpath)
fpath = join(dpath, fname)
if key == 'log_probs':
np.savez(fpath.replace('.png', ''), img)
else:
imutil.imwrite(fpath, img)
return test_dump_dpath
def make_parts(prep, fullres, scale=1, clear=False):
"""
Slices the fullres images into smaller parts that fit into the network
but are at the original resolution (or higher).
>>> from clab.tasks.urban_mapper_3d import *
>>> task = UrbanMapper3D(root='~/remote/aretha/data/UrbanMapper3D', workdir='~/data/work/urban_mapper')
>>> task.prepare_fullres_inputs()
>>> fullres = task.fullres
>>> datadir = ub.ensuredir((task.workdir, 'data'))
>>> prep = Preprocessor(datadir)
>>> scale = 1
>>> clear = False
>>> lowres = prep.make_parts(fullres, scale)
"""
part_config = prep.part_config
hashid = hashutil.hash_data(ub.repr2(part_config), hashlen=8)
shapestr = '_'.join(list(map(str, prep.input_shape)))
mode = 'part-scale{}-{}-{}'.format(scale, shapestr, hashid)
parts, flag = prep._mode_new_input(mode, fullres, clear=clear)
if flag:
return parts
input_shape = prep.input_shape
overlap = part_config['overlap']
keepbound = part_config['keepbound']
records = list(fullres.iter_records())
for record in ub.ProgIter(records, label='make ' + mode):
dump_fname = basename(record['dump_fname'])
im_shape = np.array(Image.open(record['im']).size[::-1])
im_shape = tuple(np.floor(im_shape * scale).astype(np.int))
# Consolodate all channels that belong to this record
in_paths = record.get('aux').copy()
for k in ['im', 'gt']:
if k in record:
in_paths[k] = record[k]
# Read the images for this record and resize if necessary
in_images = {k: imutil.imread(v)
for k, v in in_paths.items()} # 9% of the time
if scale != 1.0:
for k in in_images.keys():
interp = cv2.INTER_LANCZOS4 if k == 'im' else cv2.INTER_NEAREST
in_images[k] = imutil.imscale(in_images[k], scale,
interp)[0]
sl_gen = imutil.image_slices(im_shape, input_shape, overlap,
keepbound)
for idx, rc_slice in enumerate(sl_gen):
rsl, csl = rc_slice
suffix = '_part{:0=4d}_{:0=3d}_{:0=3d}'.format(
idx, rsl.start, csl.start)
fname = ub.augpath(dump_fname, suffix=suffix)
for k, in_data in in_images.items():
out_data = in_data[rc_slice]
out_fpath = join(parts.dirs[k], fname)
imutil.imwrite(out_fpath, out_data) # 84% of the time
parts.paths[k].append(out_fpath)
return parts