From 4ba665b360d5707cb6d0ac51cfd9e706958ab33e Mon Sep 17 00:00:00 2001
From: glopesdev <glopesdev@users.noreply.github.com>
Date: Fri, 18 Oct 2024 01:52:30 +0100
Subject: [PATCH 1/2] Add interpolated georeference for VR walks

---
 pluma/schema/__init__.py | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/pluma/schema/__init__.py b/pluma/schema/__init__.py
index 7b182b7..9775389 100644
--- 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):

From ac339ed032787f1dcee6cfe8aac2f50832d8bb4b Mon Sep 17 00:00:00 2001
From: glopesdev <glopesdev@users.noreply.github.com>
Date: Fri, 18 Oct 2024 02:05:00 +0100
Subject: [PATCH 2/2] Ensure ZMQ stream type can be set on init

---
 pluma/stream/zeromq.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pluma/stream/zeromq.py b/pluma/stream/zeromq.py
index fe2e9df..0864b58 100644
--- 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