Skip to content
Snippets Groups Projects
Commit 81bf57f0 authored by andersct's avatar andersct
Browse files

Add pinhole camera simulation

- add scene elements for visualization
- add example that displays top down and projection to camera image
parent 0ae5c817
No related branches found
No related tags found
No related merge requests found
import numpy as np
from scene_elements.scene import Scene
from scene_elements.sphere import Sphere
from scene_elements.water import Water
from scene_elements.camera import Camera
from misc.matrix_building import rot_yaw_xyz1
import matplotlib.pyplot as plt
def main():
scene = Scene([
Sphere([1, 2], color='green'),
Sphere([-1, 2], color='red'),
Sphere([2, 6], color='yellow'),
Sphere([0, 7], color='black'),
Water(),
])
cam_f = 420. # f * px/m scaling based on sensor size
cam_rows = 480
cam_cols = 640
cam_K = np.array([
[cam_f, 0, cam_rows/2, 0],
[0, cam_f, cam_cols/2, 0],
[0, 0, 1., 0],
])
cam_Rt = np.array([
[0, 0, 1, -.5],
[0, -1, 0, 0],
[1, 0, 0, 0],
[0, 0, 0, 1],
])
camera = Camera(K=cam_K, Rt=cam_Rt, cam_cols=cam_cols, cam_rows=cam_rows)
body2world_Rt = rot_yaw_xyz1(np.pi/(2 + .2))
camera.set_body2world_Rt(body2world_Rt)
# make fig for top view and image view
fig_scene, ax_scene = plt.subplots()
scene.draw_top_view(ax_scene)
ax_scene.set_xlim([-6, 6])
ax_scene.set_ylim([-2, 10])
ax_scene.set_aspect('equal')
plt.show()
fig_scene, ax_image = plt.subplots()
scene.draw_image_view(ax_image, camera)
ax_image.set_aspect('equal')
plt.show()
if __name__ == '__main__':
main()
......@@ -11,3 +11,16 @@ def rot_2d(theta):
mat = np.array([[a, -b], [b, a]])
return mat
def rot_yaw_xyz1(theta):
"""
:param theta: rad
:return: 4, 4 | homogeneous ccw rotation matrix
"""
mat = np.eye(4)
mat[:2, :2] = rot_2d(theta)
return mat
......@@ -21,9 +21,10 @@ class MountedThrusters(object):
class DiamondMountedThrusters(object):
"""
Body coordinate frame:
^ y
|
--> x
--> x (z-up, so right-handed)
3 0
---------
......
import numpy as np
class Camera(object):
def __init__(self, K, Rt, cam_cols=None, cam_rows=None):
"""
Coordinates:
[cam] <--Rt-- [body] <--body2world_Rt-- [world]
:param K: 3, 4 | intrinsic matrix
:param Rt: 4, 4 | extrinsic matrix - body2cam
:param cam_cols:
:param cam_rows:
"""
self.cam_cols = cam_cols
self.cam_rows = cam_rows
self.K = K
self.Rt = Rt
self.P = K.dot(Rt)
self.cam_center = Rt[:3, :3].T.dot(-Rt[:-1, -1])
self.body2world_Rt = np.eye(4)
def set_body2world_Rt(self, Rt):
self.body2world_Rt = Rt
def project(self, x):
return self.P.dot(self.body2world_Rt.T.dot(x))
def get_camera_center_world(self):
return self.body2world_Rt[:3, :3].dot(self.cam_center)
import numpy as np
class Scene(object):
def __init__(self, element_list):
self.element_list = element_list
def draw_top_view(self, ax):
ax.grid(True)
for element in self.element_list:
element.draw_top_view(ax)
def draw_image_view(self, ax, camera):
ax.set_xlim([0, camera.cam_cols])
ax.set_ylim([0, camera.cam_rows])
# ax.grid(True)
# apply world to cam body Rt
cam_center_world = camera.get_camera_center_world()[:2]
dist_from_cam = [np.linalg.norm(element.xy - cam_center_world)
if len(element.xy) > 0 else np.inf
for element in self.element_list]
zorder_list = len(self.element_list) - np.argsort(np.asarray(dist_from_cam))
for i in range(len(self.element_list)):
self.element_list[i].draw_image_view(
ax, camera, zorder=zorder_list[i])
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as pa
class Sphere(object):
def __init__(self, xy, z=0., radius=0.1, color=None):
self.xy = np.asarray(xy)
self.z = z
self.radius = radius
self.color = color
def draw_top_view(self, ax):
circle = pa.Circle(tuple(self.xy), radius=self.radius, facecolor=self.color)
ax.add_artist(circle)
def draw_image_view(self, ax, camera, zorder=0):
"""
Do not draw if behind camera.
(Out of frame objects can be drawn, they will just be cut off)
:param ax:
:param camera:
:param zorder: ordering for fake depth
:return:
"""
xyz1 = np.hstack([self.xy, self.z, 1.])
top_bot_xyz1 = np.array([xyz1, xyz1]).T
top_bot_xyz1[2, :] += [self.radius, -self.radius]
top_bot_xy1 = camera.project(top_bot_xyz1)
if np.any(top_bot_xy1[-1] <= 0):
return
top_bot_xy1 /= top_bot_xy1[-1]
xy = top_bot_xy1[:-1].mean(axis=1)
r = np.linalg.norm(top_bot_xy1[:-1, 0] - xy)/2
# reverse xy to place into plot coordinates, since cam x is vertical
circle = pa.Circle(
tuple(xy)[::-1], radius=r,
facecolor=self.color, zorder=zorder
)
ax.add_artist(circle)
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as pa
class Water(object):
def __init__(self, z=0., bound_xy=()):
self.xy = ()
self.z = z
self.color = '#80ebff'
self.bound_xy = bound_xy
if len(bound_xy) == 0:
theta = np.linspace(0, 2*np.pi, num=720, endpoint=False)
self.bound_xy = np.array([
np.cos(theta), np.sin(theta)
]) * 1e5
def draw_top_view(self, ax):
m = 200
rect = pa.Rectangle((-m/2, -m/2), width=m, height=m, facecolor=self.color, zorder=-1)
ax.add_artist(rect)
def draw_image_view(self, ax, camera, zorder=-1):
"""
Do not draw if behind camera.
(Out of frame objects can be drawn, they will just be cut off)
:param ax:
:param camera:
:param zorder: ordering for fake depth
:return:
"""
bound_xyz1 = np.vstack([self.bound_xy, 0*self.bound_xy])
bound_xyz1[2, :] = self.z
bound_xyz1[3, :] = 1
bound_xy1 = camera.project(bound_xyz1)
bound_xy1 = bound_xy1[:, bound_xy1[-1] > 1e-2]
if bound_xy1.size == 0:
return
bound_xy1 /= bound_xy1[-1]
# use furthest point on each side of camera fov
# lower two points given by assumption we are surrounded by water
ind_min = bound_xy1[1].argmin()
ind_max = bound_xy1[1].argmax()
verts = np.asarray([
[0, 0],
bound_xy1[:2, ind_min],
bound_xy1[:2, ind_max],
[0, camera.cam_cols],
])
poly = pa.Polygon(verts[:, ::-1], facecolor=self.color, zorder=zorder)
ax.add_artist(poly)
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment