refactor: reduced code duplication in marker creation from CSV, SVG o…

…r QR code

src/modules/sbstudio/api/base.py

@@ -386,7 +386,7 @@ def export(
     def create_formation_from_svg(
         self,
         source: str,
-        n: int,
+        num_points: int,
         size: float,
     ) -> Tuple[List[Point3D], List[Color3D]]:
         """Samples the path objects of an SVG string into a list of coordinates
@@ -407,7 +407,7 @@ def create_formation_from_svg(
             "parameters": {
                 "version": 1,
                 "source": source,
-                "n": n,
+                "n": num_points,
                 "size": size,
             },
         }

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

@@ -1,13 +1,13 @@
 from bpy.props import EnumProperty, FloatProperty, StringProperty
 
-from sbstudio.plugin.model.formation import add_points_to_formation
+from numpy import array
 
-from .base import FormationOperator
+from .base import StaticMarkerCreationOperator, PointsAndColors
 
 __all__ = ("AddMarkersFromQRCodeOperator",)
 
 
-class AddMarkersFromQRCodeOperator(FormationOperator):
+class AddMarkersFromQRCodeOperator(StaticMarkerCreationOperator):
     """Adds new markers to a formation from the shape of a QR code."""
 
     bl_idname = "skybrush.add_markers_from_qr_code"
@@ -42,8 +42,7 @@ class AddMarkersFromQRCodeOperator(FormationOperator):
     def invoke(self, context, event):
         return context.window_manager.invoke_props_dialog(self)
 
-    def execute_on_formation(self, formation, context):
-        from numpy import array
+    def _create_points(self, context) -> PointsAndColors:
         from sbstudio.vendor.qrcode import (
             QRCode,
             ERROR_CORRECT_L,
@@ -78,16 +77,5 @@ def execute_on_formation(self, formation, context):
             points.shape = (0, 3)
 
         points *= self.spacing
-        min_y, max_y = points[:, 1].min(), points[:, 1].max()
-        min_z, max_z = points[:, 2].min(), points[:, 2].max()
-        delta = array(context.scene.cursor.location, dtype=float) - array(
-            [0, (max_y - min_y) / 2, -(max_z - min_z) / 2], dtype=float
-        )  # type: ignore
-        points += delta
 
-        if points.shape[0] < 1:
-            self.report({"ERROR"}, "Formation would be empty, nothing was created")
-        else:
-            add_points_to_formation(formation, points.tolist())
-
-        return {"FINISHED"}
+        return PointsAndColors(points)

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

@@ -1,20 +1,15 @@
 import csv
 import logging
 
-from dataclasses import dataclass
-from itertools import chain
-from typing import Dict
+from numpy import array, zeros
+from numpy.typing import NDArray
+from typing import Dict, Tuple
 
 from bpy.path import ensure_ext
 from bpy.props import StringProperty
 from bpy_extras.io_utils import ImportHelper
 
-from sbstudio.model.color import Color3D
-from sbstudio.model.point import Point3D
-
-from sbstudio.plugin.model.formation import add_points_to_formation
-
-from .base import FormationOperator
+from .base import StaticMarkerCreationOperator, PointsAndColors
 
 __all__ = ("AddMarkersFromStaticCSVOperator",)
 
@@ -25,13 +20,7 @@
 #############################################################################
 
 
-@dataclass
-class ImportedData:
-    position: Point3D
-    color: Color3D
-
-
-class AddMarkersFromStaticCSVOperator(FormationOperator, ImportHelper):
+class AddMarkersFromStaticCSVOperator(StaticMarkerCreationOperator, ImportHelper):
     """Adds markers from a Skybrush-compatible static .csv file (containing a
     single formation snapshot) to the currently selected formation.
     """
@@ -44,61 +33,28 @@ class AddMarkersFromStaticCSVOperator(FormationOperator, ImportHelper):
     filter_glob = StringProperty(default="*.csv", options={"HIDDEN"})
     filename_ext = ".csv"
 
-    def execute_on_formation(self, formation, context):
-        filepath = ensure_ext(self.filepath, self.filename_ext)
-
-        # get positions and colors from a .csv file
-        try:
-            imported_data = parse_static_csv_zip(filepath, context)
-        except RuntimeError as error:
-            self.report({"ERROR"}, str(error))
-            return {"CANCELLED"}
-
-        # try to figure out the start frame of this formation
-        storyboard_entry = (
-            context.scene.skybrush.storyboard.get_first_entry_for_formation(formation)
-        )
-        frame_start = (
-            storyboard_entry.frame_start
-            if storyboard_entry
-            else context.scene.frame_start
-        )
-        duration = storyboard_entry.duration if storyboard_entry else 1
-
-        # create new markers for the points
-        points = [item.position.as_vector() for item in imported_data.values()]
-        if not points:
-            self.report({"ERROR"}, "Formation would be empty, nothing was created")
-        else:
-            add_points_to_formation(formation, points)
-
-        # add a light effect from the imported colors
-        light_effects = context.scene.skybrush.light_effects
-        if light_effects:
-            light_effects.append_new_entry(
-                name=formation.name,
-                frame_start=frame_start,
-                duration=duration,
-                select=True,
-            )
-            light_effect = light_effects.active_entry
-            light_effect.output = "INDEXED_BY_FORMATION"
-            colors = [item.color.as_vector() for item in imported_data.values()]
-            image = light_effect.create_color_image(
-                name="Image for light effect '{}'".format(formation.name),
-                width=1,
-                height=len(colors),
-            )
-            image.pixels.foreach_set(list(chain(*colors)))
-
-        return {"FINISHED"}
-
     def invoke(self, context, event):
         context.window_manager.fileselect_add(self)
         return {"RUNNING_MODAL"}
 
+    def _create_points(self, context) -> PointsAndColors:
+        filepath = ensure_ext(self.filepath, self.filename_ext)
+        point_color_pairs = parse_static_csv_zip(filepath)
+
+        points = zeros((len(point_color_pairs), 3), dtype=float)
+        colors = zeros((len(point_color_pairs), 3), dtype=float) + 255
+
+        for index, (p, c) in enumerate(point_color_pairs.values()):
+            points[index, :] = p
+            colors[index, :] = c / 255
+
+        return PointsAndColors(points, colors)
+
 
-def parse_static_csv_zip(filename: str, context) -> Dict[str, ImportedData]:
+Item = Tuple[NDArray[float], NDArray[int]]
+
+
+def parse_static_csv_zip(filename: str) -> Dict[str, Item]:
     """Parse a Skybrush static .csv file (containing a list of static positions
     and colors)
 
@@ -113,20 +69,21 @@ def parse_static_csv_zip(filename: str, context) -> Dict[str, ImportedData]:
     Raises:
         RuntimeError: on parse errors
     """
-    result: Dict[str, ImportedData] = {}
+    result: Dict[str, Item] = {}
     header_passed: bool = False
 
     with open(filename, "r") as csv_file:
-        lines = list(csv_file)
-        for row in csv.reader(lines, delimiter=","):
+        for row in csv.reader(csv_file, delimiter=","):
             # skip empty lines
             if not row:
                 continue
+
             # skip possible header line (starting with "Name")
             if not header_passed:
                 header_passed = True
                 if row[0].lower().startswith("name"):
                     continue
+
             # parse line and check for errors
             try:
                 name = row[0]
@@ -146,7 +103,6 @@ def parse_static_csv_zip(filename: str, context) -> Dict[str, ImportedData]:
                 raise RuntimeError(f"Duplicate object name in input CSV file: {name}")
 
             # store position and color entry
-            data = ImportedData(position=Point3D(x, y, z), color=Color3D(r, g, b))
-            result[name] = data
+            result[name] = array((x, y, z), dtype=float), array((r, g, b), dtype=int)
 
     return result

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

@@ -1,21 +1,18 @@
 import logging
 
-from dataclasses import dataclass
-from itertools import chain
-from typing import List, Tuple
+from numpy import array
+from numpy.typing import NDArray
+from typing import Tuple
 
 from bpy.path import ensure_ext
 from bpy.props import FloatProperty, IntProperty, StringProperty
 from bpy_extras.io_utils import ImportHelper
 
 from sbstudio.api.errors import SkybrushStudioAPIError
-from sbstudio.model.color import Color3D
-from sbstudio.model.point import Point3D
 from sbstudio.plugin.api import get_api
-from sbstudio.plugin.model.formation import add_points_to_formation
 from sbstudio.plugin.selection import Collections
 
-from .base import FormationOperator
+from .base import StaticMarkerCreationOperator, PointsAndColors
 
 __all__ = ("AddMarkersFromSVGOperator",)
 
@@ -26,13 +23,7 @@
 #############################################################################
 
 
-@dataclass
-class ImportedData:
-    position: Point3D
-    color: Color3D
-
-
-class AddMarkersFromSVGOperator(FormationOperator, ImportHelper):
+class AddMarkersFromSVGOperator(StaticMarkerCreationOperator, ImportHelper):
     """Adds markers to the currently selected formatio from sampling an SVG file
     with colored path objects. Operator calls the backend API to get the
     sampled positions and colors from the SVG file.
@@ -65,58 +56,6 @@ class AddMarkersFromSVGOperator(FormationOperator, ImportHelper):
     filter_glob = StringProperty(default="*.svg", options={"HIDDEN"})
     filename_ext = ".svg"
 
-    def execute_on_formation(self, formation, context):
-        filepath = ensure_ext(self.filepath, self.filename_ext)
-
-        # get positions and colors from an .svg file
-        try:
-            positions, colors = parse_svg(
-                filepath, n=self.count, size=self.size, context=context
-            )
-        except RuntimeError as error:
-            self.report({"ERROR"}, str(error))
-            return {"CANCELLED"}
-
-        # try to figure out the start frame of this formation
-        storyboard_entry = (
-            context.scene.skybrush.storyboard.get_first_entry_for_formation(formation)
-        )
-        frame_start = (
-            storyboard_entry.frame_start
-            if storyboard_entry
-            else context.scene.frame_start
-        )
-        duration = storyboard_entry.duration if storyboard_entry else 1
-
-        # create new markers for the points
-        points = [point.as_vector() for point in positions]
-        if not points:
-            self.report({"ERROR"}, "Formation would be empty, nothing was created")
-        else:
-            add_points_to_formation(formation, points)
-
-        # add a light effect from the imported colors
-        light_effects = context.scene.skybrush.light_effects
-        if light_effects:
-            light_effects.append_new_entry(
-                name=formation.name,
-                frame_start=frame_start,
-                duration=duration,
-                select=True,
-            )
-            light_effect = light_effects.active_entry
-            light_effect.output = "INDEXED_BY_FORMATION"
-            image = light_effect.create_color_image(
-                name="Image for light effect '{}'".format(formation.name),
-                width=1,
-                height=len(colors),
-            )
-            image.pixels.foreach_set(
-                list(chain(*[list(color.as_vector()) for color in colors]))
-            )
-
-        return {"FINISHED"}
-
     def invoke(self, context, event):
         drones = Collections.find_drones(create=False)
         if drones:
@@ -125,16 +64,23 @@ def invoke(self, context, event):
         context.window_manager.fileselect_add(self)
         return {"RUNNING_MODAL"}
 
+    def _create_points(self, context) -> PointsAndColors:
+        filepath = ensure_ext(self.filepath, self.filename_ext)
+        points, colors = parse_svg(
+            filepath, num_points=self.count, size=self.size, context=context
+        )
+        return PointsAndColors(points, colors)
+
 
 def parse_svg(
-    filename: str, n: int, size: float, context
-) -> Tuple[List[Point3D], List[Color3D]]:
+    filename: str, num_points: int, size: float, context
+) -> Tuple[NDArray[float], NDArray[float]]:
     """Parse an .svg file (containing a list of static positions and colors)
     using the backend API
 
     Args:
         filename: the name of the .svg input file
-        n: the number of points to generate
+        num_points: the number of points to generate
         size: the maximum extent of the points along the main axes
         context: the Blender context
 
@@ -143,23 +89,22 @@ def parse_svg(
 
     Raises:
         SkybrushStudioAPIError: on API parse errors
-
     """
     with open(filename, "r") as svg_file:
         source = svg_file.read()
 
         try:
             points, colors = get_api().create_formation_from_svg(
-                source=source, n=n, size=size
+                source=source, num_points=num_points, size=size
             )
         except Exception as ex:
             if not isinstance(ex, SkybrushStudioAPIError):
                 raise SkybrushStudioAPIError from ex
             else:
                 raise
 
-    # rotate from XY to ZY plane and shift to cursor position
-    center = context.scene.cursor.location
-    points = [Point3D(center.x, p.y + center.y, p.x + center.z) for p in points]
+    # rotate from XY to ZY plane
+    points = array((0, p.y, p.x) for p in points)
+    colors = array((c.r / 255, c.g / 255, c.b / 255) for c in colors)
 
     return points, colors