| import numpy as np |
| import os |
| import argparse |
| import open3d as o3d |
| import glob |
| import cv2 |
| import copy |
|
|
| def get_roll_rot(angle): |
| ca=np.cos(angle) |
| sa=np.sin(angle) |
| rot=np.array([ |
| [ca,-sa,0,0], |
| [sa,ca,0,0], |
| [0,0,1,0], |
| [0,0,0,1] |
| ]) |
| return rot |
| def rotate_mat(direction): |
| if direction == 'Up': |
| return np.eye(4) |
| elif direction == 'Left': |
| rot_mat=get_roll_rot(np.pi/2) |
| elif direction == 'Right': |
| rot_mat=get_roll_rot(-np.pi/2) |
| elif direction == 'Down': |
| rot_mat=get_roll_rot(np.pi) |
| else: |
| raise Exception(f'No such direction (={direction}) rotation') |
| return rot_mat |
|
|
| def rotate_K(K,direction): |
| if direction == 'Up' or direction=="Down": |
| new_K4=np.eye(4) |
| new_K4[0:3,0:3]=copy.deepcopy(K) |
| return new_K4 |
| elif direction == 'Left' or direction =="Right": |
| fx,fy,cx,cy=K[0,0],K[1,1],K[0,2],K[1,2] |
| new_K4 = np.array([ |
| [fy, 0, cy, 0], |
| [0, fx, cx, 0], |
| [0, 0, 1, 0], |
| [0, 0, 0, 1] |
| ]) |
| return new_K4 |
|
|
| def rotate_bbox(bbox,direction, H,W): |
|
|
| x_min,y_min,x_max,y_max=bbox[0:4] |
| if direction == 'Up': |
| return bbox |
| elif direction == 'Left': |
| |
| new_bbox=[min(H-bbox[1],H-bbox[3]),bbox[0],max(H-bbox[1],H-bbox[3]),bbox[2]] |
| elif direction == 'Right': |
| new_bbox=[bbox[1],min(W-bbox[0],W-bbox[2]),bbox[3],max(W-bbox[0],W-bbox[2])] |
| elif direction == 'Down': |
| new_bbox=[min(W-x_min,W-x_max),min(H-y_min,H-y_max),max(W-x_min,W-x_max),max(H-y_min,H-y_max)] |
| else: |
| raise Exception(f'No such direction (={direction}) rotation') |
| return new_bbox |
|
|
| def rotate_image(img, direction): |
| if direction == 'Up': |
| pass |
| elif direction == 'Left': |
| img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE) |
| elif direction == 'Right': |
| img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE) |
| elif direction == 'Down': |
| img = cv2.rotate(img, cv2.ROTATE_180) |
| else: |
| raise Exception(f'No such direction (={direction}) rotation') |
| return img |
|
|
| parser=argparse.ArgumentParser() |
| parser.add_argument("--data_folder",type=str,required=True) |
| parser.add_argument("--save_dir",type=str,default=r"../example_process_data") |
| parser.add_argument("--debug",action="store_true",default=False) |
| args=parser.parse_args() |
|
|
| print("processing %s"%(args.data_folder)) |
|
|
| data_folder=args.data_folder |
| scene_name=os.path.basename(data_folder) |
| save_folder=os.path.join(args.save_dir,scene_name) |
| os.makedirs(save_folder,exist_ok=True) |
| color_folder=os.path.join(data_folder,"color") |
| depth_folder=os.path.join(data_folder,"depth") |
| pose_folder=os.path.join(data_folder,"pose") |
|
|
| print(color_folder) |
|
|
| color_list=glob.glob(color_folder+"/*.jpg") |
| image_id_list=[os.path.basename(item)[0:-4] for item in color_list] |
| image_id_list.sort() |
|
|
| bbox_path=os.path.join(data_folder,"objects.npy") |
| bboxes_dict=np.load(bbox_path,allow_pickle=True).item() |
|
|
| intrinsic_path=os.path.join(data_folder,"intrinsic","intrinsic_color.txt") |
| K=np.loadtxt(intrinsic_path) |
|
|
| align_path=os.path.join(data_folder,"alignment_matrix.txt") |
| align_matrix=np.loadtxt(align_path) |
| if align_matrix.shape[0]==3: |
| new_align_matrix=np.eye(4) |
| new_align_matrix[0:3,0:3]=align_matrix |
| align_matrix=new_align_matrix |
|
|
| mesh_path=os.path.join(data_folder,"fused_mesh.ply") |
| o3d_mesh=o3d.io.read_triangle_mesh(mesh_path) |
| o3d_vertices = np.array(o3d_mesh.vertices) |
| o3d_vert_homo=np.concatenate([o3d_vertices,np.ones([o3d_vertices.shape[0],1])],axis=1) |
| align_o3d_vertices = np.dot(o3d_vert_homo,align_matrix)[:,0:3] |
| o3d_mesh.vertices = o3d.utility.Vector3dVector(align_o3d_vertices) |
| align_mesh_save_path=os.path.join(save_folder,"align_mesh.ply") |
| o3d.io.write_triangle_mesh(align_mesh_save_path,o3d_mesh) |
|
|
| x=np.linspace(-1,1,10) |
| y=np.linspace(-1,1,10) |
| z=np.linspace(-1,1,10) |
| X,Y,Z=np.meshgrid(x,y,z,indexing='ij') |
| vox_coor=np.concatenate([X[:,:,:,np.newaxis],Y[:,:,:,np.newaxis],Z[:,:,:,np.newaxis]],axis=-1) |
| vox_coor=np.reshape(vox_coor,(-1,3)) |
| |
|
|
| pre_proj_mates={} |
| obj_points_dict={} |
| trans_mats={} |
| point_save_folder=os.path.join(save_folder,"5_partial_points") |
| os.makedirs(point_save_folder,exist_ok=True) |
| tran_save_folder=os.path.join(save_folder,"10_tran_matrix") |
| os.makedirs(tran_save_folder,exist_ok=True) |
| for object_id in bboxes_dict: |
| object = bboxes_dict[object_id] |
| category = object['category'] |
| sizes = object['size'] |
| sizes *= 1.1 |
| transform_matrix_t = np.array(object['transform']).reshape([4, 4]) |
| translate = transform_matrix_t[:3, 3] |
| rotation = transform_matrix_t[:3, :3] |
|
|
| bbox_o3d = o3d.geometry.OrientedBoundingBox(translate.reshape([3, 1]), |
| rotation, |
| np.array(sizes).reshape([3, 1])) |
| crop_pcd = o3d_mesh.crop(bbox_o3d) |
| crop_vert = np.asarray(crop_pcd.vertices) |
| org_crop_vert = crop_vert[:, :] |
| crop_vert = crop_vert - translate |
| crop_vert = np.dot(crop_vert,np.linalg.inv(rotation).T) |
| crop_vert[:, 2] *= -1 |
| bb_min, bb_max = np.amin(crop_vert, axis=0), np.amax(crop_vert, axis=0) |
| max_length = (bb_max - bb_min).max() |
| center = (bb_max + bb_min) / 2 |
| crop_vert = (crop_vert - center) / max_length * 2 |
|
|
| obj_points_dict[object_id]=crop_vert |
| crop_pcd.vertices=o3d.utility.Vector3dVector(crop_vert) |
| save_path=os.path.join(point_save_folder,category+"_%d.ply"%(object_id)) |
| o3d.io.write_triangle_mesh(save_path,crop_pcd) |
|
|
| proj_mat = np.eye(4) |
| scale_tran = np.eye(4) |
| scale_tran[0, 0], scale_tran[1, 1], scale_tran[2, 2] = max_length / 2, max_length / 2, max_length / 2 |
| proj_mat = np.dot(proj_mat, scale_tran) |
| center_tran = np.eye(4) |
| center_tran[0:3, 3] = center |
| proj_mat = np.dot(center_tran, proj_mat) |
| invert_mat = np.eye(4) |
| invert_mat[2, 2] *= -1 |
| proj_mat = np.dot(invert_mat, proj_mat) |
| proj_mat[0:3, 0:3] = np.dot(rotation,proj_mat[0:3, 0:3]) |
| translate_mat = np.eye(4) |
| translate_mat[0:3, 3] = translate |
| proj_mat = np.dot(translate_mat, proj_mat) |
|
|
| '''tran mat is to align output to scene space''' |
| tran_mat=copy.deepcopy(proj_mat) |
| trans_mats[object_id]=tran_mat |
| tran_save_path=os.path.join(tran_save_folder,category+"_%d.npy"%(object_id)) |
| np.save(tran_save_path,tran_mat) |
|
|
| unalign_mat = np.linalg.inv(align_matrix) |
| proj_mat = np.dot(unalign_mat.T, proj_mat) |
| pre_proj_mates[object_id]=proj_mat |
|
|
| ref=np.array([ |
| [0,1.0], |
| [-1.0,0], |
| [0,1.0], |
| [0.0,-1.0] |
| ]) |
| dir_list=[ |
| "Down", |
| "Left", |
| "Right", |
| "Up" |
| ] |
|
|
| for image_id in image_id_list: |
| color_path=os.path.join(color_folder,image_id+".jpg") |
| depth_path=os.path.join(depth_folder,image_id+".png") |
| pose_path=os.path.join(pose_folder,image_id+".txt") |
|
|
| color=cv2.imread(color_path) |
| height,width=color.shape[0:2] |
| depth=cv2.imread(depth_path,cv2.IMREAD_ANYCOLOR|cv2.IMREAD_ANYDEPTH)/1000.0 |
| pose=np.loadtxt(pose_path) |
| for object_id in bboxes_dict: |
| object=bboxes_dict[object_id] |
| category=object['category'] |
| sizes=object['size'] |
| object_vox_coor=vox_coor*sizes[np.newaxis,:] |
| |
| |
|
|
| prev_proj_mat=pre_proj_mates[object_id] |
| wrd2cam_pose = np.linalg.inv(pose) |
| current_proj_mat = np.dot(wrd2cam_pose, prev_proj_mat) |
| K4=np.eye(4) |
| K4[0:3,0:3]=K |
|
|
| '''calibrate proj_mat''' |
| up_vectors = np.array([[0, 0, 0, 1.0], |
| [0, 0.5, 0, 1.0]]) |
| up_vec_inimg = np.dot(up_vectors, current_proj_mat.T) |
| up_vec_inimg = np.dot(up_vec_inimg,K4.T) |
| up_x = up_vec_inimg[:, 0] / up_vec_inimg[:, 2] |
| up_y = up_vec_inimg[:, 1] / up_vec_inimg[:, 2] |
| pt1 = np.array((up_x[0], up_y[0])) |
| pt2 = np.array((up_x[1], up_y[1])) |
| up_dir = pt2 - pt1 |
| |
|
|
| product = np.sum(up_dir[np.newaxis, :] * ref, axis=1) |
| max_ind = np.argmax(product) |
| direction = dir_list[max_ind] |
| sky_rot = rotate_mat(direction) |
| |
|
|
| vox_homo=np.concatenate([object_vox_coor,np.ones((object_vox_coor.shape[0],1))],axis=1) |
| vox_proj=np.dot(vox_homo,current_proj_mat.T) |
| vox_proj=np.dot(vox_proj,K4.T) |
| vox_x=vox_proj[:,0]/vox_proj[:,2] |
| vox_y=vox_proj[:,1]/vox_proj[:,2] |
|
|
| if np.mean(vox_proj[:,2])>5: |
| continue |
|
|
| inside_mask=((vox_x<width-1) &(vox_x>0) &(vox_y<height-1) &(vox_y>0)).astype(np.float32) |
| infrustum_ratio=np.sum(inside_mask)/vox_x.shape[0] |
| if infrustum_ratio < 0.4 and category in ["chair", "stool"]: |
| continue |
| elif infrustum_ratio <0.2: |
| continue |
| |
|
|
| '''objects visibility check for every frame''' |
| vox_x_inside=vox_x[inside_mask>0].astype(np.int32) |
| vox_y_inside=vox_y[inside_mask>0].astype(np.int32) |
| vox_depth=vox_proj[inside_mask>0,2] |
| |
| depth_sample=depth[vox_y_inside,vox_x_inside] |
| depth_mask=(depth_sample>0)&(depth_sample<10.0) |
| depth_sample=depth_sample[depth_mask] |
| vox_depth=vox_depth[depth_mask] |
|
|
| if vox_depth.shape[0]<100: |
| continue |
|
|
| occluded_ratio=np.sum(((vox_depth-depth_sample)>0.2).astype(np.float32))/vox_depth.shape[0] |
| if occluded_ratio>0.6 and category in ["chair"]: |
| continue |
|
|
| depth_near_ratio = np.sum((np.abs(vox_depth - depth_sample) < sizes.max() * 0.5).astype(np.float32)) / \ |
| vox_depth.shape[0] |
| if depth_near_ratio < 0.2: |
| continue |
|
|
| '''make sure in every image, the object is upward''' |
| bbox=(np.amin(vox_x_inside),np.amin(vox_y_inside),np.amax(vox_x_inside),np.amax(vox_y_inside)) |
| rot_image=rotate_image(color,direction) |
| bbox = rotate_bbox(bbox, direction, height, width) |
| crop_image=rot_image[bbox[1]:bbox[3],bbox[0]:bbox[2]] |
| crop_h, crop_w = crop_image.shape[0:2] |
| max_length = max(crop_h, crop_w) |
| if max_length<100: |
| continue |
| pad_image = np.zeros((max_length, max_length, 3)) |
| if crop_h > crop_w: |
| margin = crop_h - crop_w |
| pad_image[:, margin // 2:margin // 2 + crop_w] = crop_image[:, :, :] |
| x_start, x_end = bbox[0] - margin // 2, margin // 2 + bbox[2] |
| y_start, y_end = bbox[1], bbox[3] |
| else: |
| margin = crop_w - crop_h |
| pad_image[margin // 2:margin // 2 + crop_h, :] = crop_image[:, :, :] |
|
|
| y_start, y_end = bbox[1] - margin // 2, bbox[3] + margin // 2 |
| x_start, x_end = bbox[0], bbox[2] |
|
|
| pad_image=cv2.resize(pad_image,dsize=(224,224),interpolation=cv2.INTER_LINEAR) |
| image_save_folder = os.path.join(save_folder, "6_images", category + "_%d" % (object_id)) |
| os.makedirs(image_save_folder, exist_ok=True) |
| image_save_path=os.path.join(image_save_folder,image_id+".jpg") |
| |
| cv2.imwrite(image_save_path,pad_image) |
|
|
| proj_mat=np.dot(sky_rot,current_proj_mat) |
| new_K4 = rotate_K(K, direction) |
| new_K4[0, 2] -= x_start |
| new_K4[1, 2] -= y_start |
| new_K4[0] = new_K4[0] / max_length * 224 |
| new_K4[1] = new_K4[1] / max_length * 224 |
| proj_mat = np.dot(new_K4, proj_mat) |
|
|
| proj_save_folder=os.path.join(save_folder,"8_proj_matrix",category+"_%d"%(object_id)) |
| os.makedirs(proj_save_folder,exist_ok=True) |
| proj_save_path=os.path.join(proj_save_folder,image_id+".npy") |
| np.save(proj_save_path,proj_mat) |
|
|
| '''debug proj matrix''' |
| if args.debug: |
| proj_save_folder=os.path.join(save_folder,"9_proj_images",category+"_%d"%(object_id)) |
| os.makedirs(proj_save_folder,exist_ok=True) |
| canvas=copy.deepcopy(pad_image) |
| par_points=obj_points_dict[object_id] |
| par_homo=np.concatenate([par_points,np.ones((par_points.shape[0],1))],axis=1) |
| par_inimg=np.dot(par_homo,proj_mat.T) |
| x=par_inimg[:,0]/par_inimg[:,2] |
| y=par_inimg[:,1]/par_inimg[:,2] |
| x=np.clip(x,a_min=0,a_max=223).astype(np.int32) |
| y=np.clip(y,a_min=0,a_max=223).astype(np.int32) |
| canvas[y,x]=np.array([[0,255,0]]) |
| proj_save_path=os.path.join(proj_save_folder,image_id+".jpg") |
| cv2.imwrite(proj_save_path,canvas) |
|
|
|
|
|
|
|
|
