feat: yaw control basics are added (WIP)

src/modules/sbstudio/api/base.py

@@ -21,6 +21,7 @@
 from sbstudio.model.time_markers import TimeMarkers
 from sbstudio.model.trajectory import Trajectory
 from sbstudio.model.types import Coordinate3D
+from sbstudio.model.yaw import YawSetpointList
 from sbstudio.utils import create_path_and_open
 
 from .constants import COMMUNITY_SERVER_URL
@@ -292,6 +293,7 @@ def export(
         validation: SafetyCheckParams,
         trajectories: Dict[str, Trajectory],
         lights: Optional[Dict[str, LightProgram]] = None,
+        yaw_setpoints: Optional[Dict[str, YawSetpointList]] = None,
         output: Optional[Path] = None,
         show_title: Optional[str] = None,
         show_type: str = "outdoor",
@@ -307,6 +309,7 @@ def export(
             validation: safety check parameters
             trajectories: dictionary of trajectories indexed by drone names
             lights: dictionary of light programs indexed by drone names
+            yaw_setpoints: dictionary of yaw setpoints indexed by drone names
             output: the file path where the output should be saved or `None`
                 if the output must be returned instead of saving it to a file
             show_title: arbitrary show title; `None` if no title is needed
@@ -335,6 +338,9 @@ def export(
         if lights is None:
             lights = {name: LightProgram() for name in trajectories.keys()}
 
+        if yaw_setpoints is None:
+            yaw_setpoints = {name: YawSetpointList() for name in trajectories.keys()}
+
         environment = {"type": show_type}
 
         if time_markers is None:
@@ -363,6 +369,9 @@ def export(
                                     "trajectory": trajectories[name].as_dict(
                                         ndigits=ndigits, version=0
                                     ),
+                                    "yawControl": yaw_setpoints[name].as_dict(
+                                        ndigits=ndigits
+                                    ),
                                 },
                             }
                             for name in natsorted(trajectories.keys())

src/modules/sbstudio/model/types.py

@@ -1,10 +1,13 @@
 from typing import Tuple
 
-__all__ = ("Coordinate3D", "RGBAColor")
+__all__ = ("Coordinate3D", "RGBAColor", "Rotation3D")
 
 
 #: Type alias for simple 3D coordinates
 Coordinate3D = Tuple[float, float, float]
 
 #: Type alias for RGBA color tuples used by Blender
 RGBAColor = Tuple[float, float, float, float]
+
+#: Type alias for simple 3D rotations
+Rotation3D = Tuple[float, float, float]

src/modules/sbstudio/model/yaw.py

@@ -0,0 +1,120 @@
+from dataclasses import dataclass
+from operator import attrgetter
+from typing import (
+    Sequence,
+    TypeVar,
+)
+
+__all__ = (
+    "YawSetpointList",
+    "YawSetpoint",
+)
+
+
+C = TypeVar("C", bound="YawSetpointList")
+
+
+@dataclass
+class YawSetpoint:
+    """The simplest representation of a yaw setpoint."""
+
+    time: float
+    """The timestamp associated to the yaw setpoint, in seconds."""
+
+    angle: float
+    """The yaw angle associated to the yaw setpoint, in degrees."""
+
+
+class YawSetpointList:
+    """Simplest representation of a causal yaw setpoint list in time.
+
+    Setpoints are assumed to be linear, i.e. yaw rate is constant
+    between setpoints.
+    """
+
+    def __init__(self, setpoints: Sequence[YawSetpoint] = []):
+        self.setpoints = sorted(setpoints, key=attrgetter("time"))
+
+    def append(self, setpoint: YawSetpoint) -> None:
+        """Add a setpoint to the end of the setpoint list."""
+        if self.setpoints and self.setpoints[-1].time >= setpoint.time:
+            raise ValueError("New setpoint must come after existing setpoints in time")
+        self.setpoints.append(setpoint)
+
+    def as_dict(self, ndigits: int = 3):
+        """Create a Skybrush-compatible dictionary representation of this
+        instance.
+
+        Parameters:
+            ndigits: round floats to this precision
+
+        Return:
+            dictionary of this instance, to be converted to JSON later
+
+        """
+        return {
+            "setpoints": [
+                [
+                    round(setpoint.time, ndigits=ndigits),
+                    round(setpoint.angle, ndigits=ndigits),
+                ]
+                for setpoint in self.setpoints
+            ],
+            "version": 1,
+        }
+
+    def shift(
+        self: C,
+        delta: float,
+    ) -> C:
+        """Translates the yaw setpoints with the given delta angle. The
+        setpoint list will be manipulated in-place.
+
+        Args:
+            delta: the translation angle
+
+        Returns:
+            The shifted yaw setpoint list
+        """
+
+        self.setpoints = [YawSetpoint(p.time, p.angle + delta) for p in self.setpoints]
+
+        return self
+
+    def simplify(self: C) -> C:
+        """Simplify yaw setpoints in place.
+
+        Returns:
+            the simplified yaw setpoint list
+        """
+        if not self.setpoints:
+            return self
+
+        # set first yaw in the [0, 360) range and shift entire list accordingly
+        angle = self.setpoints[0].angle % 360
+        delta = angle - self.setpoints[0].angle
+        if delta:
+            self.shift(delta)
+
+        # remove intermediate points on constant angular speed segments
+        new_setpoints: list[YawSetpoint] = []
+        last_angular_speed = -1e12
+        for setpoint in self.setpoints:
+            if not new_setpoints:
+                new_setpoints.append(setpoint)
+            else:
+                dt = setpoint.time - new_setpoints[-1].time
+                if dt <= 0:
+                    raise RuntimeError(
+                        f"Yaw timestamps are not causal ({setpoint.time} <= {new_setpoints[-1].time})"
+                    )
+                angular_speed = (setpoint.angle - new_setpoints[-1].angle) / dt
+                if abs(angular_speed - last_angular_speed) < 1e-6:
+                    new_setpoints[-1] = setpoint
+                else:
+                    new_setpoints.append(setpoint)
+                last_angular_speed = angular_speed
+
+        self.setpoints = new_setpoints
+
+        return self

src/modules/sbstudio/plugin/operators/export_to_skyc.py

@@ -1,4 +1,4 @@
-from bpy.props import StringProperty, IntProperty
+from bpy.props import BoolProperty, IntProperty, StringProperty
 
 from sbstudio.model.file_formats import FileFormat
 
@@ -37,6 +37,13 @@ class SkybrushExportOperator(ExportOperator):
         description="Number of samples to take from light programs per second",
     )
 
+    # yaw control enable/disable
+    use_yaw_control = BoolProperty(
+        name="Use yaw control",
+        description="Specifies whether yaw control should be used during the show",
+        default=False,
+    )
+
     def get_format(self) -> FileFormat:
         """Returns the file format that the operator uses. Must be overridden
         in subclasses.
@@ -50,4 +57,5 @@ def get_settings(self):
         return {
             "output_fps": self.output_fps,
             "light_output_fps": self.light_output_fps,
+            "use_yaw_control": self.use_yaw_control,
         }

src/modules/sbstudio/plugin/operators/utils.py

@@ -14,6 +14,7 @@
 from sbstudio.model.light_program import LightProgram
 from sbstudio.model.safety_check import SafetyCheckParams
 from sbstudio.model.trajectory import Trajectory
+from sbstudio.model.yaw import YawSetpointList
 from sbstudio.plugin.constants import Collections
 from sbstudio.plugin.errors import SkybrushStudioExportWarning
 from sbstudio.model.file_formats import FileFormat
@@ -26,6 +27,8 @@
     sample_colors_of_objects,
     sample_positions_of_objects,
     sample_positions_and_colors_of_objects,
+    sample_positions_and_yaw_of_objects,
+    sample_positions_colors_and_yaw_of_objects,
 )
 from sbstudio.plugin.utils.time_markers import get_time_markers_from_context
 
@@ -149,6 +152,85 @@ def _get_trajectories_and_lights(
     return trajectories, lights
 
 
+@with_context
+def _get_trajectories_lights_and_yaw_setpoints(
+    drones,
+    settings: Dict,
+    bounds: Tuple[int, int],
+    *,
+    context: Optional[Context] = None,
+) -> Tuple[Dict[str, Trajectory], Dict[str, LightProgram], Dict[str, YawSetpointList]]:
+    """Get trajectories, LED lights and yaw setpoints of all selected/picked objects.
+
+    Parameters:
+        context: the main Blender context
+        drones: the list of drones to export
+        settings: export settings
+        bounds: the frame range used for exporting
+
+    Returns:
+        dictionary of Trajectory, LightProgram and YawSetpointList objects indexed by object names
+    """
+    trajectory_fps = settings.get("output_fps", 4)
+    light_fps = settings.get("light_output_fps", 4)
+
+    trajectories: Dict[str, Trajectory]
+    lights: Dict[str, LightProgram]
+    yaw_setpoints: Dict[str, YawSetpointList]
+
+    if trajectory_fps == light_fps:
+        # This is easy, we can iterate over the show once
+        with suspended_safety_checks():
+            result = sample_positions_colors_and_yaw_of_objects(
+                drones,
+                frame_range(bounds[0], bounds[1], fps=trajectory_fps, context=context),
+                context=context,
+                by_name=True,
+                simplify=True,
+            )
+
+        trajectories = {}
+        lights = {}
+        yaw_setpoints = {}
+
+        for key, (trajectory, light_program, yaw_curve) in result.items():
+            trajectories[key] = trajectory
+            lights[key] = light_program.simplify()
+            yaw_setpoints[key] = yaw_curve
+
+    else:
+        # We need to iterate over the show twice, once for the trajectories
+        # and yaw setpoints, once for the lights
+        with suspended_safety_checks():
+            with suspended_light_effects():
+                result = sample_positions_and_yaw_of_objects(
+                    drones,
+                    frame_range(
+                        bounds[0], bounds[1], fps=trajectory_fps, context=context
+                    ),
+                    context=context,
+                    by_name=True,
+                    simplify=True,
+                )
+
+                trajectories = {}
+                yaw_setpoints = {}
+
+                for key, (trajectory, yaw_curve) in result.items():
+                    trajectories[key] = trajectory
+                    yaw_setpoints[key] = yaw_curve
+
+            lights = sample_colors_of_objects(
+                drones,
+                frame_range(bounds[0], bounds[1], fps=light_fps, context=context),
+                context=context,
+                by_name=True,
+                simplify=True,
+            )
+
+    return trajectories, lights, yaw_setpoints
+
+
 def export_show_to_file_using_api(
     api: SkybrushStudioAPI,
     context: Context,
@@ -197,11 +279,26 @@ def export_show_to_file_using_api(
                 "There are no objects to export; export cancelled"
             )
 
-    # get trajectories
-    log.info("Getting object trajectories and light programs")
-    trajectories, lights = _get_trajectories_and_lights(
-        drones, settings, frame_range, context=context
-    )
+    # get yaw control enabled state
+    use_yaw_control = settings.get("use_yaw_control", False)
+
+    # get trajectories, light programs and yaw setpoints
+    if use_yaw_control:
+        log.info("Getting object trajectories, light programs and yaw setpoints")
+        (
+            trajectories,
+            lights,
+            yaw_setpoints,
+        ) = _get_trajectories_lights_and_yaw_setpoints(
+            drones, settings, frame_range, context=context
+        )
+    else:
+        log.info("Getting object trajectories and light programs")
+        (
+            trajectories,
+            lights,
+        ) = _get_trajectories_and_lights(drones, settings, frame_range, context=context)
+        yaw_setpoints = None
 
     # get automatic show title
     show_title = str(basename(filepath).split(".")[0])
@@ -280,6 +377,7 @@ def export_show_to_file_using_api(
             validation=validation,
             trajectories=trajectories,
             lights=lights,
+            yaw_setpoints=yaw_setpoints,
             output=filepath,
             time_markers=time_markers,
             renderer=renderer,

src/modules/sbstudio/plugin/tasks/safety_check.py

@@ -10,6 +10,7 @@
 from typing import Iterator
 
 from sbstudio.math.nearest_neighbors import find_nearest_neighbors
+from sbstudio.plugin.utils.evaluator import get_position_of_object
 from sbstudio.plugin.constants import Collections
 from sbstudio.utils import LRUCache
 
@@ -37,10 +38,7 @@ def create_position_snapshot_for_drones_in_collection(collection, *, frame):
     """Create a dictionary mapping the names of the drones in the given
     collection to their positions.
     """
-    return {
-        drone.name: tuple(drone.matrix_world.translation)
-        for drone in collection.objects
-    }
+    return {drone.name: get_position_of_object(drone) for drone in collection.objects}
 
 
 def estimate_velocities_of_drones_at_frame(snapshot, *, frame, scene):