NYU hand pose dataset visualization example

NYU hand pose dataset visualization example

# This is a simple script to visualize a training example
import matplotlib.pyplot as plt
import numpy as np
import os

from os.path import join
from scipy.io import loadmat
from mpl_toolkits.mplot3d import Axes3D

np.random.seed(1024)

class NYU_HPD(object):
	"""docstring for NYU hand pose dataset"""
	def __init__(self,
				 dataset_dir='./train/',
				 image_index=1,
				 kinect_index=1):
		super(NYU_HPD, self).__init__()

		filename_prefix = '%d_%07d' % (kinect_index, image_index)

		self.kinect_index = kinect_index
		self.image_index = image_index
		self.rgb_path = join(dataset_dir,
									 'rgb_%s.png' % filename_prefix)
		self.depth_path = join(dataset_dir,
									   'depth_%s.png' % filename_prefix)
		self.synthdepth_path = join(dataset_dir,
											'synthdepth_%s.png' % filename_prefix)
		self.joint_path = join(dataset_dir, 'joint_data.mat')

	# loading very large images
	def load_image(self, img_path):
		assert os.path.isfile(img_path), 'missing image: %s' % img_path
		img = plt.imread(img_path)

		if np.max(img) <= 1.5:
			np.clip(img, a_min=0, a_max=1, out=img)
			# this command split should reduce mount of required memory
			np.multiply(img, 255, out=img)
			img = img.astype(np.uint8, copy=False)

		return img

	# Load and display an RGB example
	def show_rgb(self):
		rgb = self.load_image(self.rgb_path)
		fig = plt.figure()
		plt.imshow(rgb)
		plt.axis('off')
		plt.show()

	# Load and display a synthdepth/depth example
	def show_depth(self, synth=False):
		# The top 8 bits of depth are packed into green 
		# and the lower 8 bits into blue
		if synth:
			depth = self.load_image(self.synthdepth_path)
		else:
			depth = self.load_image(self.depth_path)
		depth = depth[:, :, 2].astype('uint16') + \
					 (depth[:, :, 1].astype('uint16') << 8)
		ind = np.nonzero(depth)
		plt.imshow(depth,
				   cmap='gray',
				   vmax=np.max(depth[ind]) + 10,
				   vmin=np.min(depth[ind]) - 10)
		plt.axis('off')
		plt.show()

	# Load the UVD Coefficients and display them on the depth image
	def show_joints(self):
		dat = loadmat(self.joint_path)
		jnt_uvd = dat['joint_uvd'][self.kinect_index - 1, self.image_index - 1, :, :]

		jnt_colors = np.random.rand(jnt_uvd.shape[0], 3)

		fig = plt.figure()
		plt.scatter(jnt_uvd[:, 0], jnt_uvd[:, 1],
					c=jnt_colors,
					alpha=0.75)

		self.show_depth()

	# Visualize the hand and the joints in 3D
	def show_3D_hand(self):
		for idx in range(2):
			if idx:
				depth = self.load_image(self.synthdepth_path)
			else:
				depth = self.load_image(self.depth_path)
			depth = depth[:, :, 2].astype('uint16') + \
					(depth[:, :, 1].astype('uint16') << 8)

			uvd = self.depth2uvd(depth)
			xyz = self.uvd2xyz(uvd)

			points = xyz.reshape([xyz.shape[0] * \
								  xyz.shape[1], 3])
			dat = loadmat(self.joint_path)
			jnt_uvd = dat['joint_uvd'][self.kinect_index - 1, self.image_index - 1, :, :]

			jnt_colors = np.random.rand(jnt_uvd.shape[0], 3)

			hand_points = self.uvd2xyz(np.expand_dims(jnt_uvd,
													  axis=0))
			hand_points = np.squeeze(hand_points)
			# Collect the points within the AABBOX of the hand
			axis_bounds = [min(hand_points[:, 0]), max(hand_points[:, 0]),
						   min(hand_points[:, 1]), max(hand_points[:, 1]),
						   min(hand_points[:, 2]), max(hand_points[:, 2])]
			axis_bounds = np.array(axis_bounds)
			axis_bounds[0:6:2] = axis_bounds[0:6:2] - 20
			axis_bounds[1:6:2] = axis_bounds[1:6:2] + 20

			ipnts = [points[i, 0] >= axis_bounds[0] and
					 points[i, 0] <= axis_bounds[1] and
					 points[i, 1] >= axis_bounds[2] and
					 points[i, 1] <= axis_bounds[3] and
					 points[i, 2] >= axis_bounds[4] and
					 points[i, 2] <= axis_bounds[5]
					 for i in range(points.shape[0])]

			points = points[ipnts, :]

			ax = plt.axes(projection='3d')
			# Visualize the hand point cloud
			ax.plot3D(points[:, 0],
					  points[:, 2],
					  points[:, 1],
					  '.')
			# ax.view_init(elev=0, azim=0)
			ax.set_xlabel('x')
			ax.set_ylabel('y')
			ax.set_zlabel('z')

			ax.scatter(hand_points[:, 0],
					   hand_points[:, 2],
					   hand_points[:, 1],
					   c=jnt_colors)

			plt.show()

	# Visualize the depth in 3D
	def show_3D_body(self):
		depth = self.load_image(self.depth_path)
		depth = depth[:, :, 2].astype('uint16') + \
				(depth[:, :, 1].astype('uint16') << 8)

		uvd = self.depth2uvd(depth)
		xyz = self.uvd2xyz(uvd)
		# uvd_ = self.xyz2uvd(xyz)

		# Decimate the image (otherwise rendering will be too slow)
		decimation = 4
		xyz_decimated = xyz[0:xyz.shape[0]:decimation,
						0:xyz.shape[1]:decimation, :]
		points = xyz_decimated.reshape([xyz_decimated.shape[0] * \
										xyz_decimated.shape[1], 3])

		# Collect the points within the AABBOX of the non-background points
		body_points = points[points[:, 2] < 2000, :]
		axis_bounds = [min(body_points[:, 0]), max(body_points[:, 0]),
					   min(body_points[:, 1]), max(body_points[:, 1]),
					   min(body_points[:, 2]), max(body_points[:, 2])]

		fig = plt.figure()
		ax = plt.axes(projection='3d')
		# Visualize the entire point cloud
		ax.plot3D(body_points[:, 0],
				  body_points[:, 2],
				  body_points[:, 1],
				  'b.')
		# ax.view_init(elev=0, azim=0)
		ax.set_xlabel('x')
		ax.set_ylabel('y')
		ax.set_zlabel('z')

		plt.xlim(axis_bounds[0:2])
		plt.ylim(axis_bounds[4:6])
		# plt.zlim(axis_bounds[4:6])

		plt.show()

	def depth2uvd(self, depth):
		v, u = np.mgrid[0:depth.shape[0], 0:depth.shape[1]]
		uvd = np.dstack((u, v, depth))

		return uvd

	def uvd2xyz(self, uvd):
		xRes = 640
		yRes = 480
		xzFactor = 1.08836710
		yzFactor = 0.817612648

		normalizedX = uvd[:,:, 0].astype('float') / xRes - 0.5
		normalizedY = 0.5 - uvd[:,:, 1].astype('float') / yRes

		xyz = np.zeros_like(uvd, dtype='float')
		xyz[:, :, 2] = uvd[:,:, 2].astype('float')
		xyz[:, :, 1] = normalizedY * xyz[:, :, 2] * yzFactor
		xyz[:, :, 0] = normalizedX * xyz[:, :, 2] * xzFactor

		return xyz

	def xyz2uvd(self, xyz):
		halfResX = 640 / 2
		halfResY = 480 / 2
		coeffX = 588.036865
		coeffY = 587.075073

		uvd = np.zeros_like(xyz)
		uvd[:, :, 0] = coeffX * xyz[:, :, 0] / xyz[:,:,2] + halfResX
		uvd[:, :, 1] = halfResY - coeffY * xyz[:,:, 1] / xyz[:, :, 2]
		uvd[:, :, 2] = xyz[:, :, 2]

		return uvd


def main():
	dataset_dir = './train/'
	image_index = 1
	kinect_index = 1

	nyu = NYU_HPD(dataset_dir, image_index, kinect_index)
	# nyu.show_rgb()
	nyu.show_3D_hand()
	# nyu.show_joints()


if __name__ == '__main__':
	main()