From 5414d5456a8611ef22d20d9d10e2adb86356eb6e Mon Sep 17 00:00:00 2001
From: glopesdev <glopesdev@users.noreply.github.com>
Date: Thu, 12 Oct 2023 22:19:57 +0100
Subject: [PATCH] Add support for grouping pose tracking by model

---
 src/Aeon.Acquisition/LogPoseTracking.bonsai | 134 ++++++++++------
 src/Aeon.Acquisition/PoseTracking.bonsai    | 162 ++++++++++++--------
 2 files changed, 181 insertions(+), 115 deletions(-)

diff --git a/src/Aeon.Acquisition/LogPoseTracking.bonsai b/src/Aeon.Acquisition/LogPoseTracking.bonsai
index 77ac20a..6b992b5 100644
--- a/src/Aeon.Acquisition/LogPoseTracking.bonsai
+++ b/src/Aeon.Acquisition/LogPoseTracking.bonsai
@@ -1,9 +1,9 @@
 <?xml version="1.0" encoding="utf-8"?>
 <WorkflowBuilder Version="2.8.0"
                  xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+                 xmlns:rx="clr-namespace:Bonsai.Reactive;assembly=Bonsai.Core"
                  xmlns:aeon="clr-namespace:Aeon.Acquisition;assembly=Aeon.Acquisition"
                  xmlns:sys="clr-namespace:System;assembly=mscorlib"
-                 xmlns:rx="clr-namespace:Bonsai.Reactive;assembly=Bonsai.Core"
                  xmlns="https://bonsai-rx.org/2018/workflow">
   <Description>Chunks and logs a timestamped pose tracking stream into the base data log.</Description>
   <Workflow>
@@ -14,63 +14,95 @@
       <Expression xsi:type="ExternalizedMapping">
         <Property Name="Address" />
         <Property Name="IdentityIndex" />
+        <Property Name="LogName" />
       </Expression>
-      <Expression xsi:type="Combinator">
-        <Combinator xsi:type="aeon:FormatPose">
-          <aeon:Address>255</aeon:Address>
-          <aeon:IdentityIndex xsi:nil="true" />
-        </Combinator>
-      </Expression>
-      <Expression xsi:type="ExternalizedMapping">
-        <Property Name="Value" DisplayName="LogName" />
-      </Expression>
-      <Expression xsi:type="PropertySource" TypeArguments="WorkflowProperty(sys:String),sys:String">
-        <MemberName>LogName</MemberName>
-        <Value>PoseTrackingData</Value>
-      </Expression>
-      <Expression xsi:type="ExternalizedMapping">
-        <Property Name="Name" />
-      </Expression>
-      <Expression xsi:type="SubscribeSubject">
-        <Name>PoseTrackingMetadata</Name>
-      </Expression>
-      <Expression xsi:type="Combinator">
-        <Combinator xsi:type="rx:Zip" />
-      </Expression>
-      <Expression xsi:type="Combinator">
-        <Combinator xsi:type="rx:Take">
-          <rx:Count>1</rx:Count>
-        </Combinator>
-      </Expression>
-      <Expression xsi:type="Format">
-        <Format>{0}_{1}</Format>
-        <Selector>Item1,Item2.LogName</Selector>
-      </Expression>
-      <Expression xsi:type="PropertyMapping">
-        <PropertyMappings>
-          <Property Name="LogName" />
-        </PropertyMappings>
-      </Expression>
-      <Expression xsi:type="IncludeWorkflow" Path="Aeon.Acquisition:LogHarp.bonsai">
-        <Heartbeats xsi:nil="true" />
-        <ClosingDuration xsi:nil="true" />
-        <LogName>PoseTrackingData</LogName>
+      <Expression xsi:type="rx:SelectMany">
+        <Name>LogPoseModel</Name>
+        <Workflow>
+          <Nodes>
+            <Expression xsi:type="WorkflowInput">
+              <Name>Source1</Name>
+            </Expression>
+            <Expression xsi:type="rx:AsyncSubject">
+              <Name>TrackingGroup</Name>
+            </Expression>
+            <Expression xsi:type="SubscribeSubject">
+              <Name>TrackingGroup</Name>
+            </Expression>
+            <Expression xsi:type="Combinator">
+              <Combinator xsi:type="rx:Merge" />
+            </Expression>
+            <Expression xsi:type="ExternalizedMapping">
+              <Property Name="Address" />
+              <Property Name="IdentityIndex" />
+            </Expression>
+            <Expression xsi:type="Combinator">
+              <Combinator xsi:type="aeon:FormatPose">
+                <aeon:Address>255</aeon:Address>
+                <aeon:IdentityIndex xsi:nil="true" />
+              </Combinator>
+            </Expression>
+            <Expression xsi:type="ExternalizedMapping">
+              <Property Name="Value" DisplayName="LogName" />
+            </Expression>
+            <Expression xsi:type="PropertySource" TypeArguments="WorkflowProperty(sys:String),sys:String">
+              <MemberName>LogName</MemberName>
+              <Value>PoseTrackingData</Value>
+            </Expression>
+            <Expression xsi:type="SubscribeSubject">
+              <Name>TrackingGroup</Name>
+            </Expression>
+            <Expression xsi:type="MemberSelector">
+              <Selector>Key</Selector>
+            </Expression>
+            <Expression xsi:type="Combinator">
+              <Combinator xsi:type="rx:Zip" />
+            </Expression>
+            <Expression xsi:type="Combinator">
+              <Combinator xsi:type="rx:Take">
+                <rx:Count>1</rx:Count>
+              </Combinator>
+            </Expression>
+            <Expression xsi:type="Format">
+              <Format>{0}_{1}</Format>
+              <Selector>Item1,Item2.LogName</Selector>
+            </Expression>
+            <Expression xsi:type="PropertyMapping">
+              <PropertyMappings>
+                <Property Name="LogName" />
+              </PropertyMappings>
+            </Expression>
+            <Expression xsi:type="IncludeWorkflow" Path="Aeon.Acquisition:LogHarp.bonsai">
+              <Heartbeats xsi:nil="true" />
+              <ClosingDuration xsi:nil="true" />
+              <LogName>PoseTrackingData</LogName>
+            </Expression>
+            <Expression xsi:type="WorkflowOutput" />
+          </Nodes>
+          <Edges>
+            <Edge From="0" To="1" Label="Source1" />
+            <Edge From="2" To="3" Label="Source1" />
+            <Edge From="3" To="5" Label="Source1" />
+            <Edge From="4" To="5" Label="Source2" />
+            <Edge From="5" To="14" Label="Source1" />
+            <Edge From="6" To="7" Label="Source1" />
+            <Edge From="7" To="10" Label="Source1" />
+            <Edge From="8" To="9" Label="Source1" />
+            <Edge From="9" To="10" Label="Source2" />
+            <Edge From="10" To="11" Label="Source1" />
+            <Edge From="11" To="12" Label="Source1" />
+            <Edge From="12" To="13" Label="Source1" />
+            <Edge From="13" To="14" Label="Source2" />
+            <Edge From="14" To="15" Label="Source1" />
+          </Edges>
+        </Workflow>
       </Expression>
       <Expression xsi:type="WorkflowOutput" />
     </Nodes>
     <Edges>
       <Edge From="0" To="2" Label="Source1" />
       <Edge From="1" To="2" Label="Source2" />
-      <Edge From="2" To="11" Label="Source1" />
-      <Edge From="3" To="4" Label="Source1" />
-      <Edge From="4" To="7" Label="Source1" />
-      <Edge From="5" To="6" Label="Source1" />
-      <Edge From="6" To="7" Label="Source2" />
-      <Edge From="7" To="8" Label="Source1" />
-      <Edge From="8" To="9" Label="Source1" />
-      <Edge From="9" To="10" Label="Source1" />
-      <Edge From="10" To="11" Label="Source2" />
-      <Edge From="11" To="12" Label="Source1" />
+      <Edge From="2" To="3" Label="Source1" />
     </Edges>
   </Workflow>
 </WorkflowBuilder>
\ No newline at end of file
diff --git a/src/Aeon.Acquisition/PoseTracking.bonsai b/src/Aeon.Acquisition/PoseTracking.bonsai
index 849f6e5..2fb01b9 100644
--- a/src/Aeon.Acquisition/PoseTracking.bonsai
+++ b/src/Aeon.Acquisition/PoseTracking.bonsai
@@ -1,8 +1,8 @@
 <?xml version="1.0" encoding="utf-8"?>
 <WorkflowBuilder Version="2.8.0"
                  xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
-                 xmlns:aeon="clr-namespace:Aeon.Acquisition;assembly=Aeon.Acquisition"
                  xmlns:rx="clr-namespace:Bonsai.Reactive;assembly=Bonsai.Core"
+                 xmlns:aeon="clr-namespace:Aeon.Acquisition;assembly=Aeon.Acquisition"
                  xmlns:sleap="clr-namespace:Bonsai.Sleap;assembly=Bonsai.Sleap"
                  xmlns:harp="clr-namespace:Bonsai.Harp;assembly=Bonsai.Harp"
                  xmlns="https://bonsai-rx.org/2018/workflow">
@@ -12,84 +12,118 @@
       <Expression xsi:type="WorkflowInput">
         <Name>Source1</Name>
       </Expression>
-      <Expression xsi:type="MemberSelector">
-        <Selector>Value.Image</Selector>
-      </Expression>
-      <Expression xsi:type="ExternalizedMapping">
-        <Property Name="Name" DisplayName="PathPrefix" Description="The optional root folder containing all pretrained models." />
-      </Expression>
-      <Expression xsi:type="SubscribeSubject">
-        <Name />
-      </Expression>
       <Expression xsi:type="ExternalizedMapping">
+        <Property Name="PathPrefix" />
         <Property Name="ModelPath" />
-      </Expression>
-      <Expression xsi:type="Combinator">
-        <Combinator xsi:type="aeon:CreatePoseTrackingMetadata">
-          <aeon:ModelPath />
-        </Combinator>
-      </Expression>
-      <Expression xsi:type="ExternalizedMapping">
-        <Property Name="Name" DisplayName="PoseTrackingMetadata" Description="The name of the subject storing model metadata." />
-      </Expression>
-      <Expression xsi:type="rx:AsyncSubject">
-        <Name>PoseTrackingMetadata</Name>
-      </Expression>
-      <Expression xsi:type="PropertyMapping">
-        <PropertyMappings>
-          <Property Name="ModelFileName" Selector="ModelFileName" />
-          <Property Name="TrainingConfig" Selector="TrainingConfig" />
-        </PropertyMappings>
-      </Expression>
-      <Expression xsi:type="ExternalizedMapping">
         <Property Name="CentroidMinConfidence" />
         <Property Name="IdentityMinConfidence" />
         <Property Name="PartMinConfidence" />
       </Expression>
-      <Expression xsi:type="Combinator">
-        <Combinator xsi:type="sleap:PredictPoseIdentities">
-          <sleap:CentroidMinConfidence xsi:nil="true" />
-          <sleap:IdentityMinConfidence xsi:nil="true" />
-          <sleap:PartMinConfidence xsi:nil="true" />
-          <sleap:ScaleFactor xsi:nil="true" />
-          <sleap:ColorConversion xsi:nil="true" />
-        </Combinator>
-      </Expression>
-      <Expression xsi:type="MemberSelector">
-        <Selector>Seconds</Selector>
-      </Expression>
-      <Expression xsi:type="Combinator">
-        <Combinator xsi:type="rx:Zip" />
-      </Expression>
-      <Expression xsi:type="Combinator">
-        <Combinator xsi:type="harp:CreateTimestamped" />
+      <Expression xsi:type="rx:Defer">
+        <Workflow>
+          <Nodes>
+            <Expression xsi:type="ExternalizedMapping">
+              <Property Name="Name" DisplayName="PathPrefix" Description="The optional root folder containing all pretrained models." />
+            </Expression>
+            <Expression xsi:type="SubscribeSubject">
+              <Name />
+            </Expression>
+            <Expression xsi:type="ExternalizedMapping">
+              <Property Name="ModelPath" />
+            </Expression>
+            <Expression xsi:type="Combinator">
+              <Combinator xsi:type="aeon:CreatePoseTrackingMetadata">
+                <aeon:ModelPath />
+              </Combinator>
+            </Expression>
+            <Expression xsi:type="rx:AsyncSubject">
+              <Name>PoseTrackingMetadata</Name>
+            </Expression>
+            <Expression xsi:type="WorkflowInput">
+              <Name>Source1</Name>
+            </Expression>
+            <Expression xsi:type="MemberSelector">
+              <Selector>Value.Image</Selector>
+            </Expression>
+            <Expression xsi:type="SubscribeSubject">
+              <Name>PoseTrackingMetadata</Name>
+            </Expression>
+            <Expression xsi:type="PropertyMapping">
+              <PropertyMappings>
+                <Property Name="ModelFileName" Selector="ModelFileName" />
+                <Property Name="TrainingConfig" Selector="TrainingConfig" />
+              </PropertyMappings>
+            </Expression>
+            <Expression xsi:type="ExternalizedMapping">
+              <Property Name="CentroidMinConfidence" />
+              <Property Name="IdentityMinConfidence" />
+              <Property Name="PartMinConfidence" />
+            </Expression>
+            <Expression xsi:type="Combinator">
+              <Combinator xsi:type="sleap:PredictPoseIdentities">
+                <sleap:CentroidMinConfidence xsi:nil="true" />
+                <sleap:IdentityMinConfidence xsi:nil="true" />
+                <sleap:PartMinConfidence xsi:nil="true" />
+                <sleap:ScaleFactor xsi:nil="true" />
+                <sleap:ColorConversion xsi:nil="true" />
+              </Combinator>
+            </Expression>
+            <Expression xsi:type="MemberSelector">
+              <Selector>Seconds</Selector>
+            </Expression>
+            <Expression xsi:type="Combinator">
+              <Combinator xsi:type="rx:Zip" />
+            </Expression>
+            <Expression xsi:type="Combinator">
+              <Combinator xsi:type="harp:CreateTimestamped" />
+            </Expression>
+            <Expression xsi:type="SubscribeSubject">
+              <Name>PoseTrackingMetadata</Name>
+            </Expression>
+            <Expression xsi:type="Combinator">
+              <Combinator xsi:type="rx:CombineLatest" />
+            </Expression>
+            <Expression xsi:type="rx:GroupBy">
+              <rx:KeySelector>Item2</rx:KeySelector>
+              <rx:ElementSelector>Item1</rx:ElementSelector>
+            </Expression>
+            <Expression xsi:type="WorkflowOutput" />
+          </Nodes>
+          <Edges>
+            <Edge From="0" To="1" Label="Source1" />
+            <Edge From="1" To="3" Label="Source1" />
+            <Edge From="2" To="3" Label="Source2" />
+            <Edge From="3" To="4" Label="Source1" />
+            <Edge From="5" To="6" Label="Source1" />
+            <Edge From="5" To="11" Label="Source1" />
+            <Edge From="6" To="10" Label="Source1" />
+            <Edge From="7" To="8" Label="Source1" />
+            <Edge From="8" To="10" Label="Source2" />
+            <Edge From="9" To="10" Label="Source3" />
+            <Edge From="10" To="12" Label="Source1" />
+            <Edge From="11" To="12" Label="Source2" />
+            <Edge From="12" To="13" Label="Source1" />
+            <Edge From="13" To="15" Label="Source1" />
+            <Edge From="14" To="15" Label="Source2" />
+            <Edge From="15" To="16" Label="Source1" />
+            <Edge From="16" To="17" Label="Source1" />
+          </Edges>
+        </Workflow>
       </Expression>
       <Expression xsi:type="ExternalizedMapping">
         <Property Name="Name" DisplayName="PoseTrackingEvents" Description="The name of the output sequence containing all pose tracking data." />
       </Expression>
-      <Expression xsi:type="rx:PublishSubject">
+      <Expression xsi:type="rx:BehaviorSubject">
         <Name>PoseTrackingEvents</Name>
       </Expression>
       <Expression xsi:type="WorkflowOutput" />
     </Nodes>
     <Edges>
-      <Edge From="0" To="1" Label="Source1" />
-      <Edge From="0" To="11" Label="Source1" />
-      <Edge From="1" To="10" Label="Source1" />
-      <Edge From="2" To="3" Label="Source1" />
-      <Edge From="3" To="5" Label="Source1" />
-      <Edge From="4" To="5" Label="Source2" />
-      <Edge From="5" To="7" Label="Source1" />
-      <Edge From="6" To="7" Label="Source2" />
-      <Edge From="7" To="8" Label="Source1" />
-      <Edge From="8" To="10" Label="Source2" />
-      <Edge From="9" To="10" Label="Source3" />
-      <Edge From="10" To="12" Label="Source1" />
-      <Edge From="11" To="12" Label="Source2" />
-      <Edge From="12" To="13" Label="Source1" />
-      <Edge From="13" To="15" Label="Source1" />
-      <Edge From="14" To="15" Label="Source2" />
-      <Edge From="15" To="16" Label="Source1" />
+      <Edge From="0" To="2" Label="Source1" />
+      <Edge From="1" To="2" Label="Source2" />
+      <Edge From="2" To="4" Label="Source1" />
+      <Edge From="3" To="4" Label="Source2" />
+      <Edge From="4" To="5" Label="Source1" />
     </Edges>
   </Workflow>
 </WorkflowBuilder>
\ No newline at end of file