import argparse import math from pathlib import Path import json from PIL import Image import cv2 import torch from torchvision import transforms as torch_trans import numpy as np from evalmde.utils.proj import depth_to_xyz from evalmde.utils.common import assign_item_to_dict, pathlib_file from evalmde.utils.image import resize from evalmde.utils.image import imread_rgb from evalmde.utils.depth import load_data from evalmde.utils.np_and_th import get_shifted_data @torch.no_grad() def compute_grid_lb_ub(data, i, j): ''' . . . ------------- |(0,0)|(0,1)| . ------.------ . |(1,0)|(1,1)| ------------- . . . ''' if i == 0 and j == 0: x00 = .25 * (data[:-1, :-1] + data[:-1, 1:] + data[1:, :-1] + data[1:, 1:]) x01 = .5 * (data[:-1, 1:] + data[1:, 1:]) x10 = .5 * (data[1:, :-1] + data[1:, 1:]) x11 = 1. * data[1:, 1:] elif i == 0 and j == 1: x00 = .5 * (data[:-1, :-1] + data[1:, :-1]) x01 = .25 * (data[:-1, :-1] + data[:-1, 1:] + data[1:, :-1] + data[1:, 1:]) x10 = 1. * data[1:, :-1] x11 = .5 * (data[1:, :-1] + data[1:, 1:]) elif i == 1 and j == 0: x00 = .5 * (data[:-1, :-1] + data[:-1, 1:]) x01 = 1. * data[:-1, 1:] x10 = .25 * (data[:-1, :-1] + data[:-1, 1:] + data[1:, :-1] + data[1:, 1:]) x11 = .5 * (data[:-1, 1:] + data[1:, 1:]) else: x00 = 1. * data[:-1, :-1] x01 = .5 * (data[:-1, :-1] + data[:-1, 1:]) x10 = .5 * (data[:-1, :-1] + data[1:, :-1]) x11 = .25 * (data[:-1, :-1] + data[:-1, 1:] + data[1:, :-1] + data[1:, 1:]) x = torch.stack([x00, x01, x10, x11], dim=-1) lb, ub = x.min(dim=-1).values, x.max(dim=-1).values # (H - 1, W - 1), (H - 1, W - 1) del x return lb, ub @torch.no_grad() def compute_high_res_idx(high_res_shape, data_low_res, valid, valid_high_res, gap, val_lb): ''' :param high_res_shape: (Hu, Wu) :param data_low_res: shape (Hl, Wl) :param valid_high_res: shape (Hl, Wl) :param gap: :param val_lb: :return: res_high_res res_high_res: shape (Hu, Wu) ''' Hu, Wu = high_res_shape # fill invalid pixels with neighbor means data_low_res = data_low_res.clone() nb_data_sum = torch.zeros_like(data_low_res) nb_data_cnt = torch.zeros_like(data_low_res) for di in [-1, 0, 1]: for dj in [-1, 0, 1]: nb_valid = get_shifted_data(valid, di, dj) nb_data = get_shifted_data(data_low_res, di, dj) nb_data_sum[nb_valid] += nb_data[nb_valid] nb_data_cnt[nb_valid] += 1 nb_data_sum[nb_data_cnt < .5] = 0 data_low_res[~valid] = (nb_data_sum / nb_data_cnt.clamp(min=1))[~valid] data_high_res = torch_trans.functional.resize(data_low_res[None], (Hu, Wu), torch_trans.InterpolationMode.BILINEAR)[0] res_high_res = -torch.ones((Hu, Wu), dtype=torch.int32, device=data_high_res.device) for i in range(2): for j in range(2): lb, ub = compute_grid_lb_ub(data_high_res, i, j) lb_i = torch.clip(torch.ceil((lb - val_lb) / gap), min=0).to(res_high_res.dtype) ub_i = torch.clip(torch.floor((ub - val_lb) / gap), max=2e9).to(res_high_res.dtype) multi_line_mask = (lb_i < ub_i) | ((lb_i == ub_i) & (res_high_res[1 - i: Hu - i, 1 - j: Wu - j] != -1)) single_line_mask = (lb_i == ub_i) & (res_high_res[1 - i: Hu - i, 1 - j: Wu - j] == -1) res_high_res[1 - i: Hu - i, 1 - j: Wu - j][single_line_mask] = lb_i[single_line_mask] multi_line_upd_idx = torch.clip(torch.round((data_high_res[1 - i: Hu - i, 1 - j: Wu - j] - val_lb) / gap), min=0, max=2e9).to(res_high_res.dtype) multi_line_upd_idx = torch.where(multi_line_upd_idx < lb_i, lb_i, multi_line_upd_idx) multi_line_upd_idx = torch.where(multi_line_upd_idx > ub_i, ub_i, multi_line_upd_idx) multi_line_upd_mask = ((res_high_res[1 - i: Hu - i, 1 - j: Wu - j] == -1) | ( torch.abs(data_high_res[1 - i: Hu - i, 1 - j: Wu - j] - (res_high_res[1 - i: Hu - i, 1 - j: Wu - j] * gap + val_lb)) > torch.abs(data_high_res[1 - i: Hu - i, 1 - j: Wu - j] - (multi_line_upd_idx * gap + val_lb)) )) & multi_line_mask res_high_res[1 - i: Hu - i, 1 - j: Wu - j][multi_line_upd_mask] = multi_line_upd_idx[multi_line_upd_mask] res_high_res[~valid_high_res] = -1 return res_high_res def get_contour_line_gap(data: torch.Tensor, valid: torch.Tensor, num_gap, qt): if not valid.any(): return 1 qt_lb = data[valid].quantile(qt).item() qt_ub = data[valid].quantile(1 - qt).item() gap = (qt_ub - qt_lb) / (num_gap * (1 - qt * 2)) return gap @torch.no_grad() def gen_contour_line(rgb_high_res, data, valid, valid_high_res, is_z, num_gap, shift, thickness=0, qt=0.05, colormap=cv2.COLORMAP_JET): device = 'cuda' if torch.cuda.is_available() else 'cpu' data = torch.from_numpy(data).to(device) valid = torch.from_numpy(valid).to(device) valid_high_res = torch.from_numpy(valid_high_res).to(device) if is_z: data = 1 / data gap = get_contour_line_gap(data, valid, num_gap, qt) data_lb = data[valid].min().item() if valid.any() else 0 val_lb = data_lb + gap * shift high_res_shape = rgb_high_res.shape[:2] res_high_res = compute_high_res_idx(high_res_shape, data, valid, valid_high_res, gap, val_lb) res = res_high_res.clone() dlt_rng = int(math.floor(thickness)) for di in range(-dlt_rng, dlt_rng + 1): for dj in range(-dlt_rng, dlt_rng + 1): if di * di + dj * dj > thickness * thickness: continue nb_res = get_shifted_data(res_high_res, di, dj) upd_mask = get_shifted_data(valid_high_res, di, dj) & (res == -1) & (nb_res != -1) & valid_high_res res[upd_mask] = nb_res[upd_mask] if (res != -1).any(): res[res != -1] -= res[res != -1].min() res = res.cpu().numpy() num_val = max(2, res.max().item() + 1) valid_high_res = valid_high_res.cpu().numpy() base_col = cv2.applyColorMap(np.arange(256, dtype=np.uint8)[None], colormap)[0].astype(np.float32) # (256, 3) idx = np.arange(num_val, dtype=np.float32) / (num_val - 1) * 255 # (itr,) idx_lb = np.floor(idx).astype(np.int32) # (itr,) coef_lb = (idx_lb.astype(np.float32) + 1 - idx)[:, None] # (itr, 1) col = base_col[idx_lb] * coef_lb + base_col[np.clip(idx_lb + 1, a_min=None, a_max=255)] * (1 - coef_lb) # (itr, 3) col = np.round(col).astype(np.uint8) img = np.zeros_like(rgb_high_res) non_colored_mask = valid_high_res & (res == -1) img[non_colored_mask] = rgb_high_res[non_colored_mask] colored_mask = valid_high_res & (res != -1) img[colored_mask] = col[res[colored_mask]] return img, colored_mask, col def pil_ds(img: np.ndarray, H, W): pil_img = Image.fromarray(img, mode='RGB') pil_img = pil_img.resize((W, H), Image.Resampling.LANCZOS) return np.array(pil_img) def render_contour_line_imgs(xyz: np.ndarray, valid: np.ndarray, rgb_low_res: np.ndarray, save_shape, out_root): ''' :param xyz: :param valid: :param rgb_low_res: :param save_shape: (H, W) :param out_root: :return: ''' # hyperparams texture_strength = 0.8 draw_dim_lb = 4 * np.linalg.norm([1920, 1080]) out_root = pathlib_file(out_root) dim = np.linalg.norm(rgb_low_res.shape[:2]) us_sc = int(math.ceil(draw_dim_lb / dim)) us_shape = (us_sc * rgb_low_res.shape[0], us_sc * rgb_low_res.shape[1]) rgb_high_res = np.round(texture_strength * cv2.resize(rgb_low_res, (us_shape[1], us_shape[0]))).astype(np.uint8) valid_high_res = torch_trans.functional.resize(torch.from_numpy(valid)[None], rgb_high_res.shape[:2], torch_trans.InterpolationMode.NEAREST_EXACT)[0].numpy() summary = {} for thickness in [5 * np.linalg.norm(rgb_high_res.shape[:2]) / (4 * np.linalg.norm([1920, 1080]))]: for rel_num_gap in [0.015, 0.03, 0.06, 0.09, 0.12, 0.24, 0.42, 0.6]: num_gap = int(dim * rel_num_gap) for shift in [0.5]: imgs, colored_masks, col_maps = {}, {}, {} for i, name in enumerate(['x', 'y', 'z']): imgs[name], colored_masks[name], col_maps[name] = \ gen_contour_line(rgb_high_res, xyz[..., i], valid, valid_high_res, name == 'z', num_gap, shift, thickness) out_f = out_root / name / f'thickness__{thickness:.1f}___num_gap__{num_gap}___shift__{shift:.2f}.png' out_f.parent.mkdir(parents=True, exist_ok=True) cv2.imwrite(out_f.as_posix(), pil_ds(imgs[name][us_sc:-us_sc, us_sc:-us_sc, ::-1].copy(), save_shape[0], save_shape[1])) print(f'Saved to {out_f.resolve()}') assign_item_to_dict(summary, [name, thickness, num_gap, shift], str(out_f.resolve().relative_to(out_root.resolve()))) img_xy = rgb_high_res.copy() img_xy[np.logical_and(colored_masks['x'], colored_masks['y'])] = np.round(.5 * (imgs['x'].astype(np.float32) + imgs['y'].astype(np.float32))).astype(np.uint8)[np.logical_and(colored_masks['x'], colored_masks['y'])] img_xy[np.logical_and(colored_masks['x'], np.logical_not(colored_masks['y']))] = imgs['x'][np.logical_and(colored_masks['x'], np.logical_not(colored_masks['y']))] img_xy[np.logical_and(np.logical_not(colored_masks['x']), colored_masks['y'])] = imgs['y'][np.logical_and(np.logical_not(colored_masks['x']), colored_masks['y'])] img_xy[~valid_high_res] = 0 # img_xy = caption_img_xy(img_xy, col_maps) out_f = out_root / 'xy' / f'thickness__{thickness:.1f}___num_gap__{num_gap}___shift__{shift:.2f}.png' out_f.parent.mkdir(parents=True, exist_ok=True) cv2.imwrite(out_f.as_posix(), pil_ds(img_xy[us_sc:-us_sc, us_sc:-us_sc, ::-1].copy(), save_shape[0], save_shape[1])) print(f'Saved to {out_f.resolve()}') assign_item_to_dict(summary, ['xy', thickness, num_gap, shift], str(out_f.resolve().relative_to(out_root.resolve()))) with (out_root / 'summary.json').open('w') as F: json.dump(summary, F) def get_out_dir(work_dir, depth_f): return work_dir / 'contour_line' / str((work_dir / depth_f).resolve().relative_to(work_dir.resolve()))[:-4].replace('/', '_') def main(args): save_dim_ub = args.save_dim_ub root = args.root rgb_f = root / args.rgb_f data_f = root / args.depth_f raw_rgb = imread_rgb(rgb_f) save_sc = int(math.floor(save_dim_ub / np.linalg.norm(raw_rgb.shape[:2]))) save_shape = (save_sc * raw_rgb.shape[0], save_sc * raw_rgb.shape[1]) depth, intr, valid = load_data(data_f) xyz = depth_to_xyz(intr, depth) render_contour_line_imgs(xyz, valid, raw_rgb, save_shape, get_out_dir(root, args.depth_f)) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("root", type=Path) parser.add_argument("--depth_f", type=str, help='Path to depth file, relative to root.') parser.add_argument('--rgb_f', type=str, nargs='?', const=None, default='rgb.png', help='Path to rgb file, relative to root.') parser.add_argument("--save_dim_ub", type=float, default=np.linalg.norm([1920, 1080])) args = parser.parse_args() main(args)