From fac9ff1d865d565bd36ba9c04195dbcec118b1aa Mon Sep 17 00:00:00 2001
From: Lorenzo <lorenzobertoni9@gmail.com>
Date: Tue, 29 Sep 2020 20:45:22 +0200
Subject: [PATCH] update angles output

---
 monstereo/network/net.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/monstereo/network/net.py b/monstereo/network/net.py
index 1db1102..39ad6ea 100644
--- a/monstereo/network/net.py
+++ b/monstereo/network/net.py
@@ -212,10 +212,15 @@ class Loco:
             dic_out['uv_shoulders'].append(uv_shoulder)
             dic_out['uv_heads'].append(uv_head)
 
+            # For MonStereo / MonoLoco++
+            try:
+                dic_out['angles'].append(float(dic_in['yaw'][0][idx]))  # Predicted angle
+                dic_out['angles_egocentric'].append(float(dic_in['yaw'][1][idx]))  # Egocentric angle
+            except KeyError:
+                continue
+
             # Only for MonStereo
             try:
-                angle = float(dic_in['yaw'][0][idx])  # Predicted angle
-                dic_out['angles'].append(angle)
                 dic_out['aux'].append(float(dic_in['aux'][idx]))
             except KeyError:
                 continue