Merge pull request #67 from emotional-cities/unity-georeference

Add interpolated georeference for VR walks
emotional-cities · Oct 18, 2024 · e834b8a · e834b8a
2 parents d983e72 + ac339ed
commit e834b8a
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 1 deletion.
diff --git a/pluma/schema/__init__.py b/pluma/schema/__init__.py
@@ -6,9 +6,13 @@
 from dotmap import DotMap
 from typing import Union, Optional, Callable
 
+from pandas import DataFrame
+from sklearn.linear_model import LinearRegression
+
 from pluma.export.streams import shift_stream_index
 from pluma.schema.outdoor import build_schema
 
+from pluma.stream.zeromq import UnityGeoreferenceStream, UnityTransformStream
 from pluma.sync.ubx2harp import (
     SyncLookup,
     get_clockcalibration_model,
@@ -90,6 +94,32 @@ def add_ubx_georeference(
         if calibrate_clock is True:
             self.georeference.clockreference.referenceid = ubxstream.clockreference.referenceid
 
+    def add_unity_georeference(
+        self, positionstream: UnityTransformStream, georeferencestream: UnityGeoreferenceStream, strip=True
+    ):
+        if positionstream.streamtype != StreamType.UNITY:
+            raise TypeError("Position must be a Unity Stream.")
+
+        if georeferencestream.streamtype != StreamType.UNITY:
+            raise TypeError("Georeference must be a Unity Stream.")
+
+        posdata = positionstream.data
+        geodata = georeferencestream.data
+        lon_model = LinearRegression().fit(geodata[["TargetPositionZ"]].values, geodata[["TargetLongitude"]])
+        lat_model = LinearRegression().fit(geodata[["TargetPositionX"]].values, geodata[["TargetLatitude"]])
+        alt_model = LinearRegression().fit(geodata[["TargetPositionY"]].values, geodata[["TargetAltitude"]])
+        navdata = DataFrame(
+            data={
+                "Longitude": lon_model.predict(posdata[["Transform.Position.Z"]]).reshape(-1),
+                "Latitude": lat_model.predict(posdata[["Transform.Position.X"]]).reshape(-1),
+                "Elevation": alt_model.predict(posdata[["Transform.Position.Y"]]).reshape(-1),
+            },
+            index=posdata.index,
+        )
+        self.georeference.from_dataframe(navdata)
+        if strip is True:
+            self.georeference.strip()
+
     @staticmethod
     def _iter_schema_streams(schema: Union[DotMap, Stream, None] = None):
         if isinstance(schema, Stream):

diff --git a/pluma/stream/zeromq.py b/pluma/stream/zeromq.py
@@ -16,10 +16,10 @@ def __init__(
         dtypes: list[tuple[str, type]],
         **kw,
     ):
-        self.streamtype = streamtype
         self.filenames = filenames
         self.dtypes = dtypes
         super(ZmqStream, self).__init__(eventcode, **kw)
+        self.streamtype = streamtype
 
     def resample(self):
         pass