From 5d3f7bc93e4584f69b652d1ff54c7e4aa39ac571 Mon Sep 17 00:00:00 2001
From: neworderofjamie <neworderofjamie@gmail.com>
Date: Fri, 8 Mar 2024 15:26:41 +0000
Subject: [PATCH] build lists of users indices and lighting  in DVS gesture

---
 tonic/datasets/dvsgesture.py | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/tonic/datasets/dvsgesture.py b/tonic/datasets/dvsgesture.py
index a6b48a5a..21e027c5 100644
--- a/tonic/datasets/dvsgesture.py
+++ b/tonic/datasets/dvsgesture.py
@@ -86,13 +86,21 @@ def __init__(
         if not self._check_exists():
             self.download()
 
+        self.users = []
+        self.lighting = []
         file_path = os.path.join(self.location_on_system, self.folder_name)
         for path, dirs, files in os.walk(file_path):
-            dirs.sort()
-            for file in files:
-                if file.endswith("npy"):
-                    self.data.append(path + "/" + file)
-                    self.targets.append(int(file[:-4]))
+            rel_path = os.path.relpath(path, file_path)
+            if rel_path != ".":
+                user, lighting = rel_path.split("_", 1)
+                user = int(user[4:])
+                dirs.sort()
+                for file in files:
+                    if file.endswith("npy"):
+                        self.data.append(os.path.join(path, file))
+                        self.targets.append(int(file[:-4]))
+                        self.users.append(user)
+                        self.lighting.append(lighting)
 
     def __getitem__(self, index):
         """