Commit db6d85c9 authored by lucas_miranda's avatar lucas_miranda
Browse files

Implemented get_angles method

parent d59239c5
This diff is collapsed.
...@@ -12,6 +12,7 @@ from sklearn.preprocessing import StandardScaler ...@@ -12,6 +12,7 @@ from sklearn.preprocessing import StandardScaler
from source.utils import * from source.utils import *
from tqdm import tqdm from tqdm import tqdm
import os import os
import networkx as nx
import warnings import warnings
...@@ -35,6 +36,8 @@ class get_coordinates: ...@@ -35,6 +36,8 @@ class get_coordinates:
center_coords=True, center_coords=True,
distances=False, distances=False,
ego=False, ego=False,
angles=False,
connectivity=None,
): ):
self.path = path self.path = path
...@@ -56,6 +59,8 @@ class get_coordinates: ...@@ -56,6 +59,8 @@ class get_coordinates:
self.center_coords = center_coords self.center_coords = center_coords
self.distances = distances self.distances = distances
self.ego = ego self.ego = ego
self.angles = angles
self.connectivity = connectivity
self.scales = self.get_scale self.scales = self.get_scale
assert [re.findall("(.*)_", vid)[0] for vid in self.videos] == [ assert [re.findall("(.*)_", vid)[0] for vid in self.videos] == [
...@@ -181,31 +186,69 @@ class get_coordinates: ...@@ -181,31 +186,69 @@ class get_coordinates:
return distance_dict return distance_dict
def get_angles(self, velocities=0): def get_angles(self, table_dict):
"""Computes the angles between all selected bodyparts over time. """
If ego is provided, it only returns angles to a specified bodypart"""
Computes all the angles between adjacent bodypart trios per video and per frame in the data.
Parameters (from self):
connectivity (dictionary): dict stating to which bodyparts each bodypart is connected;
table_dict (dict of dataframes): tables loaded from the data;
Output:
angle_dict (dictionary): dict containing angle dataframes per vido
"""
bp_net = nx.Graph(self.connectivity)
cliques = nx.enumerate_all_cliques(bp_net)
cliques = [i for i in cliques if len(i) == 3]
angle_dict = {}
for key, tab in table_dict.items():
dats = []
for clique in cliques:
dat = pd.DataFrame(
angle_trio(np.array(tab[clique]).reshape(3, tab.shape[0], 2))
).T
orders = [[0, 1, 2], [0, 2, 1], [1, 0, 2]]
dat.columns = [tuple(clique[i] for i in order) for order in orders]
dats.append(dat)
dats = pd.concat(dats, axis=1)
angle_dict[key] = dats
return angle_dict
def run(self, verbose=1): def run(self, verbose=1):
"""Generates a dataset using all the options specified during initialization""" """Generates a dataset using all the options specified during initialization"""
tables, quality = self.load_tables(verbose) tables, quality = self.load_tables(verbose)
distances = None distances = None
angles = None
if self.distances: if self.distances:
distances = self.get_distances(tables, verbose) distances = self.get_distances(tables, verbose)
if self.angles:
angles = self.get_angles(tables)
if verbose == 1: if verbose == 1:
print("Done!") print("Done!")
return coordinates( return coordinates(
tables, tables=tables,
self.videos, videos=self.videos,
self.arena, arena=self.arena,
self.arena_dims, arena_dims=self.arena_dims,
self.scales, scales=self.scales,
quality, quality=quality,
self.exp_conditions, exp_conditions=self.exp_conditions,
distances, distances=distances,
angles=angles,
) )
...@@ -220,9 +263,11 @@ class coordinates: ...@@ -220,9 +263,11 @@ class coordinates:
quality, quality,
exp_conditions=None, exp_conditions=None,
distances=None, distances=None,
angles=None,
): ):
self._tables = tables self._tables = tables
self.distances = distances self.distances = distances
self.angles = angles
self._videos = videos self._videos = videos
self._exp_conditions = exp_conditions self._exp_conditions = exp_conditions
self._arena = arena self._arena = arena
...@@ -272,6 +317,18 @@ class coordinates: ...@@ -272,6 +317,18 @@ class coordinates:
"Distances not computed. Read the documentation for more details" "Distances not computed. Read the documentation for more details"
) )
def get_angles(self, degrees=False):
if self.angles is not None:
if degrees == True:
return table_dict(
{key: np.degrees(tab) for key, tab in self.angles.items()},
typ="angles",
)
else:
return table_dict(self.angles, typ="angles")
raise ValueError("Angles not computed. Read the documentation for more details")
def get_videos(self, play=False): def get_videos(self, play=False):
if play: if play:
raise NotImplementedError raise NotImplementedError
......
...@@ -12,6 +12,7 @@ import scipy ...@@ -12,6 +12,7 @@ import scipy
import seaborn as sns import seaborn as sns
from itertools import cycle, combinations from itertools import cycle, combinations
from joblib import Parallel, delayed from joblib import Parallel, delayed
from numpy.core.umath_tests import inner1d
from scipy import spatial from scipy import spatial
from sklearn import mixture from sklearn import mixture
from statsmodels.tsa.holtwinters import ExponentialSmoothing from statsmodels.tsa.holtwinters import ExponentialSmoothing
...@@ -59,6 +60,24 @@ def bpart_distance(frame, labels, absdist_arena, pixdist_arena): ...@@ -59,6 +60,24 @@ def bpart_distance(frame, labels, absdist_arena, pixdist_arena):
return dist return dist
def angle(a, b, c):
ba = a - b
bc = c - b
cosine_angle = inner1d(ba, bc) / (
np.linalg.norm(ba, axis=1) * np.linalg.norm(bc, axis=1)
)
angle = np.arccos(cosine_angle)
return angle
def angle_trio(array, degrees=False):
a, b, c = array
return np.array([angle(a, b, c), angle(a, c, b), angle(b, a, c),])
def smooth_boolean_array(a): def smooth_boolean_array(a):
"""Returns a boolean array in which isolated appearances of a feature are smoothened""" """Returns a boolean array in which isolated appearances of a feature are smoothened"""
for i in range(1, len(a) - 1): for i in range(1, len(a) - 1):
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment