Add transducer tracking result to session definition (OpenwaterHealth#182)

Author: sadhana-r

src/openlifu/db/session.py

@@ -5,7 +5,7 @@
 from dataclasses import asdict, dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Optional
 
 import numpy as np
 
@@ -28,6 +28,26 @@ class ArrayTransform:
     (In order to apply the transform to transducer points,
     first represent the points in these units.)"""
 
+@dataclass
+class TransducerTrackingResult:
+    """
+    Class representing the results of running the transducer tracking
+    algorithm.
+    """
+
+    photoscan_id: str
+    """ID of the photoscan object used for transducer tracking"""
+
+    transducer_to_photoscan_transform: ArrayTransform
+    """Transform output by transducer tracking algorithm to register the transducer surface to the photoscan model"""
+
+    photoscan_to_volume_transform: ArrayTransform
+    """Transform output by the transducer tracking algorithm to register the photoscan model the volume's skin segmentation"""
+
+    transducer_tracking_approved: Optional[bool] = False
+    """Approval state of transducer tracking result. `True` means the user has provided some kind of
+    confirmation that the transducer transforms in this result agrees with reality."""
+
 @dataclass
 class Session:
     """
@@ -83,9 +103,8 @@ class Session:
     only. None of the other transforms in the list are considered to be approved.
     """
 
-    transducer_tracking_approved: bool | None = False
-    """Approval state of transducer tracking. `True` means the user has provided some kind of
-    confirmation that the transducer transform in this session agrees with reality."""
+    transducer_tracking_results: Optional[List[TransducerTrackingResult]] = field(default_factory=list)
+    """List of any transducer tracking results"""
 
     def __post_init__(self):
         if self.id is None and self.name is None:
@@ -130,6 +149,11 @@ def from_dict(d:Dict):
             raise ValueError("Sessions no longer recognize a volume attribute -- it is now volume_id.")
         if 'array_transform' in d:
             d['array_transform'] = ArrayTransform(np.array(d['array_transform']['matrix']), d['array_transform']['units'])
+        if 'transducer_tracking_results' in d:
+            d['transducer_tracking_results'] = [TransducerTrackingResult(t['photoscan_id'],
+                                                                         ArrayTransform(np.array(t['transducer_to_photoscan_transform']['matrix']),t['transducer_to_photoscan_transform']['units']),
+                                                                         ArrayTransform(np.array(t['photoscan_to_volume_transform']['matrix']), t['photoscan_to_volume_transform']['units']),
+                                                                         t['transducer_tracking_approved']) for t in d['transducer_tracking_results']]
         if isinstance(d['targets'], list):
             if len(d['targets'])>0 and isinstance(d['targets'][0], dict):
                 d['targets'] = [Point.from_dict(p) for p in d['targets']]
@@ -165,12 +189,13 @@ def to_dict(self):
         d['markers'] = [p.to_dict() for p in d['markers']]
 
         d['array_transform'] = asdict(d['array_transform'])
-
         for target_id,(approval,transforms) in d['virtual_fit_results'].items():
             d['virtual_fit_results'][target_id] = (
                 approval,
                 [asdict(t) for t in transforms],
             )
+        
+        d['transducer_tracking_results'] = [asdict(t) for t in d['transducer_tracking_results']]
 
         return d
 

tests/resources/example_db/subjects/example_subject/sessions/example_session/example_session.json

@@ -43,5 +43,28 @@
       ]
     ]
   },
-  "transducer_tracking_approved": false
+  "transducer_tracking_results": [
+    {
+      "photoscan_id": "example_photoscan",
+      "transducer_to_photoscan_transform": {
+        "matrix": [
+          [1.0, 0.0, 0.0, 0.0],
+          [0.0, 1.0, 0.0, 0.0],
+          [0.0, 0.0, 1.0, 0.0],
+          [0.0, 0.0, 0.0, 1.0]
+        ],
+        "units": "mm"
+      },
+      "photoscan_to_volume_transform": {
+        "matrix": [
+          [1.0, 0.0, 0.0, 0.0],
+          [0.0, 1.0, 0.0, 0.0],
+          [0.0, 0.0, 1.0, 0.0],
+          [0.0, 0.0, 0.0, 1.0]
+        ],
+        "units": "mm"
+      },
+      "transducer_tracking_approved": false
+    }
+  ]
 }

tests/test_database.py

@@ -18,6 +18,7 @@
 from openlifu.db.database import Database, OnConflictOpts
 from openlifu.db.session import ArrayTransform
 from openlifu.photoscan import Photoscan
+from openlifu.db.session import ArrayTransform, TransducerTrackingResult
 from openlifu.plan import Protocol, Run
 from openlifu.xdc import Transducer
 
@@ -55,6 +56,11 @@ def test_new_database(tmp_path:Path):
     assert len(db1.get_subject_ids()) == 0
     assert len(db1.get_transducer_ids()) == 0
 
+@pytest.fixture()
+def example_transducer_tracking_result() -> TransducerTrackingResult:
+    return TransducerTrackingResult(photoscan_id="example_photoscan",
+                                    transducer_to_photoscan_transform = ArrayTransform(np.eye(4),"mm"),
+                                    photoscan_to_volume_transform = ArrayTransform(np.eye(4),"mm"))
 
 def test_write_protocol(example_database: Database):
     protocol = Protocol(name="bleh", id="a_protocol_called_bleh")
@@ -304,6 +310,12 @@ def test_write_session_associated_object_structure_created(example_database: Dat
     assert example_database.get_solutions_filename(example_subject.id, session.id).is_file()
     assert example_database.get_runs_filename(example_subject.id, session.id).is_file()
 
+def test_write_session_with_transducer_tracking_results(example_database: Database, example_subject: Subject, example_transducer_tracking_result):
+    """ Test that when there is a transducer tracking result class associated with a session, the session
+    is correctly written to file."""
+    session = Session(name="bleh", id='a_session',subject_id=example_subject.id)
+    session.transducer_tracking_results = [example_transducer_tracking_result]
+    example_database.write_session(example_subject, session)
 
 def test_write_run(example_database: Database, tmp_path:Path):
     subject_id = "example_subject"