hsr_centersnap_shapo_demo.py

import sys
import argparse
import cv2
import numpy as np

import rosbag
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import numpy as np

bag_file = 'hsr_kinect.bag'

color_topic = '/hsrb/head_rgbd_sensor/rgb/image_rect_color'
depth_topic = '/hsrb/head_rgbd_sensor/depth_registered/image_rect_raw'
color_info_topic = '/hsrb/head_rgbd_sensor/rgb/camera_info'
depth_info_topic = '/hsrb/head_rgbd_sensor/depth_registered/camera_info'

OOB = np.array([0.0, 0.0, 0.0])
def interp2d_color(image, x, y):
  height, width, _ = image.shape
  x_low, y_low = int(np.floor(x)), int(np.floor(y))
  x_alpha = 1 - (x - x_low)
  y_alpha = 1 - (y - y_low)
  x = x_low
  y = y_low
  if x < 0 or x >= (width-1):
    return OOB

  if y < 0 or y >= (height-1):
    return OOB

  val = (image[y  ,x  ,:] * x_alpha *     y_alpha + 
         image[y  ,x+1,:] * (1-x_alpha) * y_alpha +
         image[y+1,x  ,:] * x_alpha *     (1 - y_alpha) + 
         image[y+1,x+1,:] * (1-x_alpha) * (1 - y_alpha))
  return val

def interp2d_depth(image, x, y):
  height, width = image.shape
  x_low, y_low = int(np.floor(x)), int(np.floor(y))
  x_alpha = 1 - (x - x_low)
  y_alpha = 1 - (y - y_low)
  x = x_low
  y = y_low
  if x < 0 or x >= (width-1):
    return 0.0

  if y < 0 or y >= (height-1):
    return 0.0

  val0, w0 = image[y  ,x  ],  x_alpha *     y_alpha
  val1, w1 = image[y  ,x+1],  (1-x_alpha) * y_alpha
  val2, w2 = image[y+1,x  ],  x_alpha *     (1 - y_alpha)
  val3, w3 = image[y+1,x+1],  (1-x_alpha) * (1 - y_alpha)

  if val0 < 0.01: w0 = 0.0
  if val1 < 0.01: w1 = 0.0
  if val2 < 0.01: w2 = 0.0
  if val3 < 0.01: w3 = 0.0
  wtotal = w0 + w1 + w2 + w3
  if wtotal < 0.01:
    return 0.0
  val = val0 * w0 + val1 * w1 + val2 * w2 + val3 * w3
  val /= wtotal
  return val

class Camera:
  fx, fy, cx, cy = 0.0, 0.0, 0.0, 0.0

  def deproject(self, pixels):
    assert pixels.shape[-1] == 2
    point_x = (pixels[...,0] - self.cx) / self.fx
    point_y = (pixels[...,1] - self.cy) / self.fy
    points = np.stack([point_x, point_y], axis=-1)
    assert points.shape[-1] == 2
    return points

  def project(self, points):
    assert points.shape[-1] == 2
    pixels_x = points[...,0] * self.fx + self.cx
    pixels_y = points[...,1] * self.fy + self.cy
    pixels = np.stack([pixels_x, pixels_y], axis=-1)
    assert pixels.shape[-1] == 2
    return pixels

  def reproject(self, input_cam, input_img, is_depth=False):
    if is_depth:
      output_img = np.zeros((self.height, self.width), dtype=np.float64)
    else:
      output_img = np.zeros((self.height, self.width, 3), dtype=np.uint8)
    for row in range(self.height):
      for col in range(self.width):
        out_pixel = np.array([col, row], dtype=np.float64)
        point = self.deproject(out_pixel)
        in_pixel = input_cam.project(point)
        if is_depth:
          output_img[row][col] = interp2d_depth(input_img, in_pixel[0], in_pixel[1])
        else:
          output_img[row][col] = interp2d_color(input_img, in_pixel[0], in_pixel[1])
    return output_img

  def __repr__(self):
    return f'Camera(fx={self.fx}, fy={self.fy}, cx={self.cx}, cy={self.cy}, width={self.width}, height={self.height}'


class NocsReal(Camera):
  height = 480
  width = 640
  fx = 591.0125
  fy = 590.16775
  cx = 322.525
  cy = 244.11084


def make_camera_from_info(msg):
  cam = Camera()
  cam.fx = msg.K[0]
  cam.cx = msg.K[2]
  cam.fy = msg.K[4]
  cam.cy = msg.K[5]
  cam.height = msg.height
  cam.width = msg.width
  return cam

def main():
  bag = rosbag.Bag(bag_file, "r")
  bridge = CvBridge()
  state = 0
  cam_in = None
  cam_out = NocsReal()
  last_color = None
  last_depth = None
  count = 0
  for idx, (topic, msg, t) in enumerate(bag.read_messages()):
    if state == 0:
      if topic != color_info_topic:
        continue
      cam_in = make_camera_from_info(msg)
      print('Found color input camera info')
      print('source camera model:', cam_in)
      print('destination camera model:', cam_out)
      state = 1

    elif state == 1:
      if topic != color_topic:
        continue
      img = bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
      t = msg.header.stamp.to_sec()
      last_color = t, img
      state = 2

    elif state == 2:
      if topic != depth_topic:
        continue
      img = bridge.imgmsg_to_cv2(msg, desired_encoding="passthrough")
      t = msg.header.stamp.to_sec()
      last_depth = t, img
      state = 1

    else:
      pass

    if last_color is None or last_depth is None:
      continue
    dt = np.abs(last_color[0] - last_depth[0])
    max_dt = (1/30.0) * 0.5
    if dt >= max_dt:
      continue
    print('found match!', last_color[0], last_depth[0], dt)
    count += 1
    if count < 10:
      last_color = None
      last_depth = None
      continue

    output_color = cam_out.reproject(cam_in, last_color[1])
    input_depth = last_depth[1] / 1000.0
    output_depth = cam_out.reproject(cam_in, input_depth, is_depth=True)
    output_depth = np.round(output_depth * 1000.0).astype(np.uint16)
    cv2.imwrite('color.png', output_color)
    cv2.imwrite('depth.png', output_depth)
    break
    last_color = None
    last_depth = None
  bag.close()