diff --git a/pyrad/viewer/server/cameras.py b/pyrad/viewer/server/cameras.py
index 8aef6b802..5a453d010 100644
--- a/pyrad/viewer/server/cameras.py
+++ b/pyrad/viewer/server/cameras.py
@@ -59,7 +59,7 @@ def get_camera_wireframe(scale: float = 0.3, f: int = 4, w: int = 1.5, h: int =
     ll = np.array([-w, -h, -f])
     lr = np.array([w, -h, -f])
     C = np.zeros(3)
-    camera_points = [C, ul, C, ur, C, ll, C, lr, C]
+    camera_points = [C, ul, C, ur, C, ll, C, lr, C, ul, ur, ul, lr, ur, lr, ll, ll, ul]
     lines = np.stack(camera_points) * scale
     return lines
 
@@ -131,6 +131,6 @@ def frustum(
     colors = np.array([color for _ in range(N)])
     line_segments = LineSegments(
         PointsGeometry(position=camera_wireframe_lines.astype(np.float32).T, color=colors.astype(np.float32).T),
-        LineBasicMaterial(vertexColors=True, linewidth=10.0),
+        LineBasicMaterial(vertexColors=True),
     )
     return line_segments
diff --git a/pyrad/viewer/server/geometry.py b/pyrad/viewer/server/geometry.py
index a858a2bbb..c78052a2b 100644
--- a/pyrad/viewer/server/geometry.py
+++ b/pyrad/viewer/server/geometry.py
@@ -393,7 +393,8 @@ class Object(SceneElement):
         self.material = material
         self._type = self._type
 
-    def lower(self, object_data):
+    # pylint: arguments-differ
+    def lower(self):
         data = {
             "metadata": {
                 "version": 4.5,
diff --git a/pyrad/viewer/server/viewer_utils.py b/pyrad/viewer/server/viewer_utils.py
index e509f4af5..e7766e1dd 100644
--- a/pyrad/viewer/server/viewer_utils.py
+++ b/pyrad/viewer/server/viewer_utils.py
@@ -276,9 +276,9 @@ class VisualizerState:
             set_output_options(self.vis, list(outputs.keys()))
             self.outputs_set = True
         # gross hack to get the image key, depending on which keys the graph uses
-        if output_type == "default":
-            output_type = "rgb" if "rgb" in outputs else "rgb_fine"
-        image_output = outputs[output_type].cpu().numpy() * 255
+        if self.prev_output_type == "default":
+            self.prev_output_type = "rgb" if "rgb" in outputs else "rgb_fine"
+        image_output = outputs[self.prev_output_type].cpu().numpy() * 255
         image = (image_output).astype("uint8")
         self.vis["/Cameras/Main Camera"].set_image(image)
 
@@ -526,6 +526,27 @@ def set_camera(vis, camera: Camera):
 def draw_aabb(vis, aabb, name="aabb"):
     """Draw the axis-aligned bounding box."""
     lengths = aabb[1] - aabb[0]
-    vis[name].set_object(g.Box(lengths.tolist()), material=g.MeshPhongMaterial(color=0xFF0000, opacity=0.1))
+
+    w = 1
+    aaa = np.array([w, w, w])
+    aab = np.array([w, w, -w])
+    aba = np.array([w, -w, w])
+    baa = np.array([-w, w, w])
+    abb = np.array([w, -w, -w])
+    bba = np.array([-w, -w, w])
+    bab = np.array([-w, w, -w])
+    bbb = np.array([-w, -w, -w])
+    camera_points = [aaa, aab, aaa, aba, aab, abb, aba, abb]
+    camera_points.extend([baa, bab, baa, bba, bab, bbb, bba, bbb])
+    camera_points.extend([aaa, baa, aab, bab, aba, bba, abb, bbb])
+    lines = np.stack(camera_points)
+
+    lines = lines * np.array(lengths) / 2.0
+
+    line_segments = g.LineSegments(
+        g.PointsGeometry(position=lines.astype(np.float32).T),
+        g.LineBasicMaterial(vertexColors=True),
+    )
+    vis[name].set_object(line_segments)
     center = aabb[0] + lengths / 2.0
     vis[name].set_transform(get_translation_matrix(center.tolist()))