model_visualizer.py

# Don't run as script/executable, just run as python file
import rospy
from rospy.numpy_msg import numpy_msg
from std_msgs.msg import Int32, Float32, String, Float64MultiArray, Float64
import sys, os
# sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# from hierarchy_roach.msg import camera_message, obs_message
from nn_dynamics_roach.msg import camera_message, obs_message

import numpy as np
# import cv2
from rospy.exceptions import ROSException
from sensor_msgs.msg import Image
# from cv_bridge import CvBridge, CvBridgeError
import pickle
import IPython
# IPython.embed()
import matplotlib.pyplot as plt
from PIL import Image
from gcg.misc import utils
from gcg.misc.utils import import_params
import collections
from panda3d.core import LPoint2, LPoint3
import matplotlib.colors as colors
import matplotlib.cm as cmx
# import argparse
"""
Troubleshoot:
Train model with prediction Horizon = 1
Train without yaw input (i.e. does the history of yaw's give it info about dynamics; you can't compute that here)
Train with history len = 1

Extras to implement:
Visualize H > 1 

https://github.com/alexlee-gk/citysim3d

TODO: thing that greg said about the danger of directly using get_output fn's on model
"""
#latest_data = None
data_folder_path = "/media/anagabandi/f1e71f04-dc4b-4434-ae4c-fcb16447d5b3/gcg_roach/data"
class ModelVisualizer:
	def __init__(self, exp_name, eval_itr, mode, viz_mode):
		self._eval_itr = eval_itr
		self._mode = mode
		self._viz_mode = viz_mode

		# load config
		py_config_path = os.path.join(data_folder_path, exp_name, "params.py")
		assert(os.path.exists(py_config_path))
		params = import_params(py_config_path)
		# with open(py_config_path, 'r') as f:
		#     params_txt = ''.join(f.readlines())

		### create environments
		env_eval_params=params['env_eval'] if params['env_eval'] else params['env']
		self._env_eval = env_eval_params['class'](params=env_eval_params['kwargs'])
		self._env_eval_params = env_eval_params

		self._target_image_height = self._env_eval.spec.observation_im_space.shape[0]
		self._target_image_width = self._env_eval.spec.observation_im_space.shape[1]
		self._target_aspect = self._target_image_width/self._target_image_height

		self._policy_params = params['policy']
		### create policy
		self._policy = self._policy_params['class'](
			env_spec=self._env_eval.spec,
			exploration_strategies=[],
			inference_only=True,
			**self._policy_params['kwargs']
		)

		### Load policy
		policy_fname = os.path.join(data_folder_path, params['exp_name'], 'itr_{}_inference_policy.ckpt'.format(self._eval_itr))
		self._policy.restore(policy_fname, train=False)

		# Make hist_len queues for image
		self._img_queue = collections.deque([])
		self._history_len = self._policy._history_len
		self._action_hor = self._policy._H

		self._obs_queue = collections.deque([])
		self._actions_queue = collections.deque([])

		# Doing computation outside of callback loop (reduce lag)
		self._yaw_max = self._env_eval._yaw_limits[1]
		self._yaw_min = self._env_eval._yaw_limits[0]
		self._num_samples = 1
		self._yaw_samples = np.linspace(self._yaw_min, self._yaw_max, self._num_samples)

		speed = self._env_eval._speed_range[0]
		self._height = self._env_eval._height_range[0]
		x = speed*-np.sin(np.radians(self._yaw_samples))
		y = speed*np.cos(np.radians(self._yaw_samples))

		self._lens = self._env_eval._panda._camera.node().getLens()
		points3d = np.ones((4, self._num_samples))
		points3d[0,:] = x
		points3d[1,:] = y
		points3d[2,:] = np.ones(self._num_samples)*self._height
		self._x, self._y = self.points3d_to_xy(points3d)
		# IPython.embed()

		# plt.scatter(self._x, self._y)
		# plt.ylim(0, self._target_image_height)
		# plt.xlim(0, self._target_image_width)
		# IPython.embed()

		self._obs_vecs = np.tile(np.array([0, 0, 0, 0]), (self._num_samples, self._history_len, 1)) # ['heading', 'coll']
		self._prev_actions = np.tile(np.array([0]), (self._num_samples, self._history_len-1, 1))
		self._goal = np.zeros((self._num_samples, 1))
		self._future_actions = np.zeros((self._num_samples, self._action_hor, 1))
		self._future_actions[:,0, 0] = self._yaw_samples

		##### Setup plotting
		# For some reason: subplots(2) throws a threading error: "main thread is not in main loop" from plt.pause and tkinter
		# self._fig, self._axes = plt.subplots(2)
		bwr = plt.get_cmap('bwr')
		cmn = colors.Normalize(vmin=0, vmax=1)
		self._cm = cmx.ScalarMappable(norm=cmn, cmap=bwr) 
		
	def points3d_to_xy(self, points3d):
		# Takes 3D point in camera coordinates to (x, y) pixel coordinates on image
		b_alt = self.points3d_to_points2d(points3d)
		xy = self.points2d_to_xy(self._lens, b_alt.T)
		return xy[:,0], xy[:,1]

	def points3d_to_points2d(self, points3d):
		point2d = LPoint2()
		b_alt = np.zeros((2, points3d.shape[1]))
		for i in range(points3d.shape[1]):
			point3d = LPoint3(points3d[0][i], points3d[1][i], points3d[2][i])
			if self._lens.project(point3d, point2d):
				b_alt[:,i]=np.array(point2d)
				# print("could project")
			else:
				# can't project because lies off page
				# print("couldn't project")
				b_alt[:,i]=np.array([np.nan, np.nan])
				# raise Exception("Couldn't project 3D point to 2D image space")
		return b_alt

	def points2d_to_xy(self, lens, points2d):
		points2d = (points2d + 1)/2.0
		points2d[: ,1] = 1 - points2d[:,1]
		target_film_size = np.array([self._target_image_width, self._target_image_height])
		points2d = points2d*target_film_size 

		return points2d

	def resize(self, image, dim):
		# image is numpy array	
		# crops and resizes to aspect ratio	
		image_width = dim[1]
		image_height = dim[0]
		aspect = image_width/image_height
		if aspect > self._target_aspect:
			# Then crop the top and b edges:
			new_width = int(self._target_aspect * image_height)
			offset = (image_width - new_width) / 2
			resize = (offset, 0, image_width - offset, height)
		else:
			# ... crop the top and bottom:
			new_height = int(image_width / self._target_aspect)
			offset = (image_height - new_height) / 2
			resize = (0, offset, image_width, image_height - offset)

		thumb = image.crop(resize).resize((self._target_image_width, self._target_image_height), Image.LANCZOS)
		return thumb

	def manual_camera_callback(self, msg):
		numpy_array = np.reshape(np.asarray(msg.list), msg.dim)   
		pil_image = Image.fromarray(numpy_array.astype('uint8'), 'RGB')
		bw_pil_image = pil_image.convert('L')	
		obs_im = self.resize(bw_pil_image, msg.dim)
		obs_im_flat = np.array(obs_im).flatten()

		plt.cla()
		plt.imshow(np.array(obs_im)[::-1,:], origin="lower")
		if len(self._img_queue)==self._history_len:
			self._img_queue.popleft()
		self._img_queue.append(obs_im_flat)
		if len(self._img_queue) == self._history_len:
			"""Making overlay"""
			obs_ims = np.tile(np.array(self._img_queue), (self._num_samples, 1, 1))
			if self._viz_mode == "rays": 
				yhats, action_values = self._policy.get_model_outputs(obs_ims, self._obs_vecs, self._prev_actions, self._goal, self._future_actions)
				pred_coll = yhats["coll"]

				# colors = np.random.uniform(0, 1, 20)
				colors = np.mean(pred_coll, axis=1)
				plt.scatter(self._x, self._y, c=colors, s=5, cmap="bwr") # b=0, r=1
			elif self._viz_mode == "rand_proj":
				rand_future_actions = np.random.uniform(low=self._yaw_min, high=self._yaw_max, size=(self._num_samples, self._action_hor, 1))
				yhats, action_values = self._policy.get_model_outputs(obs_ims, self._obs_vecs, self._prev_actions, self._goal, rand_future_actions)
				pred_coll = yhats["coll"]
				# TODO: weird that padding with 0 makes no difference
				pred_x = np.insert(yhats["dx"], 0, np.zeros(self._num_samples), axis=1)
				pred_y = np.insert(yhats["dy"], 0, np.zeros(self._num_samples), axis=1)
			
				points3d = np.zeros((3, self._num_samples*(self._action_hor+1)))
				points3d[0,:] = pred_x.flatten()
				points3d[1,:] = pred_y.flatten()
				points3d[2,:] = self._height*np.ones((1, self._num_samples*(self._action_hor+1)))
				x, y = self.points3d_to_xy(points3d)
				x = np.reshape(x, (self._num_samples, self._action_hor+1)).T
				y = np.reshape(y, (self._num_samples, self._action_hor+1)).T 

				# y = 5*(y - self._target_image_height/2.0) + self._target_image_height/2.0

				colors = self._cm.to_rgba(pred_coll.flatten())
				plt.scatter(x.flatten(), y.flatten(), c=colors)
				plt.plot(x, y, color='b') # b=0, r=1"
			elif self._viz_mode == "rand_unproj":
				rand_future_actions = np.random.uniform(low=self._yaw_min, high=self._yaw_max, size=(self._num_samples, self._action_hor, 1))
				yhats, action_values = self._policy.get_model_outputs(obs_ims, self._obs_vecs, self._prev_actions, self._goal, rand_future_actions)
				pred_coll = yhats["coll"]
				pred_x = yhats["dx"]
				pred_y = yhats["dy"]
				
				scale = 5				
				unproj_x = scale*pred_x + self._target_image_width/2.0 # assume [x, y] = (0, 0) in camera coords
				unproj_y = scale*pred_y + self._target_image_height/4.0
				colors = self._cm.to_rgba(pred_coll.flatten())
				plt.scatter(unproj_x.flatten(), unproj_y.flatten(), c=colors)
				plt.plot(unproj_x.T, unproj_y.T, color='b') # b=0, r=1"""

		plt.pause(0.00001)
		plt.draw()

	def sim_replay_callback(self, msg):

		bw_array = np.squeeze(np.reshape(np.asarray(msg.obs_im), msg.obs_im_dim))
		bw_pil_image = Image.fromarray(bw_array.astype('uint8'), 'L')
		bw_resh_array = np.array(bw_pil_image.resize((self._target_image_width, self._target_image_height), Image.LANCZOS))
		obs_im = bw_resh_array
		obs_im_flat = obs_im.flatten()		

		plt.imshow(np.array(obs_im))
		if len(self._img_queue)==self._history_len:
			self._img_queue.popleft()
		if len(self._obs_queue)==self._history_len:
			self._obs_queue.popleft()
		if len(self._actions_queue)==(self._history_len-1):
			self._actions_queue.popleft()

		self._img_queue.append(obs_im_flat)
		self._obs_queue.append(np.array(msg.obs_vec))
		self._actions_queue.append(np.array(msg.action))
		if len(self._img_queue) == self._history_len:
			"""Making overlay"""
			"""resize and recolor input image, flatten """
			obs_ims = np.tile(np.array(self._img_queue), (self._num_samples, 1, 1))
			obs_vecs = np.tile(np.array(self._obs_queue), (self._num_samples, 1, 1))
			prev_actions = np.tile(np.array(self._actions_queue), (self._num_samples, 1, 1))
			if self._viz_mode == "rays":
				yhats, action_value = self._policy.get_model_outputs(obs_ims, obs_vecs, prev_actions, self._goal, self._future_actions)
				pred_coll = yhats["coll"]

				# Creating image with plotted overlay
				colors = np.mean(pred_coll, axis=1)
				plt.scatter(self._x, self._y, c=colors, s=5, cmap="bwr") # b=0, r=1
			elif self._viz_mode == "rand_proj":
				rand_future_actions = np.random.uniform(low=self._yaw_min, high=self._yaw_max, size=(self._num_samples, self._action_hor, 1))
				yhats, action_values = self._policy.get_model_outputs(obs_ims, obs_vecs, prev_actions, self._goal, rand_future_actions)
				pred_coll = yhats["coll"]
				pred_x = np.insert(yhats["dx"], 0, np.zeros(self._num_samples), axis=1)
				pred_y = np.insert(yhats["dy"], 0, np.zeros(self._num_samples), axis=1)
			
				points3d = np.zeros((3, self._num_samples*(self._action_hor+1)))
				points3d[0,:] = pred_x.flatten()
				points3d[1,:] = pred_y.flatten()
				points3d[2,:] = self._height*np.ones((1, self._num_samples*self._action_hor))
				x, y = self.points3d_to_xy(points3d)
				x = np.reshape(x, (self._num_samples, self._action_hor))
				y = np.reshape(y, (self._num_samples, self._action_hor)) 

				# y = 5*(y - self._target_image_height/2.0) + self._target_image_height/2.0

				colors = self._cm.to_rgba(pred_coll.flatten())
				plt.scatter(x.flatten(), y.flatten(), c=colors)
				plt.plot(x, y, color='b')
			elif self._viz_mode == "rand_unproj":
				rand_future_actions = np.random.uniform(low=self._yaw_min, high=self._yaw_max, size=(self._num_samples, self._action_hor, 1))
				yhats, action_values = self._policy.get_model_outputs(obs_ims, obs_vecs, prev_actions, self._goal, rand_future_actions)
				pred_coll = yhats["coll"]
				pred_x = yhats["dx"]
				pred_y = yhats["dy"]
				
				scale = 5				
				unproj_x = scale*pred_x + self._target_image_width/2.0 # assume [x, y] = (0, 0) in camera coords
				unproj_y = scale*pred_y + self._target_image_height/4.0
				colors = self._cm.to_rgba(pred_coll.flatten())
				plt.scatter(unproj_x.flatten(), unproj_y.flatten(), c=colors)
				plt.plot(unproj_x.T, unproj_y.T, color='b') # b=0, r=1"""

		plt.pause(0.00001)
		plt.draw()

	def run(self): 
		rospy.init_node('model_visualizer', anonymous=True) 
		if self._mode == "sim_replay":
			sub = rospy.Subscriber("obs", obs_message, self.sim_replay_callback)
		elif self._mode == "manual_camera":
			sub = rospy.Subscriber("live_camera_image", camera_message, self.manual_camera_callback)
		elif self._mode == "joystick_roach":
			raise NotImplementedError
		else:
			raise Exception("Not a recognized mode")
		rospy.spin()

if __name__ == '__main__':
	""" To use: fill out mode, the config path, and the training iteration we should load the model from
	Edit ln 65 to folder where you keep the trained models.
	Modes: 
	1) sim_replay: for visualizing on rollouts from sim. Do the rollouts first, record data, publish it via the sim_replay node, and visualize. 
	You can record data like so:
    data = {'obs_ims': [], 'env_infos': [], 'obs_vecs': [], 'actions': []}
	mypickle.dump(data, "/home/siminliu/Desktop/point_mass_cyl_hex_video_debug_{timestamp:s}.pkl".format(timestamp=timestamp))
	2) manual_camera: for moving the camera around by hand. Recommend running camera_publisher, then capturing a rosbag of the video feed, and then playing that back at 0.5 speed and running this node. Bit easier to watch and analyze this way, as opposed to live video feed.  
	3) joystick_roach: <not_implemented_yet> for driving the roach around via joystick and visualizing. 
	Would probably use a rosbag for this too. 

	Other relevant files:
	camera_publisher.py
	sim_replay.py

	Viz mode:
	1) rays: visualize mean collision prob for following rays outward from the origin.
	Advised to use large num_samples (i.e. 20)
	2) rand_proj: visualize step-wise collision prob for following rand action sequences. projected means it plots points where the point mass would be seen on the image. 
	Advised to use small num_samples (i.e. 1)
	3) rand_unproj: same as above, except unproj means it plots points points assuming the x-y plane IS the image plane (in reality, the x-y plane is tilted away from the image plane). This one is easier to look at but may be a bit misleading.
	Advised to use small num_samples (i.e. 10)
	"""
	mode = "manual_camera" 
	exp_name = "roach/non_adaptation/point_planner_H_12_reward"
	eval_itr ="0032"
	viz_mode = "rand_proj"
	mv = ModelVisualizer(exp_name, eval_itr, mode, viz_mode)
	mv.run()