show_3d renderer refactoring
This commit is contained in:
parent
be52513c96
commit
deccff302d
|
@ -2,6 +2,9 @@
|
|||
|
||||
This folder has the source code for environment bundle, which includes viewer renderer, multichannel renderer and physics simulation engine
|
||||
|
||||
### Requirement
|
||||
You need CUDA 8.0 in order to run the environment.
|
||||
|
||||
### Run everything at once
|
||||
Open three different shells, and follow the sequence
|
||||
```shell
|
|
@ -0,0 +1 @@
|
|||
transfer.c
|
|
@ -9,12 +9,14 @@ Install opencv, pytorch and others
|
|||
conda install -c menpo opencv -y
|
||||
conda install pytorch torchvision cuda80 -c soumith
|
||||
conda install progressbar pillow
|
||||
pip install zmq
|
||||
```
|
||||
|
||||
Build cython dependencies
|
||||
Build cpp, cython dependencies
|
||||
```shell
|
||||
bash build.sh
|
||||
bash build_cuda.sh
|
||||
python setup.py build_ext --inplace
|
||||
```
|
||||
|
||||
Download the necessary helper files
|
|
@ -0,0 +1,430 @@
|
|||
from __future__ import print_function
|
||||
import numpy as np
|
||||
import ctypes as ct
|
||||
import cv2
|
||||
import sys
|
||||
import torch
|
||||
import argparse
|
||||
import time
|
||||
import utils
|
||||
import transforms3d
|
||||
import json
|
||||
import zmq
|
||||
|
||||
from torchvision import datasets, transforms
|
||||
from torch.autograd import Variable
|
||||
from numpy import cos, sin
|
||||
from profiler import Profiler
|
||||
from multiprocessing.dummy import Process
|
||||
|
||||
from datasets import ViewDataSet3D
|
||||
from completion import CompletionNet
|
||||
|
||||
cuda_pc = np.ctypeslib.load_library('render_cuda_f','.')
|
||||
coords = np.load('coord.npy')
|
||||
|
||||
class InImg(object):
|
||||
def __init__(self):
|
||||
self.grid = 768
|
||||
|
||||
def getpixel(self, key):
|
||||
corrx, corry = key[0], key[1]
|
||||
|
||||
indx = int(corrx / self.grid)
|
||||
indy = int(corry / self.grid)
|
||||
|
||||
remx = int(corrx % self.grid)
|
||||
remy = int(corry % self.grid)
|
||||
|
||||
if (indy == 0):
|
||||
return (0, remx, remy)
|
||||
elif (indy == 2):
|
||||
return (5, remx, remy)
|
||||
else:
|
||||
return (indx + 1, remx, remy)
|
||||
|
||||
class ViewRenderer(object):
|
||||
def __init__(self):
|
||||
return
|
||||
|
||||
|
||||
class PCRenderer:
|
||||
def __init__(self, port):
|
||||
self.roll, self.pitch, self.yaw = 0, 0, 0
|
||||
self.quat = [1, 0, 0, 0]
|
||||
self.x, self.y, self.z = 0, 0, 0
|
||||
self.fps = 0
|
||||
self.mousex, self.mousey = 0.5, 0.5
|
||||
#self.changed = True
|
||||
self.org_pitch, self.org_yaw, self.org_roll = 0, 0, 0
|
||||
self.org_x, self.org_y, self.org_z = 0, 0, 0
|
||||
self.clickstart = (0,0)
|
||||
self.mousedown = False
|
||||
self.fps = 0
|
||||
self.rotation_const = np.array([[0,1,0,0],[0,0,1,0],[-1,0,0,0],[0,0,0,1]])
|
||||
self.overlay = False
|
||||
self.show_depth = False
|
||||
self._context_phys = zmq.Context()
|
||||
self.socket_phys = self._context_phys.socket(zmq.REP)
|
||||
self.socket_phys.connect("tcp://localhost:%d" % port)
|
||||
|
||||
def _onmouse(self, *args):
|
||||
if args[0] == cv2.EVENT_LBUTTONDOWN:
|
||||
self.org_pitch, self.org_yaw, self.org_x, self.org_y, self.org_z =\
|
||||
self.pitch,self.yaw,self.x,self.y,self.z
|
||||
self.clickstart = (self.mousex, self.mousey)
|
||||
|
||||
if args[0] == cv2.EVENT_RBUTTONDOWN:
|
||||
self.org_roll = self.roll
|
||||
self.clickstart = (self.mousex, self.mousey)
|
||||
|
||||
if (args[3] & cv2.EVENT_FLAG_LBUTTON):
|
||||
self.pitch = self.org_pitch + (self.mousex - self.clickstart[0])/10
|
||||
self.yaw = self.org_yaw + (self.mousey - self.clickstart[1])
|
||||
#self.changed=True
|
||||
|
||||
if (args[3] & cv2.EVENT_FLAG_RBUTTON):
|
||||
self.roll = self.org_roll + (self.mousex - self.clickstart[0])/50
|
||||
#self.changed=True
|
||||
|
||||
my=args[1]
|
||||
mx=args[2]
|
||||
self.mousex=mx/float(256)
|
||||
self.mousey=my/float(256 * 2)
|
||||
|
||||
def _updateStateFromKeyboard(self):
|
||||
cmd=cv2.waitKey(5)%256
|
||||
if cmd==ord('q'):
|
||||
return False
|
||||
elif cmd == ord('w'):
|
||||
self.x -= 0.05
|
||||
#self.changed = True
|
||||
elif cmd == ord('s'):
|
||||
self.x += 0.05
|
||||
#self.changed = True
|
||||
elif cmd == ord('a'):
|
||||
self.y += 0.05
|
||||
#self.changed = True
|
||||
elif cmd == ord('d'):
|
||||
self.y -= 0.05
|
||||
#self.changed = True
|
||||
elif cmd == ord('z'):
|
||||
self.z += 0.01
|
||||
#self.changed = True
|
||||
elif cmd == ord('x'):
|
||||
self.z -= 0.01
|
||||
#self.changed = True
|
||||
elif cmd == ord('r'):
|
||||
self.pitch,self.yaw,self.x,self.y,self.z = 0,0,0,0,0
|
||||
self.roll = 0
|
||||
#self.changed = True
|
||||
elif cmd == ord('t'):
|
||||
pose = poses[0]
|
||||
RT = pose.reshape((4,4))
|
||||
R = RT[:3,:3]
|
||||
T = RT[:3,-1]
|
||||
self.x,self.y,self.z = np.dot(np.linalg.inv(R),T)
|
||||
self.roll, self.pitch, self.yaw = (utils.rotationMatrixToEulerAngles(R))
|
||||
#self.changed = True
|
||||
elif cmd == ord('o'):
|
||||
self.overlay = not self.overlay
|
||||
elif cmd == ord('f'):
|
||||
self.show_depth = not self.show_depth
|
||||
elif cmd == ord('v'):
|
||||
cv2.imwrite('save.jpg', show_rgb)
|
||||
return True
|
||||
|
||||
def _getPoseOrientationFromPhysics(self):
|
||||
receive = self.socket_phys.recv().decode("utf-8")
|
||||
new_pos, new_quat = json.loads(receive)
|
||||
self.socket_phys.send(json.dumps({"received": True}))
|
||||
return new_pos, new_quat
|
||||
|
||||
def _getNewPoseFromPhysics(self, view_pose):
|
||||
'''Return pos(xyz), quat(wxyz)
|
||||
'''
|
||||
view_pose['quat'] = utils.quat_wxyz_to_xyzw(view_pose['quat']).tolist()
|
||||
self.socket_phys.send(json.dumps(view_pose))
|
||||
new_pos, new_quat = json.loads(self.socket_phys.recv().decode("utf-8"))
|
||||
return new_pos, utils.quat_xyzw_to_wxyz(new_quat)
|
||||
|
||||
def _sendInitialPoseToPhysics(self, pose):
|
||||
pose[1] = utils.quat_wxyz_to_xyzw(pose[1]).tolist()
|
||||
receive = str(self.socket_phys.recv().decode("utf-8"))
|
||||
if (receive == "Initial"):
|
||||
self.socket_phys.send(json.dumps(pose))
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
def _getViewerRelativePose(self):
|
||||
cpose = np.eye(4)
|
||||
gamma = self.yaw
|
||||
alpha = self.pitch
|
||||
beta = -self.roll
|
||||
#cpose[:3, :3] = transforms3d.euler.euler2mat(alpha, beta, gamma)
|
||||
cpose[:3, :3] = transforms3d.quaternions.quat2mat(self.quat)
|
||||
cpose[ 0, -1] = self.x
|
||||
cpose[ 1, -1] = self.y
|
||||
cpose[ 2, -1] = self.z
|
||||
return cpose
|
||||
|
||||
def _getViewerAbsolutePose(self, target_pose):
|
||||
v_cam2world = target_pose
|
||||
v_cam2cam = self._getViewerRelativePose()
|
||||
p = v_cam2world.dot(np.linalg.inv(v_cam2cam))
|
||||
p = p.dot(np.linalg.inv(self.rotation_const))
|
||||
pos = utils.mat_to_posi_xyz(p)
|
||||
quat_wxyz = utils.quat_xyzw_to_wxyz(utils.mat_to_quat_xyzw(p))
|
||||
return pos, quat_wxyz
|
||||
|
||||
|
||||
def render(self, imgs, depths, pose, model, poses, target_pose, show, target_depth, opengl_arr):
|
||||
t0 = time.time()
|
||||
|
||||
v_cam2world = target_pose
|
||||
p = (v_cam2world).dot(np.linalg.inv(pose))
|
||||
p = p.dot(np.linalg.inv(self.rotation_const))
|
||||
s = utils.mat_to_str(p)
|
||||
|
||||
with Profiler("Depth request round-trip"):
|
||||
socket_mist.send(s)
|
||||
message = socket_mist.recv()
|
||||
|
||||
with Profiler("Read from framebuffer and make pano"):
|
||||
wo, ho = 768 * 4, 768 * 3
|
||||
|
||||
# Calculate height and width of output image, and size of each square face
|
||||
h = wo/3
|
||||
w = 2*h
|
||||
n = ho/3
|
||||
opengl_arr = np.array(np.frombuffer(message, dtype=np.float32)).reshape((h, w))
|
||||
|
||||
|
||||
def _render_depth(opengl_arr):
|
||||
with Profiler("Render Depth"):
|
||||
cv2.imshow('target depth', opengl_arr/16.)
|
||||
|
||||
def _render_pc(opengl_arr):
|
||||
with Profiler("Render pointcloud"):
|
||||
scale = 100. # 512
|
||||
target_depth = np.int32(opengl_arr * scale)
|
||||
show[:] = 0
|
||||
poses_after = [
|
||||
pose.dot(np.linalg.inv(poses[i])).astype(np.float32)
|
||||
for i in range(len(imgs))]
|
||||
|
||||
for i in range(len(imgs)):
|
||||
cuda_pc.render(ct.c_int(imgs[i].shape[0]),
|
||||
ct.c_int(imgs[i].shape[1]),
|
||||
imgs[i].ctypes.data_as(ct.c_void_p),
|
||||
depths[i].ctypes.data_as(ct.c_void_p),
|
||||
poses_after[i].ctypes.data_as(ct.c_void_p),
|
||||
show.ctypes.data_as(ct.c_void_p),
|
||||
target_depth.ctypes.data_as(ct.c_void_p)
|
||||
)
|
||||
threads = [
|
||||
Process(target=_render_pc, args=(opengl_arr,)),
|
||||
Process(target=_render_depth, args=(opengl_arr,))]
|
||||
[t.start() for t in threads]
|
||||
[t.join() for t in threads]
|
||||
|
||||
if model:
|
||||
tf = transforms.ToTensor()
|
||||
before = time.time()
|
||||
source = tf(show)
|
||||
source_depth = tf(np.expand_dims(target_depth, 2).astype(np.float32)/65536 * 255)
|
||||
imgv.data.copy_(source)
|
||||
maskv.data.copy_(source_depth)
|
||||
print('Transfer time', time.time() - before)
|
||||
before = time.time()
|
||||
recon = model(imgv, maskv)
|
||||
print('NNtime:', time.time() - before)
|
||||
before = time.time()
|
||||
show2 = recon.data.cpu().numpy()[0].transpose(1,2,0)
|
||||
show[:] = (show2[:] * 255).astype(np.uint8)
|
||||
print('Transfer to CPU time:', time.time() - before)
|
||||
|
||||
t1 =time.time()
|
||||
t = t1-t0
|
||||
self.fps = 1/t
|
||||
|
||||
def renderOffScreen(self):
|
||||
return
|
||||
|
||||
def renderToScreen(self, imgs, depths, poses, model, target, tdepth, target_poses):
|
||||
cv2.namedWindow('show3d')
|
||||
cv2.namedWindow('target depth')
|
||||
cv2.moveWindow('show3d',0,0)
|
||||
cv2.setMouseCallback('show3d',self._onmouse)
|
||||
|
||||
showsz = target.shape[0]
|
||||
show = np.zeros((showsz,showsz * 2,3),dtype='uint8')
|
||||
target_depth = np.zeros((showsz,showsz * 2)).astype(np.int32)
|
||||
imgv = Variable(torch.zeros(1,3, showsz, showsz*2), volatile=True).cuda()
|
||||
maskv = Variable(torch.zeros(1,1, showsz, showsz*2), volatile=True).cuda()
|
||||
|
||||
## TODO: error handling
|
||||
pos, quat_wxyz = self._getViewerAbsolutePose(target_poses[0])
|
||||
pos = pos.tolist()
|
||||
quat_wxyz = quat_wxyz.tolist()
|
||||
assert(self._sendInitialPoseToPhysics([pos, quat_wxyz]))
|
||||
|
||||
while True:
|
||||
'''
|
||||
v_cam2world = target_poses[0]
|
||||
v_cam2cam = self._getViewerRelativePose()
|
||||
p = v_cam2world.dot(np.linalg.inv(v_cam2cam))
|
||||
p = p.dot(np.linalg.inv(self.rotation_const))
|
||||
pos = utils.mat_to_posi_xyz(p).tolist()
|
||||
quat = utils.mat_to_quat_xyzw(p).tolist()
|
||||
'''
|
||||
|
||||
#world_cpose =
|
||||
|
||||
## Query physics engine to get [x, y, z, roll, pitch, yaw]
|
||||
print("waiting for physics")
|
||||
new_pos, new_quat = self._getPoseOrientationFromPhysics()
|
||||
self.x, self.y, self.z = new_pos
|
||||
self.quat = new_quat
|
||||
|
||||
v_cam2world = target_poses[0]
|
||||
v_cam2cam = self._getViewerRelativePose()
|
||||
cpose = np.linalg.inv(np.linalg.inv(v_cam2world).dot(v_cam2cam).dot(self.rotation_const))
|
||||
|
||||
'''
|
||||
pos, quat_wxyz = self._getViewerAbsolutePose(target_poses[0])
|
||||
pos = pos.tolist()
|
||||
quat_wxyz = utils.quat_wxyz_to_xyzw(quat_wxyz).tolist()
|
||||
'''
|
||||
|
||||
#self.changed = True
|
||||
|
||||
#if self.changed:
|
||||
#new_quat = utils.mat_to_quat_xyzw(world_cpose)
|
||||
#new_quat = utils.z_up_to_y_up(utils.quat_xyzw_to_wxyz(new_quat))
|
||||
#new_posi = utils.mat_to_posi_xyz(world_cpose)
|
||||
|
||||
## Entry point for change of view
|
||||
## Optimization
|
||||
depth_buffer = np.zeros(imgs[0].shape[:2], dtype=np.float32)
|
||||
|
||||
relative_poses = np.copy(target_poses)
|
||||
for i in range(len(relative_poses)):
|
||||
relative_poses[i] = np.dot(np.linalg.inv(relative_poses[i]), target_poses[0])
|
||||
|
||||
poses_after = [cpose.dot(np.linalg.inv(relative_poses[i])).astype(np.float32) for i in range(len(imgs))]
|
||||
pose_after_distance = [np.linalg.norm(rt[:3,-1]) for rt in poses_after]
|
||||
|
||||
top5 = (np.argsort(pose_after_distance))[:5]
|
||||
imgs_top5 = [imgs[i] for i in top5]
|
||||
depths_top5 = [depths[i] for i in top5]
|
||||
relative_poses_top5 = [relative_poses[i] for i in top5]
|
||||
|
||||
|
||||
self.render(imgs_top5, depths_top5, cpose.astype(np.float32), model, relative_poses_top5, target_poses[0], show, target_depth, depth_buffer)
|
||||
|
||||
#render(imgs, depths, cpose.astype(np.float32), model, poses)
|
||||
#self.changed = False
|
||||
|
||||
if self.overlay:
|
||||
show_out = (show/2 + target/2).astype(np.uint8)
|
||||
elif self.show_depth:
|
||||
show_out = (target_depth * 10).astype(np.uint8)
|
||||
else:
|
||||
show_out = show
|
||||
|
||||
|
||||
|
||||
cv2.putText(show,'pitch %.3f yaw %.2f roll %.3f x %.2f y %.2f z %.2f'%(self.pitch, self.yaw, self.roll, self.x, self.y, self.z),(15,showsz-15),0,0.5,(255,255,255))
|
||||
cv2.putText(show,'fps %.1f'%(self.fps),(15,15),0,0.5,(255,255,255))
|
||||
|
||||
show_rgb = cv2.cvtColor(show_out, cv2.COLOR_BGR2RGB)
|
||||
cv2.imshow('show3d',show_rgb)
|
||||
|
||||
if (not self._updateStateFromKeyboard()):
|
||||
break
|
||||
|
||||
|
||||
def show_target(target_img):
|
||||
cv2.namedWindow('target')
|
||||
cv2.moveWindow('target',0,256 + 50)
|
||||
show_rgb = cv2.cvtColor(target_img, cv2.COLOR_BGR2RGB)
|
||||
cv2.imshow('target', show_rgb)
|
||||
|
||||
if __name__=='__main__':
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--debug' , action='store_true', help='debug mode')
|
||||
parser.add_argument('--datapath' , required = True, help='dataset path')
|
||||
parser.add_argument('--model_id' , type = str, default = 0, help='model id')
|
||||
parser.add_argument('--model' , type = str, default = '', help='path of model')
|
||||
|
||||
opt = parser.parse_args()
|
||||
d = ViewDataSet3D(root=opt.datapath, transform = np.array, mist_transform = np.array, seqlen = 2, off_3d = False, train = False)
|
||||
|
||||
scene_dict = dict(zip(d.scenes, range(len(d.scenes))))
|
||||
if not opt.model_id in scene_dict.keys():
|
||||
print("model not found")
|
||||
else:
|
||||
scene_id = scene_dict[opt.model_id]
|
||||
|
||||
uuids, rts = d.get_scene_info(scene_id)
|
||||
print(uuids, rts)
|
||||
|
||||
targets = []
|
||||
sources = []
|
||||
source_depths = []
|
||||
poses = []
|
||||
|
||||
for k,v in uuids:
|
||||
print(k,v)
|
||||
data = d[v]
|
||||
source = data[0][0]
|
||||
target = data[1]
|
||||
target_depth = data[3]
|
||||
source_depth = data[2][0]
|
||||
pose = data[-1][0].numpy()
|
||||
targets.append(target)
|
||||
poses.append(pose)
|
||||
sources.append(target)
|
||||
source_depths.append(target_depth)
|
||||
|
||||
model = None
|
||||
if opt.model != '':
|
||||
comp = CompletionNet()
|
||||
comp = torch.nn.DataParallel(comp).cuda()
|
||||
comp.load_state_dict(torch.load(opt.model))
|
||||
model = comp.module
|
||||
model.eval()
|
||||
print(model)
|
||||
print('target', poses, poses[0])
|
||||
#print('no.1 pose', poses, poses[1])
|
||||
# print(source_depth)
|
||||
print(sources[0].shape, source_depths[0].shape)
|
||||
|
||||
context_mist = zmq.Context()
|
||||
print("Connecting to hello world server...")
|
||||
|
||||
socket_mist = context_mist.socket(zmq.REQ)
|
||||
socket_mist.connect("tcp://localhost:5555")
|
||||
|
||||
print(coords.flatten().dtype)
|
||||
with Profiler("Transform coords"):
|
||||
new_coords = np.getbuffer(coords.flatten().astype(np.uint32))
|
||||
print(coords.shape)
|
||||
print("Count: ", coords.flatten().astype(np.uint32).size )
|
||||
print("elem [2,3,5]: ", coords[4][2][1] )
|
||||
socket_mist.send(new_coords)
|
||||
print("Sent reordering")
|
||||
message = socket_mist.recv()
|
||||
print("msg")
|
||||
|
||||
|
||||
show_target(target)
|
||||
|
||||
renderer = PCRenderer(5556)
|
||||
renderer.renderToScreen(sources, source_depths, poses, model, target, target_depth, rts)
|
||||
|
||||
|
||||
|
|
@ -6,8 +6,9 @@
|
|||
import numpy as np
|
||||
import math
|
||||
import PIL
|
||||
import transforms3d
|
||||
from numpy import cos,sin
|
||||
|
||||
from transfer import transfer2
|
||||
|
||||
# In[36]:
|
||||
|
||||
|
@ -163,4 +164,118 @@ def generate_transformation_matrix(x,y,z,yaw,pitch,roll):
|
|||
current_rt = cpose
|
||||
rotation = np.array([[0,-1,0,0],[-1,0,0,0],[0,0,1,0],[0,0,0,1]])
|
||||
current_rt = np.dot(rotation, current_rt)
|
||||
return current_rt
|
||||
return current_rt
|
||||
|
||||
|
||||
|
||||
|
||||
def mat_to_str(matrix):
|
||||
s = ""
|
||||
for row in range(4):
|
||||
for col in range(4):
|
||||
s = s + " " + str(matrix[row][col])
|
||||
return s.strip()
|
||||
|
||||
|
||||
def mat_to_posi_xyz(cpose):
|
||||
return cpose[:3, -1]
|
||||
|
||||
def mat_to_quat_xyzw(cpose):
|
||||
rot = cpose[:3, :3]
|
||||
## Return: [r_x, r_y, r_z, r_w]
|
||||
wxyz = transforms3d.quaternions.mat2quat(rot)
|
||||
return quat_wxyz_to_xyzw(wxyz)
|
||||
|
||||
|
||||
def convert_array(img_array, outimg):
|
||||
inimg = InImg()
|
||||
|
||||
wo, ho = inimg.grid * 4, inimg.grid * 3
|
||||
|
||||
# Calculate height and width of output image, and size of each square face
|
||||
h = wo/3
|
||||
w = 2*h
|
||||
n = ho/3
|
||||
|
||||
# Create new image with width w, and height h
|
||||
# outimg = np.zeros((h,w,1)) #.astype(np.uint8)
|
||||
'''
|
||||
# PHYSICS
|
||||
outimg = np.zeros((h,w,1)) #.astype(np.uint8)
|
||||
|
||||
in_imgs = None
|
||||
#print("converting images", len(img_array))
|
||||
|
||||
#print("Passed in image array", len(img_array), np.max(img_array[0]))
|
||||
in_imgs = img_array
|
||||
|
||||
# For each pixel in output image find colour value from input image
|
||||
#print(outimg.shape)
|
||||
'''
|
||||
|
||||
# todo: for some reason the image is flipped 180 degrees
|
||||
transfer2(img_array, coords, h, w, outimg)
|
||||
|
||||
# return outimg
|
||||
|
||||
|
||||
|
||||
## Blender generated poses and OpenGL default poses adhere to
|
||||
## different conventions. To better understand the nitty-gritty
|
||||
## transformations inside this file, check out this:
|
||||
## https://en.wikibooks.org/wiki/OpenGL_Programming/Modern_OpenGL_Tutorial_Navigation
|
||||
## Blender: z-is-up
|
||||
## same with: csv, rt_camera_matrix
|
||||
## OpenGL: y-is-up
|
||||
## same with: obj
|
||||
## Default camera: y is up-direction, -z facing
|
||||
|
||||
|
||||
## Quat(wxyz)
|
||||
def quat_pos_to_mat(pos, quat):
|
||||
r_w, r_x, r_y, r_z = quat
|
||||
#print("quat", r_w, r_x, r_y, r_z)
|
||||
mat = np.eye(4)
|
||||
mat[:3, :3] = transforms3d.quaternions.quat2mat([r_w, r_x, r_y, r_z])
|
||||
mat[:3, -1] = pos
|
||||
# Return: roll, pitch, yaw
|
||||
return mat
|
||||
|
||||
## Used for URDF models that are default -x facing
|
||||
## Rotate the model around its internal x axis for 90 degrees
|
||||
## so that it is at "normal" pose when applied camera_rt_matrix
|
||||
## Format: wxyz for input & return
|
||||
def z_up_to_y_up(quat_wxyz):
|
||||
## Operations (1) rotate around y for pi/2,
|
||||
## (2) rotate around z for pi/2
|
||||
to_y_up = transforms3d.euler.euler2quat(0, np.pi/2, np.pi/2)
|
||||
return transforms3d.quaternions.qmult(quat_wxyz, to_y_up)
|
||||
|
||||
## Models coming out of opengl are negative x facing
|
||||
## Transform the default to
|
||||
def y_up_to_z_up(quat_wxyz):
|
||||
to_z_up = transforms3d.euler.euler2quat(np.pi/2, 0, np.pi/2)
|
||||
return transforms3d.quaternions.qmult(to_z_up, quat_wxyz)
|
||||
|
||||
|
||||
def quat_wxyz_to_euler(wxyz):
|
||||
q0, q1, q2, q3 = wxyz
|
||||
sinr = 2 * (q0 * q1 + q2 * q3)
|
||||
cosr = 1 - 2 * (q1 * q1 + q2 * q2)
|
||||
sinp = 2 * (q0 * q2 - q3 * q1)
|
||||
siny = 2 * (q0 * q3 + q1 * q2)
|
||||
cosy = 1 - 2 * (q2 * q2 + q3 * q3)
|
||||
|
||||
roll = np.arctan2(sinr, cosr)
|
||||
pitch = np.arcsin(sinp)
|
||||
yaw = np.arctan2(siny, cosy)
|
||||
return [roll, pitch, yaw]
|
||||
|
||||
|
||||
## wxyz: numpy array format
|
||||
def quat_wxyz_to_xyzw(wxyz):
|
||||
return np.concatenate((wxyz[1:], wxyz[:1]))
|
||||
|
||||
## xyzw: numpy array format
|
||||
def quat_xyzw_to_wxyz(xyzw):
|
||||
return np.concatenate((xyzw[-1:], xyzw[:-1]))
|
|
@ -0,0 +1,8 @@
|
|||
from setuptools import setup
|
||||
|
||||
setup(name='realenv_env',
|
||||
version='0.1.0',
|
||||
install_requires=['gym',
|
||||
'pybullet',
|
||||
'zmq',
|
||||
'transforms3d'])
|
|
@ -0,0 +1,39 @@
|
|||
import gym
|
||||
from gym import error, spaces, utils
|
||||
from gym.utils import seeding
|
||||
import subprocess
|
||||
|
||||
class SimpleEnv(gym.Env):
|
||||
metadata = {'render.modes': ['human']}
|
||||
|
||||
def __init__(self):
|
||||
cmd_channel = "./channels/depth_render/depth_render --datapath data -m 11HB6XZSh1Q"
|
||||
cmd_physics = "python show_3d2.py --datapath ../data/ --idx 10"
|
||||
cmd_render = ""
|
||||
|
||||
#self.p_channel = subprocess.Popen(cmd_channel.split(), stdout=subprocess.PIPE)
|
||||
#self.p_physics = subprocess.Popen()
|
||||
#self.p_render = subprocess.Popen()
|
||||
|
||||
def testShow3D(self):
|
||||
|
||||
|
||||
def _step(self, action):
|
||||
return
|
||||
|
||||
def _reset(self):
|
||||
return
|
||||
|
||||
def _render(self, mode='human', close=False):
|
||||
return
|
||||
|
||||
def _end(self):
|
||||
#self.p_channel.kill()
|
||||
#self.p_physics.kill()
|
||||
#self.p_render.kill()
|
||||
return
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
env = SimpleEnv()
|
||||
env.
|
|
@ -1,717 +0,0 @@
|
|||
## Blender generated poses and OpenGL default poses adhere to
|
||||
## different conventions. To better understand the nitty-gritty
|
||||
## transformations inside this file, check out this:
|
||||
## https://en.wikibooks.org/wiki/OpenGL_Programming/Modern_OpenGL_Tutorial_Navigation
|
||||
## Blender: z-is-up
|
||||
## same with: csv, rt_camera_matrix
|
||||
## OpenGL: y-is-up
|
||||
## same with: obj
|
||||
## Default camera: y is up-direction, -z facing
|
||||
|
||||
|
||||
from __future__ import print_function
|
||||
import numpy as np
|
||||
import ctypes as ct
|
||||
import cv2
|
||||
import sys
|
||||
import argparse
|
||||
from datasets import ViewDataSet3D
|
||||
from completion import CompletionNet
|
||||
import torch
|
||||
from torchvision import datasets, transforms
|
||||
from torch.autograd import Variable
|
||||
import time
|
||||
from numpy import cos, sin
|
||||
import utils
|
||||
import transforms3d
|
||||
import json
|
||||
|
||||
import zmq
|
||||
from transfer import transfer2
|
||||
|
||||
from profiler import Profiler
|
||||
from multiprocessing.dummy import Process
|
||||
|
||||
PHYSICS_FIRST = True
|
||||
|
||||
class InImg(object):
|
||||
def __init__(self):
|
||||
self.grid = 768
|
||||
|
||||
def getpixel(self, key):
|
||||
corrx, corry = key[0], key[1]
|
||||
|
||||
indx = int(corrx / self.grid)
|
||||
indy = int(corry / self.grid)
|
||||
|
||||
remx = int(corrx % self.grid)
|
||||
remy = int(corry % self.grid)
|
||||
|
||||
if (indy == 0):
|
||||
return (0, remx, remy)
|
||||
elif (indy == 2):
|
||||
return (5, remx, remy)
|
||||
else:
|
||||
return (indx + 1, remx, remy)
|
||||
|
||||
|
||||
dll=np.ctypeslib.load_library('render_cuda_f','.')
|
||||
|
||||
|
||||
|
||||
def mat_to_str(matrix):
|
||||
s = ""
|
||||
for row in range(4):
|
||||
for col in range(4):
|
||||
s = s + " " + str(matrix[row][col])
|
||||
return s.strip()
|
||||
|
||||
coords = np.load('coord.npy')
|
||||
|
||||
def convert_array(img_array, outimg):
|
||||
inimg = InImg()
|
||||
|
||||
wo, ho = inimg.grid * 4, inimg.grid * 3
|
||||
|
||||
# Calculate height and width of output image, and size of each square face
|
||||
h = wo/3
|
||||
w = 2*h
|
||||
n = ho/3
|
||||
|
||||
# Create new image with width w, and height h
|
||||
# outimg = np.zeros((h,w,1)) #.astype(np.uint8)
|
||||
'''
|
||||
# PHYSICS
|
||||
outimg = np.zeros((h,w,1)) #.astype(np.uint8)
|
||||
|
||||
in_imgs = None
|
||||
#print("converting images", len(img_array))
|
||||
|
||||
#print("Passed in image array", len(img_array), np.max(img_array[0]))
|
||||
in_imgs = img_array
|
||||
|
||||
# For each pixel in output image find colour value from input image
|
||||
#print(outimg.shape)
|
||||
'''
|
||||
|
||||
# todo: for some reason the image is flipped 180 degrees
|
||||
transfer2(img_array, coords, h, w, outimg)
|
||||
|
||||
# return outimg
|
||||
|
||||
|
||||
|
||||
def mat_to_posi_xyz(cpose):
|
||||
return cpose[:3, -1]
|
||||
|
||||
def mat_to_quat_xyzw(cpose):
|
||||
rot = cpose[:3, :3]
|
||||
## Return: [r_x, r_y, r_z, r_w]
|
||||
wxyz = transforms3d.quaternions.mat2quat(rot)
|
||||
return quat_wxyz_to_xyzw(wxyz)
|
||||
|
||||
## Quat(wxyz)
|
||||
def quat_pos_to_mat(pos, quat):
|
||||
r_w, r_x, r_y, r_z = quat
|
||||
#print("quat", r_w, r_x, r_y, r_z)
|
||||
mat = np.eye(4)
|
||||
mat[:3, :3] = transforms3d.quaternions.quat2mat([r_w, r_x, r_y, r_z])
|
||||
mat[:3, -1] = pos
|
||||
# Return: roll, pitch, yaw
|
||||
return mat
|
||||
|
||||
## Used for URDF models that are default -x facing
|
||||
## Rotate the model around its internal x axis for 90 degrees
|
||||
## so that it is at "normal" pose when applied camera_rt_matrix
|
||||
## Format: wxyz for input & return
|
||||
def z_up_to_y_up(quat_wxyz):
|
||||
## Operations (1) rotate around y for pi/2,
|
||||
## (2) rotate around z for pi/2
|
||||
to_y_up = transforms3d.euler.euler2quat(0, np.pi/2, np.pi/2)
|
||||
return transforms3d.quaternions.qmult(quat_wxyz, to_y_up)
|
||||
|
||||
## Models coming out of opengl are negative x facing
|
||||
## Transform the default to
|
||||
def y_up_to_z_up(quat_wxyz):
|
||||
to_z_up = transforms3d.euler.euler2quat(np.pi/2, 0, np.pi/2)
|
||||
return transforms3d.quaternions.qmult(to_z_up, quat_wxyz)
|
||||
|
||||
|
||||
def quat_wxyz_to_euler(wxyz):
|
||||
q0, q1, q2, q3 = wxyz
|
||||
sinr = 2 * (q0 * q1 + q2 * q3)
|
||||
cosr = 1 - 2 * (q1 * q1 + q2 * q2)
|
||||
sinp = 2 * (q0 * q2 - q3 * q1)
|
||||
siny = 2 * (q0 * q3 + q1 * q2)
|
||||
cosy = 1 - 2 * (q2 * q2 + q3 * q3)
|
||||
|
||||
roll = np.arctan2(sinr, cosr)
|
||||
pitch = np.arcsin(sinp)
|
||||
yaw = np.arctan2(siny, cosy)
|
||||
return [roll, pitch, yaw]
|
||||
|
||||
|
||||
## wxyz: numpy array format
|
||||
def quat_wxyz_to_xyzw(wxyz):
|
||||
return np.concatenate((wxyz[1:], wxyz[:1]))
|
||||
|
||||
## xyzw: numpy array format
|
||||
def quat_xyzw_to_wxyz(xyzw):
|
||||
return np.concatenate((xyzw[-1:], xyzw[:-1]))
|
||||
|
||||
|
||||
## DEPRECATED
|
||||
## Talking to physics simulation
|
||||
## New state: [x, y, z, r_w, r_x, r_y, r_z]
|
||||
def hasNoCollision(posi_xyz, quat_wxyz):
|
||||
## Change quaternion basis
|
||||
#print("new_state 0", new_state)
|
||||
#print("new temp st", [new_state[4], new_state[5], new_state[6], new_state[3]], transforms3d.euler.euler2quat(-np.pi/2, 0, 0))
|
||||
|
||||
new_state = np.concatenate((posi_xyz, quat_wxyz))
|
||||
## URDF file uses z-up as front, so we need to rotate new_state around z for -90
|
||||
#print("new_state 1", new_state)
|
||||
socket_phys.send(" ".join(map(str, new_state)))
|
||||
collisions_count = int(socket_phys.recv().decode("utf-8"))
|
||||
return collisions_count == 0
|
||||
|
||||
def getPoseOrientationFromPhysics():
|
||||
receive = socket_phys.recv().decode("utf-8")
|
||||
new_pos, new_quat = json.loads(receive)
|
||||
socket_phys.send(json.dumps({"received": True}))
|
||||
return new_pos, new_quat
|
||||
|
||||
|
||||
## Return pos(xyz), quat(wxyz)
|
||||
def getNewPoseFromPhysics(view_pose):
|
||||
view_pose['quat'] = quat_wxyz_to_xyzw(view_pose['quat']).tolist()
|
||||
#print("sending over new pose", view_pose['pos'])
|
||||
#print("sending over new quat", view_pose['quat'])
|
||||
socket_phys.send(json.dumps(view_pose))
|
||||
new_pos, new_quat = json.loads(socket_phys.recv().decode("utf-8"))
|
||||
return new_pos, quat_xyzw_to_wxyz(new_quat)
|
||||
|
||||
|
||||
class PCRenderer:
|
||||
def __init__(self):
|
||||
self.roll, self.pitch, self.yaw = 0, 0, 0
|
||||
self.quat = [1, 0, 0, 0]
|
||||
self.x, self.y, self.z = 0, 0, 0
|
||||
self.fps = 0
|
||||
self.mousex, self.mousey = 0.5, 0.5
|
||||
self.changed = True
|
||||
self.org_pitch, self.org_yaw, self.org_roll = 0, 0, 0
|
||||
self.org_x, self.org_y, self.org_z = 0, 0, 0
|
||||
self.clickstart = (0,0)
|
||||
self.mousedown = False
|
||||
self.fps = 0
|
||||
self.rotation_const = np.array([[0,1,0,0],[0,0,1,0],[-1,0,0,0],[0,0,0,1]])
|
||||
|
||||
def onmouse(self, *args):
|
||||
if args[0] == cv2.EVENT_LBUTTONDOWN:
|
||||
self.org_pitch, self.org_yaw, self.org_x, self.org_y, self.org_z =\
|
||||
self.pitch,self.yaw,self.x,self.y,self.z
|
||||
self.clickstart = (self.mousex, self.mousey)
|
||||
|
||||
if args[0] == cv2.EVENT_RBUTTONDOWN:
|
||||
self.org_roll = self.roll
|
||||
self.clickstart = (self.mousex, self.mousey)
|
||||
|
||||
if (args[3] & cv2.EVENT_FLAG_LBUTTON):
|
||||
self.pitch = self.org_pitch + (self.mousex - self.clickstart[0])/10
|
||||
self.yaw = self.org_yaw + (self.mousey - self.clickstart[1])
|
||||
self.changed=True
|
||||
|
||||
if (args[3] & cv2.EVENT_FLAG_RBUTTON):
|
||||
self.roll = self.org_roll + (self.mousex - self.clickstart[0])/50
|
||||
self.changed=True
|
||||
|
||||
my=args[1]
|
||||
mx=args[2]
|
||||
self.mousex=mx/float(256)
|
||||
self.mousey=my/float(256 * 2)
|
||||
|
||||
def getViewerCpose(self):
|
||||
cpose = np.eye(4)
|
||||
alpha = self.yaw
|
||||
beta = self.pitch
|
||||
gamma = self.roll
|
||||
cpose = cpose.flatten()
|
||||
|
||||
cpose[0] = cos(alpha) * cos(beta);
|
||||
cpose[1] = cos(alpha) * sin(beta) * sin(gamma) - sin(alpha) * cos(gamma);
|
||||
cpose[2] = cos(alpha) * sin(beta) * cos(gamma) + sin(alpha) * sin(gamma);
|
||||
cpose[3] = 0
|
||||
|
||||
cpose[4] = sin(alpha) * cos(beta);
|
||||
cpose[5] = sin(alpha) * sin(beta) * sin(gamma) + cos(alpha) * cos(gamma);
|
||||
cpose[6] = sin(alpha) * sin(beta) * cos(gamma) - cos(alpha) * sin(gamma);
|
||||
cpose[7] = 0
|
||||
|
||||
cpose[8] = -sin(beta);
|
||||
cpose[9] = cos(beta) * sin(gamma);
|
||||
cpose[10] = cos(beta) * cos(gamma);
|
||||
cpose[11] = 0
|
||||
|
||||
cpose[12:16] = 0
|
||||
cpose[15] = 1
|
||||
|
||||
cpose = cpose.reshape((4,4))
|
||||
|
||||
cpose2 = np.eye(4)
|
||||
cpose2[0,3] = self.x
|
||||
cpose2[1,3] = self.y
|
||||
cpose2[2,3] = self.z
|
||||
|
||||
cpose = np.dot(cpose, cpose2)
|
||||
return cpose
|
||||
#print('new_cpose', cpose)
|
||||
|
||||
def getViewerCpose3D(self):
|
||||
""" Similar to getViewerCpose, implemented using transforms3d
|
||||
"""
|
||||
cpose = np.eye(4)
|
||||
gamma = self.yaw
|
||||
alpha = self.pitch
|
||||
beta = -self.roll
|
||||
#cpose[:3, :3] = transforms3d.euler.euler2mat(alpha, beta, gamma)
|
||||
cpose[:3, :3] = transforms3d.quaternions.quat2mat(self.quat)
|
||||
cpose[ 0, -1] = self.x
|
||||
cpose[ 1, -1] = self.y
|
||||
cpose[ 2, -1] = self.z
|
||||
|
||||
return cpose
|
||||
|
||||
def render(self, imgs, depths, pose, model, poses, target_pose, show, target_depth, opengl_arr):
|
||||
t0 = time.time()
|
||||
|
||||
v_cam2world = target_pose#.dot(poses[0])
|
||||
p = (v_cam2world).dot(np.linalg.inv(pose))
|
||||
|
||||
p = p.dot(np.linalg.inv(self.rotation_const))
|
||||
s = mat_to_str(p)
|
||||
|
||||
'''
|
||||
p = pose.dot(np.linalg.inv(poses[0])) #.dot(target_pose)
|
||||
|
||||
trans = -pose[:3, -1]
|
||||
rot = np.linalg.inv(pose[:3, :3])
|
||||
|
||||
|
||||
p2 = np.eye(4)
|
||||
p2[:3, :3] = rot
|
||||
p2[:3, -1] = trans
|
||||
|
||||
|
||||
s = mat_to_str(poses[0] * p2)
|
||||
'''
|
||||
|
||||
|
||||
with Profiler("Depth request round-trip"):
|
||||
socket_mist.send(s)
|
||||
message = socket_mist.recv()
|
||||
|
||||
with Profiler("Read from framebuffer and make pano"):
|
||||
# data = np.array(np.frombuffer(message, dtype=np.float32)).reshape((6, 768, 768, 1))
|
||||
|
||||
wo, ho = 768 * 4, 768 * 3
|
||||
|
||||
# Calculate height and width of output image, and size of each square face
|
||||
h = wo/3
|
||||
w = 2*h
|
||||
n = ho/3
|
||||
opengl_arr = np.array(np.frombuffer(message, dtype=np.float32)).reshape((h, w))
|
||||
'''
|
||||
PHYSICS
|
||||
data = np.array(np.frombuffer(message, dtype=np.float32)).reshape((6, 768, 768, 1))
|
||||
## For some reason, the img passed back from opengl is upside down.
|
||||
## This is still yet to be debugged
|
||||
data = data[:, ::-1,::,:]
|
||||
img_array = []
|
||||
for i in range(6):
|
||||
img_array.append(data[i])
|
||||
|
||||
img_array2 = [img_array[0], img_array[1], img_array[2], img_array[3], img_array[4], img_array[5]]
|
||||
|
||||
opengl_arr = convert_array(np.array(img_array2))
|
||||
opengl_arr = opengl_arr[::, ::]
|
||||
|
||||
opengl_arr_err = opengl_arr == 0
|
||||
|
||||
|
||||
opengl_arr_show = (opengl_arr * 3500.0 / 128).astype(np.uint8)
|
||||
opengl_arr_show[opengl_arr_err[:, :, 0], 1:3] = 0
|
||||
opengl_arr_show[opengl_arr_err[:, :, 0], 0] = 255
|
||||
cv2.imshow('target depth',opengl_arr_show)
|
||||
'''
|
||||
|
||||
def render_depth(opengl_arr):
|
||||
with Profiler("Render Depth"):
|
||||
cv2.imshow('target depth', opengl_arr/16.)
|
||||
|
||||
def render_pc(opengl_arr):
|
||||
'''
|
||||
PHYSICS
|
||||
#from IPython import embed; embed()
|
||||
#target_depth[:, :, 0] = (opengl_arr[:,:,0] * 100).astype(np.int32)
|
||||
target_depth[:, :] = (opengl_arr[:,:,0] * 100).astype(np.int32)
|
||||
|
||||
print("images type", depths[0].dtype, "opengl type", opengl_arr.dtype)
|
||||
print("images: min", np.min(depths) * 12800, "mean", np.mean(depths) * 12800, "max", np.max(depths) * 12800)
|
||||
print('target: min', np.min(target_depth), "mean", np.mean(target_depth), "max", np.max(target_depth))
|
||||
|
||||
show[:] = 0
|
||||
before = time.time()
|
||||
for i in range(len(imgs)):
|
||||
pose_after = pose.dot(np.linalg.inv(poses[0])).dot(poses[i]).astype(np.float32)
|
||||
|
||||
dll.render(ct.c_int(imgs[i].shape[0]),
|
||||
ct.c_int(imgs[i].shape[1]),
|
||||
imgs[i].ctypes.data_as(ct.c_void_p),
|
||||
depths[i].ctypes.data_as(ct.c_void_p),
|
||||
pose_after.ctypes.data_as(ct.c_void_p),
|
||||
show.ctypes.data_as(ct.c_void_p),
|
||||
target_depth.ctypes.data_as(ct.c_void_p)
|
||||
)
|
||||
'''
|
||||
with Profiler("Render pointcloud"):
|
||||
scale = 100. # 512
|
||||
target_depth = np.int32(opengl_arr * scale)
|
||||
show[:] = 0
|
||||
poses_after = [
|
||||
pose.dot(np.linalg.inv(poses[i])).astype(np.float32)
|
||||
for i in range(len(imgs))]
|
||||
|
||||
for i in range(len(imgs)):
|
||||
dll.render(ct.c_int(imgs[i].shape[0]),
|
||||
ct.c_int(imgs[i].shape[1]),
|
||||
imgs[i].ctypes.data_as(ct.c_void_p),
|
||||
depths[i].ctypes.data_as(ct.c_void_p),
|
||||
poses_after[i].ctypes.data_as(ct.c_void_p),
|
||||
show.ctypes.data_as(ct.c_void_p),
|
||||
target_depth.ctypes.data_as(ct.c_void_p)
|
||||
)
|
||||
threads = [
|
||||
Process(target=render_pc, args=(opengl_arr,)),
|
||||
Process(target=render_depth, args=(opengl_arr,))]
|
||||
[t.start() for t in threads]
|
||||
[t.join() for t in threads]
|
||||
|
||||
|
||||
if model:
|
||||
tf = transforms.ToTensor()
|
||||
before = time.time()
|
||||
source = tf(show)
|
||||
source_depth = tf(np.expand_dims(target_depth, 2).astype(np.float32)/65536 * 255)
|
||||
#print(source.size(), source_depth.size())
|
||||
imgv.data.copy_(source)
|
||||
maskv.data.copy_(source_depth)
|
||||
print('Transfer time', time.time() - before)
|
||||
before = time.time()
|
||||
recon = model(imgv, maskv)
|
||||
print('NNtime:', time.time() - before)
|
||||
before = time.time()
|
||||
show2 = recon.data.cpu().numpy()[0].transpose(1,2,0)
|
||||
show[:] = (show2[:] * 255).astype(np.uint8)
|
||||
print('Transfer to CPU time:', time.time() - before)
|
||||
|
||||
t1 =time.time()
|
||||
t = t1-t0
|
||||
self.fps = 1/t
|
||||
|
||||
#cv2.waitKey(5)%256
|
||||
|
||||
|
||||
def instantCheckPos(self, pos, target_pose):
|
||||
cpose = self.getViewerCpose()
|
||||
v_cam2world = target_pose.dot(poses[0])
|
||||
world_cpose = (v_cam2world).dot(np.linalg.inv(cpose)).dot(np.linalg.inv(self.rotation_const))
|
||||
new_pos = mat_to_posi_xyz(world_cpose).tolist()
|
||||
#print("Instant chec xyz", x, y, z)
|
||||
print("Instant check pose", new_pos)
|
||||
|
||||
|
||||
def showpoints(self, imgs, depths, poses, model, target, tdepth, target_poses):
|
||||
showsz = target.shape[0]
|
||||
#print('target pose', target_pose)
|
||||
|
||||
#v_cam2world = target_pose.dot(poses[0])
|
||||
v_cam2world = target_poses[0]
|
||||
cpose = self.getViewerCpose3D()
|
||||
p = (v_cam2world).dot(np.linalg.inv(cpose))
|
||||
p = p.dot(np.linalg.inv(self.rotation_const))
|
||||
pos = mat_to_posi_xyz(p).tolist()
|
||||
quat = mat_to_quat_xyzw(p).tolist()
|
||||
receive = str(socket_phys.recv().decode("utf-8"))
|
||||
if (receive == "Initial"):
|
||||
socket_phys.send(json.dumps([pos, quat]))
|
||||
else:
|
||||
return
|
||||
|
||||
initialPos = pos
|
||||
initialQuat = quat
|
||||
initialCpose = cpose
|
||||
|
||||
show=np.zeros((showsz,showsz * 2,3),dtype='uint8')
|
||||
target_depth = np.zeros((showsz,showsz * 2)).astype(np.int32)
|
||||
|
||||
overlay = False
|
||||
show_depth = False
|
||||
cv2.namedWindow('show3d')
|
||||
cv2.namedWindow('target depth')
|
||||
|
||||
cv2.moveWindow('show3d',0,0)
|
||||
cv2.setMouseCallback('show3d',self.onmouse)
|
||||
|
||||
imgv = Variable(torch.zeros(1,3, showsz, showsz*2), volatile=True).cuda()
|
||||
maskv = Variable(torch.zeros(1,1, showsz, showsz*2), volatile=True).cuda()
|
||||
|
||||
old_state = [self.x, self.y, self.z, self.roll, self.pitch, self.yaw]
|
||||
old_cpose = np.eye(4)
|
||||
|
||||
while True:
|
||||
cpose_viewer = self.getViewerCpose3D()
|
||||
#v_cam2world = target_pose.dot(poses[0])
|
||||
v_cam2world = target_poses[0]
|
||||
world_cpose = (v_cam2world).dot(np.linalg.inv(cpose_viewer)).dot(np.linalg.inv(self.rotation_const))
|
||||
cpose = np.linalg.inv(np.linalg.inv(v_cam2world).dot(cpose_viewer).dot(self.rotation_const))
|
||||
#print("initial cpose", initialCpose)
|
||||
print("current cpose", cpose)
|
||||
#print("initial pos", initialPos)
|
||||
#print("current pos", self.x, self.y, self.z)
|
||||
#print("world pose", world_cpose)
|
||||
|
||||
|
||||
## Query physics engine to get [x, y, z, roll, pitch, yaw]
|
||||
if PHYSICS_FIRST:
|
||||
pos = mat_to_posi_xyz(world_cpose).tolist()
|
||||
quat_wxyz = quat_xyzw_to_wxyz(mat_to_quat_xyzw(world_cpose)).tolist()
|
||||
|
||||
'''
|
||||
new_pos, new_quat_wxyz = getNewPoseFromPhysics({
|
||||
'changed': self.changed,
|
||||
'pos': pos,
|
||||
'quat': quat_wxyz
|
||||
})
|
||||
'''
|
||||
|
||||
#new_quat = y_up_to_z_up(new_quat_wxyz)
|
||||
#new_quat = new_quat_wxyz
|
||||
print("waiting for physics")
|
||||
#new_pos, new_euler = getPoseOrientationFromPhysics()
|
||||
new_pos, new_quat = getPoseOrientationFromPhysics()
|
||||
|
||||
self.x, self.y, self.z = new_pos
|
||||
|
||||
self.quat = new_quat
|
||||
self.changed = True
|
||||
|
||||
if self.changed:
|
||||
new_quat = mat_to_quat_xyzw(world_cpose)
|
||||
new_quat = z_up_to_y_up(quat_xyzw_to_wxyz(new_quat))
|
||||
|
||||
new_posi = mat_to_posi_xyz(world_cpose)
|
||||
|
||||
## Entry point for change of view
|
||||
## If PHYSICS_FIRST mode, then collision is already handled
|
||||
## inside physics simulator
|
||||
if PHYSICS_FIRST or hasNoCollision(new_posi, new_quat):
|
||||
## Optimization
|
||||
depth_buffer = np.zeros(imgs[0].shape[:2], dtype=np.float32)
|
||||
#render(imgs, depths, cpose.astype(np.float32), model, poses, depth_buffer)
|
||||
|
||||
|
||||
relative_poses = np.copy(target_poses)
|
||||
for i in range(len(relative_poses)):
|
||||
relative_poses[i] = np.dot(np.linalg.inv(relative_poses[i]), target_poses[0])
|
||||
|
||||
poses_after = [
|
||||
cpose.dot(np.linalg.inv(relative_poses[i])).astype(np.float32)
|
||||
for i in range(len(imgs))]
|
||||
|
||||
pose_after_distance = [np.linalg.norm(rt[:3,-1]) for rt in poses_after]
|
||||
top5 = (np.argsort(pose_after_distance))[:5]
|
||||
imgs_top5 = [imgs[i] for i in top5]
|
||||
depths_top5 = [depths[i] for i in top5]
|
||||
relative_poses_top5 = [relative_poses[i] for i in top5]
|
||||
|
||||
|
||||
self.render(imgs_top5, depths_top5, cpose.astype(np.float32), model, relative_poses_top5, target_poses[0], show, target_depth, depth_buffer)
|
||||
old_state = [self.x, self.y, self.z, self.roll, self.pitch, self.yaw]
|
||||
old_cpose = np.copy(cpose)
|
||||
else:
|
||||
self.x, self.y, self.z, self.roll, self.pitch, self.yaw = old_state
|
||||
cpose = old_cpose
|
||||
|
||||
#render(imgs, depths, cpose.astype(np.float32), model, poses)
|
||||
#old_state = [x, y, z, roll, pitch, yaw]
|
||||
self.changed = False
|
||||
|
||||
if overlay:
|
||||
show_out = (show/2 + target/2).astype(np.uint8)
|
||||
elif show_depth:
|
||||
show_out = (target_depth * 10).astype(np.uint8)
|
||||
else:
|
||||
show_out = show
|
||||
|
||||
#assert(np.sum(show) != 0)
|
||||
|
||||
|
||||
cv2.putText(show,'pitch %.3f yaw %.2f roll %.3f x %.2f y %.2f z %.2f'%(self.pitch, self.yaw, self.roll, self.x, self.y, self.z),(15,showsz-15),0,0.5,(255,255,255))
|
||||
#print("roll %f pitch %f yaw %f" % (roll, pitch, yaw))
|
||||
|
||||
#cv2.putText(show,'fps %.1f'%(fps),(15,15),0,0.5,cv2.cv.CV_RGB(255,255,255))
|
||||
cv2.putText(show,'fps %.1f'%(self.fps),(15,15),0,0.5,(255,255,255))
|
||||
|
||||
show_rgb = cv2.cvtColor(show_out, cv2.COLOR_BGR2RGB)
|
||||
cv2.imshow('show3d',show_rgb)
|
||||
|
||||
cmd=cv2.waitKey(5)%256
|
||||
|
||||
## delta = [x, y, z, roll, pitch, yaw]
|
||||
#delta = [0, 0, 0, 0, 0, 0]
|
||||
|
||||
if cmd==ord('q'):
|
||||
break
|
||||
elif cmd == ord('w'):
|
||||
self.x -= 0.05
|
||||
#delta = [-0.05, 0, 0, 0, 0, 0]
|
||||
self.changed = True
|
||||
elif cmd == ord('s'):
|
||||
self.x += 0.05
|
||||
#delta = [0.05, 0, 0, 0, 0, 0]
|
||||
self.changed = True
|
||||
elif cmd == ord('a'):
|
||||
self.y += 0.05
|
||||
#delta = [0, 0.05, 0, 0, 0, 0]
|
||||
self.changed = True
|
||||
elif cmd == ord('d'):
|
||||
self.y -= 0.05
|
||||
#delta = [0, -0.05, 0, 0, 0, 0]
|
||||
self.changed = True
|
||||
elif cmd == ord('z'):
|
||||
self.z += 0.01
|
||||
#delta = [0, 0, 0, 0, 0, 0.01]
|
||||
self.changed = True
|
||||
elif cmd == ord('x'):
|
||||
self.z -= 0.01
|
||||
#delta = [0, 0, 0, 0, 0, -0.01]
|
||||
self.changed = True
|
||||
|
||||
elif cmd == ord('r'):
|
||||
self.pitch,self.yaw,self.x,self.y,self.z = 0,0,0,0,0
|
||||
self.roll = 0
|
||||
self.changed = True
|
||||
elif cmd == ord('t'):
|
||||
pose = poses[0]
|
||||
RT = pose.reshape((4,4))
|
||||
R = RT[:3,:3]
|
||||
T = RT[:3,-1]
|
||||
|
||||
self.x,self.y,self.z = np.dot(np.linalg.inv(R),T)
|
||||
self.roll, self.pitch, self.yaw = (utils.rotationMatrixToEulerAngles(R))
|
||||
|
||||
self.changed = True
|
||||
|
||||
|
||||
elif cmd == ord('o'):
|
||||
overlay = not overlay
|
||||
elif cmd == ord('f'):
|
||||
show_depth = not show_depth
|
||||
elif cmd == ord('v'):
|
||||
cv2.imwrite('save.jpg', show_rgb)
|
||||
|
||||
'''
|
||||
state = [x, y, z, roll, pitch, yaw]
|
||||
if (hasNoCollision(state, delta)):
|
||||
x = x + delta[0]
|
||||
y = y + delta[1]
|
||||
z = z + delta[2]
|
||||
roll = roll + delta[3]
|
||||
pitch = pitch + delta[4]
|
||||
yaw = yaw + delta[5]
|
||||
'''
|
||||
#print("roll %d pitch %d yaw %d" % (roll, pitch, yaw))
|
||||
|
||||
def show_target(target_img):
|
||||
cv2.namedWindow('target')
|
||||
cv2.moveWindow('target',0,256 + 50)
|
||||
show_rgb = cv2.cvtColor(target_img, cv2.COLOR_BGR2RGB)
|
||||
cv2.imshow('target', show_rgb)
|
||||
|
||||
if __name__=='__main__':
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--debug' , action='store_true', help='debug mode')
|
||||
parser.add_argument('--datapath' , required = True, help='dataset path')
|
||||
parser.add_argument('--model_id' , type = str, default = 0, help='model id')
|
||||
parser.add_argument('--model' , type = str, default = '', help='path of model')
|
||||
|
||||
opt = parser.parse_args()
|
||||
d = ViewDataSet3D(root=opt.datapath, transform = np.array, mist_transform = np.array, seqlen = 2, off_3d = False, train = False)
|
||||
|
||||
scene_dict = dict(zip(d.scenes, range(len(d.scenes))))
|
||||
if not opt.model_id in scene_dict.keys():
|
||||
print("model not found")
|
||||
else:
|
||||
scene_id = scene_dict[opt.model_id]
|
||||
|
||||
uuids, rts = d.get_scene_info(scene_id)
|
||||
print(uuids, rts)
|
||||
|
||||
targets = []
|
||||
sources = []
|
||||
source_depths = []
|
||||
poses = []
|
||||
|
||||
for k,v in uuids:
|
||||
print(k,v)
|
||||
data = d[v]
|
||||
source = data[0][0]
|
||||
target = data[1]
|
||||
target_depth = data[3]
|
||||
source_depth = data[2][0]
|
||||
pose = data[-1][0].numpy()
|
||||
targets.append(target)
|
||||
poses.append(pose)
|
||||
sources.append(target)
|
||||
source_depths.append(target_depth)
|
||||
|
||||
model = None
|
||||
if opt.model != '':
|
||||
comp = CompletionNet()
|
||||
comp = torch.nn.DataParallel(comp).cuda()
|
||||
comp.load_state_dict(torch.load(opt.model))
|
||||
model = comp.module
|
||||
model.eval()
|
||||
print(model)
|
||||
print('target', poses, poses[0])
|
||||
#print('no.1 pose', poses, poses[1])
|
||||
# print(source_depth)
|
||||
print(sources[0].shape, source_depths[0].shape)
|
||||
|
||||
context_mist = zmq.Context()
|
||||
print("Connecting to hello world server...")
|
||||
|
||||
socket_mist = context_mist.socket(zmq.REQ)
|
||||
socket_mist.connect("tcp://localhost:5555")
|
||||
|
||||
print(coords.flatten().dtype)
|
||||
with Profiler("Transform coords"):
|
||||
new_coords = np.getbuffer(coords.flatten().astype(np.uint32))
|
||||
print(coords.shape)
|
||||
print("Count: ", coords.flatten().astype(np.uint32).size )
|
||||
print("elem [2,3,5]: ", coords[4][2][1] )
|
||||
socket_mist.send(new_coords)
|
||||
print("Sent reordering")
|
||||
message = socket_mist.recv()
|
||||
print("msg")
|
||||
|
||||
context_phys = zmq.Context()
|
||||
socket_phys = context_phys.socket(zmq.REP)
|
||||
socket_phys.connect("tcp://localhost:5556")
|
||||
|
||||
show_target(target)
|
||||
|
||||
renderer = PCRenderer()
|
||||
renderer.showpoints(sources, source_depths, poses, model, target, target_depth, rts)
|
Loading…
Reference in New Issue