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