diff --git a/structuretoolkit/visualize.py b/structuretoolkit/visualize.py
index 4f0f67078..e2634f855 100644
--- a/structuretoolkit/visualize.py
+++ b/structuretoolkit/visualize.py
@@ -165,6 +165,16 @@ def _get_box_skeleton(cell: np.ndarray):
     return all_lines @ cell
 
 
+def _draw_box_plotly(fig, structure, px, go):
+    cell = structure.cell
+    data = fig.data
+    for lines in _get_box_skeleton(cell):
+        fig = px.line_3d(**{xx: vv for xx, vv in zip(["x", "y", "z"], lines.T)})
+        fig.update_traces(line_color="#000000")
+        data = fig.data + data
+    return go.Figure(data=data)
+
+
 def _plot3d_plotly(
     structure: Atoms,
     show_cell: bool = True,
@@ -223,12 +233,7 @@ def _plot3d_plotly(
         ),
     )
     if show_cell:
-        data = fig.data
-        for lines in _get_box_skeleton(structure.cell):
-            fig = px.line_3d(**{xx: vv for xx, vv in zip(["x", "y", "z"], lines.T)})
-            fig.update_traces(line_color="#000000")
-            data = fig.data + data
-        fig = go.Figure(data=data)
+        fig = _draw_box_plotly(fig, structure, px, go)
     fig.layout.scene.camera.projection.type = camera
     rot = _get_orientation(view_plane).T
     rot[0, :] *= distance_from_camera * 1.25