From 2593f21f3a7dbb63078bd991737175bebd12625d Mon Sep 17 00:00:00 2001
From: Jan Grewe <jan.grewe@g-node.org>
Date: Tue, 6 Sep 2022 15:00:43 +0200
Subject: [PATCH] [io] io class for reading dlc h5 data files

---
 etrack/io/__init__.py |  0
 etrack/io/dlc_data.py | 78 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+)
 create mode 100644 etrack/io/__init__.py
 create mode 100644 etrack/io/dlc_data.py

diff --git a/etrack/io/__init__.py b/etrack/io/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/etrack/io/dlc_data.py b/etrack/io/dlc_data.py
new file mode 100644
index 0000000..871e103
--- /dev/null
+++ b/etrack/io/dlc_data.py
@@ -0,0 +1,78 @@
+import os
+import numpy as np
+import pandas as pd 
+import numbers as nb
+
+from ..tracking_data import TrackingData
+
+
+class DLCReader(object):
+    
+    def __init__(self, results_file, crop=(0, 0)) -> None:
+        """
+        If the video data was cropped before tracking and the tracked positions are with respect to the cropped images, we may want to correct for this to convert the data back to absolute positions in the video frame.
+
+        Parameters
+        ----------
+        crop : 2-tuple
+            tuple of (xoffset, yoffset)
+
+        Raises
+        ------
+        ValueError if crop value is not a 2-tuple
+        """
+        if not os.path.exists(results_file):
+            raise ValueError("File %s does not exist!" % results_file)
+        if not isinstance(crop, tuple) or len(crop) < 2:
+            raise ValueError("Cropping info must be a 2-tuple of (x, y)")
+        self._file_name = results_file
+        self._crop = crop
+        self._data_frame = pd.read_hdf(results_file)
+        self._level_shape = self._data_frame.columns.levshape
+        self._scorer = self._data_frame.columns.levels[0].values
+        self._bodyparts = self._data_frame.columns.levels[1].values if self._level_shape[1] > 0 else []
+        self._positions = self._data_frame.columns.levels[2].values if self._level_shape[2] > 0 else []
+
+    @property
+    def filename(self):
+        return self._file_name
+
+    @property
+    def dataframe(self):
+        return self._data_frame
+
+    @property
+    def scorer(self):
+        return self._scorer
+
+    @property
+    def bodyparts(self):
+        return self._bodyparts
+    
+    def _correct_cropping(self, orgx, orgy):
+        x = orgx + self._crop[0]
+        y = orgy + self._crop[1]
+        return x, y
+
+    def track(self, scorer=0, bodypart=0, framerate=30):
+        if isinstance(scorer, nb.Number):
+            sc = self._scorer[scorer]
+        elif isinstance(scorer, str) and scorer in self._scorer:
+            sc = scorer
+        else:
+            raise ValueError(f"Scorer {scorer} is not in dataframe!")
+        if  isinstance(bodypart, nb.Number):
+            bp = self._bodyparts[bodypart]
+        elif isinstance(bodypart, str) and bodypart in self._bodyparts:
+            bp = bodypart
+        else:
+            raise ValueError(f"Body part {bodypart} is not in dataframe!")
+
+        x = np.asarray(self._data_frame[sc][bp]["x"] if "x" in self._positions else [])
+        y = np.asarray(self._data_frame[sc][bp]["y"] if "y" in self._positions else [])
+        x, y = self._correct_cropping(x, y)
+        l = np.asarray(self._data_frame[sc][bp]["likelihood"] if "likelihood" in self._positions else [])
+
+        time = np.arange(len(x))/framerate
+
+        return TrackingData(x, y, time, l, bp, fps=framerate)
\ No newline at end of file