Skip to content

Instantly share code, notes, and snippets.

@zubair-irshad
Created July 4, 2023 18:02
Show Gist options
  • Save zubair-irshad/a6ddbef3b4112a259b463794502f8f27 to your computer and use it in GitHub Desktop.
Save zubair-irshad/a6ddbef3b4112a259b463794502f8f27 to your computer and use it in GitHub Desktop.
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment