File size: 5,752 Bytes
ef6a8c6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
import cv2
from cv2 import imread
import numpy as np
from numpy import unique, load
import argparse
import os, os.path as osp
import tqdm
import pdb
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--root', type=str, required=True, help='root dir of one identity')
parser.add_argument('--cam_file', type=str, default='camera.npz')
parser.add_argument('--depth_dir', help='depth map path', default='depth', type=str)
parser.add_argument('--out_dir', '-n', help='outputnormal path', default='normal_depth', type=str)
parser.add_argument('--mask_dir', help='mask path', default='masks', type=str)
parser.add_argument('--mask_total_dir', help='total mask path', default='masks_total', type=str)
args = parser.parse_args()
return args
def get_surface_normal_by_depth(depth, fx, fy):
"""
depth: (h, w) of float, the unit of depth is meter
K: (3, 3) of float, the depth camere's intrinsic
"""
dz_dv, dz_du = np.gradient(depth) # u, v mean the pixel coordinate in the image
# u*depth = fx*x + cx --> du/dx = fx / dept
# smooth
# dz_du = cv2.Sobel(depth, cv2.CV_64F, 1, 0, ksize=5)
# dz_dv = cv2.Sobel(depth, cv2.CV_64F, 0, 1, ksize=5)
du_dx = fx / depth # x is xyz of camera coordinate
dv_dy = fy / depth
dz_dx = dz_du * du_dx
dz_dy = dz_dv * dv_dy
# cross-product (1,0,dz_dx)X(0,1,dz_dy) = (-dz_dx, -dz_dy, 1)
normal_cross = np.dstack((-dz_dx, -dz_dy, np.ones_like(depth)))
# normalize to unit vector
normal_unit = normal_cross / np.linalg.norm(normal_cross, axis=2, keepdims=True)
# set default normal to [0, 0, 1]
# pdb.set_trace()
nan_mask = np.isfinite(normal_unit).all(2)
normal_unit[~np.isfinite(normal_unit).all(2)] = [-1, -1, -1]
return normal_unit, nan_mask
# def delete_holes(img):
# im = img.copy()
# h, w = im.shape[:2]
# mask = np.zeros((h+2, w+2), dtype=np.uint8)
# holes = cv2.floodFill(im.copy(), mask, (0, 0), 255)[1]
# holes = ~holes
# im[holes == 255] = 255
# return im
def denoise(img):
kernel = np.ones((11, 11), np.uint8)
img= cv2.erode(img, kernel, iterations=2)
img= cv2.dilate(img, kernel, iterations=2)
return img
def main():
args = parse_args()
root = args.root
cam_file = osp.join(root, args.cam_file)
depth_dir = osp.join(root, args.depth_dir)
out_dir = osp.join(root, args.out_dir)
os.makedirs(out_dir, exist_ok=True)
mask_dir = osp.join(root, args.mask_dir)
mask_total_dir = osp.join(root, args.mask_total_dir)
os.makedirs(mask_total_dir, exist_ok=True)
cam = np.load(cam_file)
fx, fy = cam['fx'], cam['fy']
assert osp.exists(depth_dir), 'depth dir not exists'
depth_paths = sorted([
osp.join(depth_dir, f) for f in os.listdir(depth_dir)])
mask_paths = sorted([osp.join(mask_dir, f) for f in os.listdir(mask_dir)])
assert depth_paths[0].endswith('.npy'), 'depth dir must be .npy'
# [-1,1] -> [0, 2]/2 -> [0,1]*255 -> [0,255]
vis_normal = lambda normal: np.uint8((normal + 1) / 2 * 255)[..., ::-1]
for idx, depth_p in enumerate(depth_paths):
# mask is [0, 1]
mask = cv2.imread(mask_paths[idx], -1)[..., 0] / 255
assert len(mask.shape) == 2, 'dims of mask are not 2'
depth = np.load(depth_p) # the unit of depth is meter
# the val of normal is between [-1, 1]
normal, nan_mask = get_surface_normal_by_depth(depth, fx, fy)
mask_total = nan_mask * mask
#mask_total = denoise(mask_total)
# normal should be filtered by mask and nan_mask
normal = vis_normal(normal) * mask_total[..., None]
cv2.imwrite(osp.join(out_dir, f'{idx:05d}.png'), normal)
cv2.imwrite(osp.join(mask_total_dir, f'{idx:05d}.png'), mask_total.astype(np.uint8) * 255)
#cv2.imwrite(osp.join(root, 'nan_mask.png'), nan_mask.astype(np.uint8))
#xxx
if __name__ == '__main__':
#os.makedirs(args.normalPath, exist_ok = True)
main()
''' if args.totalmaskPath:
os.makedirs(args.totalmaskPath, exist_ok = True)
for img_name in tqdm.tqdm(sorted(os.listdir(args.depthPath))):
img_path = os.path.join(args.depthPath, img_name)
depth = cv2.imread(img_path, -1) / 1000.0
normal, nan_mask = get_surface_normal_by_depth(depth)
vis_normal = lambda normal: np.uint8((normal + 1) / 2 * 255)[..., ::-1]
normal = vis_normal(normal)
if args.maskPath:
mask_path = os.path.join(args.maskPath, img_name[:-4] + '_mask' + img_name[-4:])
mask = cv2.imread(mask_path)
mask = cv2.rotate(mask, cv2.ROTATE_90_CLOCKWISE)
total_mask = mask[..., 0] * nan_mask
# cv2.imwrite(os.path.join(args.totalmaskPath, img_name), np.uint8(nan_mask*255))
# cv2.imwrite(os.path.join(args.totalmaskPath, img_name), total_mask)
kernel = np.ones((11, 11), np.uint8)
total_mask = cv2.erode(total_mask, kernel, iterations=1)
total_mask = cv2.dilate(total_mask, kernel, iterations=1)
# pdb.set_trace()
normal = np.uint8(total_mask[..., None]/255) * normal
if args.totalmaskPath:
# cv2.imwrite(os.path.join(args.totalmaskPath, img_name), delete_holes(total_mask))
total_mask = cv2.rotate(total_mask, cv2.ROTATE_90_COUNTERCLOCKWISE)
cv2.imwrite(os.path.join(args.totalmaskPath, img_name), total_mask)
# pdb.set_trace()
normal_path = os.path.join(args.normalPath, img_name)
normal = cv2.rotate(normal, cv2.ROTATE_90_COUNTERCLOCKWISE)
cv2.imwrite(normal_path, normal)
'''
|