diff --git a/examples/display_scene_camera.py b/examples/display_scene_camera.py
new file mode 100644
index 0000000000000000000000000000000000000000..3b118130def78f160341c0b704d7e17b6d693f76
--- /dev/null
+++ b/examples/display_scene_camera.py
@@ -0,0 +1,52 @@
+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()
diff --git a/misc/matrix_building.py b/misc/matrix_building.py
index 1a34e3f3689c63eb8cdb2b1bbfcbeb4de062a1a0..d1e32bea31bea06f00539817b269042ab3138ff0 100644
--- a/misc/matrix_building.py
+++ b/misc/matrix_building.py
@@ -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
+
+
+
+
diff --git a/rigid_body_models/boat.py b/rigid_body_models/boat.py
index ab525f577896a19cc664a3bfe57d1944b623b912..36ecf4c709a999acf3a1c4cf8f807d2d02a1be59 100644
--- a/rigid_body_models/boat.py
+++ b/rigid_body_models/boat.py
@@ -21,9 +21,10 @@ class MountedThrusters(object):
class DiamondMountedThrusters(object):
"""
+ Body coordinate frame:
^ y
|
- --> x
+ --> x (z-up, so right-handed)
3 0
---------
diff --git a/scene_elements/__init__.py b/scene_elements/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/scene_elements/camera.py b/scene_elements/camera.py
new file mode 100644
index 0000000000000000000000000000000000000000..ffbec9f910263f507e19eadb4854a08ac6271016
--- /dev/null
+++ b/scene_elements/camera.py
@@ -0,0 +1,30 @@
+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)
diff --git a/scene_elements/scene.py b/scene_elements/scene.py
new file mode 100644
index 0000000000000000000000000000000000000000..3fd3f57c4e44a5906afb599838c65614c6cee80c
--- /dev/null
+++ b/scene_elements/scene.py
@@ -0,0 +1,27 @@
+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])
diff --git a/scene_elements/sphere.py b/scene_elements/sphere.py
new file mode 100644
index 0000000000000000000000000000000000000000..f549f2add40560276d5c5527d8c023a5e9d58a41
--- /dev/null
+++ b/scene_elements/sphere.py
@@ -0,0 +1,41 @@
+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)
diff --git a/scene_elements/water.py b/scene_elements/water.py
new file mode 100644
index 0000000000000000000000000000000000000000..91a35b6adfd2fe6457731dec90a75a118099d4b3
--- /dev/null
+++ b/scene_elements/water.py
@@ -0,0 +1,52 @@
+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)