From ad7327fba3ecdf3f9b90738c568cc2bc15a35054 Mon Sep 17 00:00:00 2001 From: Tamas Nepusz Date: Fri, 14 Jul 2023 22:25:08 +0200 Subject: [PATCH] refactor: reduced code duplication in marker creation from CSV, SVG or QR code --- src/modules/sbstudio/api/base.py | 4 +- .../operators/add_markers_from_qr_code.py | 22 +---- .../operators/add_markers_from_static_csv.py | 98 +++++-------------- .../plugin/operators/add_markers_from_svg.py | 93 ++++-------------- src/modules/sbstudio/plugin/operators/base.py | 92 +++++++++++++++++ 5 files changed, 145 insertions(+), 164 deletions(-) diff --git a/src/modules/sbstudio/api/base.py b/src/modules/sbstudio/api/base.py index 23cba536..0a11793b 100644 --- a/src/modules/sbstudio/api/base.py +++ b/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, }, } diff --git a/src/modules/sbstudio/plugin/operators/add_markers_from_qr_code.py b/src/modules/sbstudio/plugin/operators/add_markers_from_qr_code.py index a5fb7830..406eed4c 100644 --- a/src/modules/sbstudio/plugin/operators/add_markers_from_qr_code.py +++ b/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) diff --git a/src/modules/sbstudio/plugin/operators/add_markers_from_static_csv.py b/src/modules/sbstudio/plugin/operators/add_markers_from_static_csv.py index 19b840b6..022addfd 100644 --- a/src/modules/sbstudio/plugin/operators/add_markers_from_static_csv.py +++ b/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 diff --git a/src/modules/sbstudio/plugin/operators/add_markers_from_svg.py b/src/modules/sbstudio/plugin/operators/add_markers_from_svg.py index 23cb8165..0a763733 100644 --- a/src/modules/sbstudio/plugin/operators/add_markers_from_svg.py +++ b/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,14 +89,13 @@ 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): @@ -158,8 +103,8 @@ def parse_svg( 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 diff --git a/src/modules/sbstudio/plugin/operators/base.py b/src/modules/sbstudio/plugin/operators/base.py index 35babc19..3dfcc129 100644 --- a/src/modules/sbstudio/plugin/operators/base.py +++ b/src/modules/sbstudio/plugin/operators/base.py @@ -1,6 +1,13 @@ +from abc import abstractmethod +from dataclasses import dataclass +from typing import Optional + from bpy.types import Collection, Operator +from numpy import array +from numpy.typing import NDArray from sbstudio.plugin.errors import StoryboardValidationError +from sbstudio.plugin.model.formation import add_points_to_formation from sbstudio.plugin.selection import select_only @@ -92,3 +99,88 @@ def poll(cls, context): def execute(self, context): entry = context.scene.skybrush.storyboard.active_entry return self.execute_on_storyboard_entry(entry, context) + + +@dataclass +class PointsAndColors: + points: NDArray[float] + """The points to create in a static marker creation operation, in a NumPy + array where each row is a point. + """ + + colors: Optional[NDArray[float]] = None + """Optional colors corresponding to the points in a marker creation + operation, in a NumPy array where the i-th row is the color of the i-th + point in RGB space; color components must be specified in the range [0; 1]. + """ + + +class StaticMarkerCreationOperator(FormationOperator): + """Base class for operators that create a set of markers for a formation, + optionally extended with a list of colors corresponding to the points. + """ + + def execute_on_formation(self, formation, context): + # Construct the point set + try: + points_and_colors = self._create_points(context) + except RuntimeError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + + points = points_and_colors.points + colors = points_and_colors.colors + + if points.shape[0] < 1: + self.report({"ERROR"}, "Formation would be empty, nothing was created") + return {"CANCELLED"} + + # Align the center of the bounding box of the point set to the origin + mins, maxs = points.min(axis=0), points.max(axis=0) + points -= (maxs + mins) / 2 + + # Move the origin of the point set to the 3D cursor + points += array(context.scene.cursor.location, dtype=float) + + # Create the markers + add_points_to_formation(formation, points.tolist()) + + # Add a light effect containing the colors of the markers if needed + if colors is not None: + # 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 + + # 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(colors.flat) + + return {"FINISHED"} + + @abstractmethod + def _create_points(self, context) -> PointsAndColors: + """Creates the points where the markers should be placed.""" + raise NotImplementedError