Fixed parameter broadcasting in arrus.utils.imaging.SelectFrames.

us4useu · Nov 8, 2021 · 86f3371 · 86f3371
1 parent 44e55eb
commit 86f3371
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 5 deletions.
diff --git a/api/python/arrus/ops/imaging.py b/api/python/arrus/ops/imaging.py
@@ -71,10 +71,10 @@ class SimpleTxRxSequence:
     downsampling_factor: int = 1
     tx_aperture_center_element: list = None
     tx_aperture_center: list = None
-    tx_aperture_size: list = None
+    tx_aperture_size: int = None
     rx_aperture_center_element: list = None
     rx_aperture_center: list = None
-    rx_aperture_size: list = None
+    rx_aperture_size: int = None
     tgc_start: float = None
     tgc_slope: float = None
     tgc_curve: list = None

diff --git a/api/python/arrus/utils/imaging.py b/api/python/arrus/utils/imaging.py
@@ -1218,10 +1218,25 @@ def prepare(self, const_metadata):
         input_n_frames, d2, d3 = input_shape
         output_shape = (n_frames, d2, d3)
         # TODO make this op less prone to changes in op implementation
-        if isinstance(seq, arrus.ops.imaging.PwiSequence):
+        if isinstance(seq, arrus.ops.imaging.SimpleTxRxSequence):
             # select appropriate angles
-            output_angles = seq.angles[self.frames]
-            new_seq = dataclasses.replace(seq, angles=output_angles)
+            angles = self._limit_params(seq.angles, self.frames)
+            tx_focus = self._limit_params(seq.tx_focus, self.frames)
+            tx_aperture_center_element = self._limit_params(seq.tx_aperture_center_element,
+                                                            self.frames)
+            tx_aperture_center = self._limit_params(seq.tx_aperture_center, self.frames)
+            rx_aperture_center_element = self._limit_params(seq.rx_aperture_center_element,
+                                                            self.frames)
+            rx_aperture_center = self._limit_params(seq.rx_aperture_center, self.frames)
+
+            new_seq = dataclasses.replace(
+                seq,
+                angles=angles,
+                tx_focus=tx_focus,
+                tx_aperture_center_element=tx_aperture_center_element,
+                tx_aperture_center=tx_aperture_center,
+                rx_aperture_center_element=rx_aperture_center_element,
+                rx_aperture_center=rx_aperture_center)
             new_context = const_metadata.context
             new_context = arrus.metadata.FrameAcquisitionContext(
                 device=new_context.device, sequence=new_seq,
@@ -1235,6 +1250,12 @@ def prepare(self, const_metadata):
     def process(self, data):
         return data[self.frames]
 
+    def _limit_params(self, value, frames):
+        if value is not None and hasattr(value, "__len__") and len(value) > 1:
+            return value[frames]
+        else:
+            return value
+
 
 class Squeeze(Operation):
     """