Skip to content
Snippets Groups Projects
Commit 156f1af2 authored by Lucas Miranda's avatar Lucas Miranda
Browse files

Implemented rotation alignment in preprocess.py

parent a201ddb9
No related branches found
No related tags found
No related merge requests found
Source diff could not be displayed: it is too large. Options to address this: view the blob.
......@@ -232,7 +232,7 @@ DLC_social_1_coords = DLC_social_1.run(verbose=True)
DLC_social_2_coords = DLC_social_2.run(verbose=True)
# Coordinates for training data
coords1 = DLC_social_1_coords.get_coords(center="B_Center", polar=False)
coords1 = DLC_social_1_coords.get_coords(center="B_Center", align="B_Nose")
distances1 = DLC_social_1_coords.get_distances()
angles1 = DLC_social_1_coords.get_angles()
coords_distances1 = merge_tables(coords1, distances1)
......@@ -241,7 +241,7 @@ dists_angles1 = merge_tables(distances1, angles1)
coords_dist_angles1 = merge_tables(coords1, distances1, angles1)
# Coordinates for validation data
coords2 = DLC_social_2_coords.get_coords(center="B_Center", polar=False)
coords2 = DLC_social_2_coords.get_coords(center="B_Center", align="B_Nose")
distances2 = DLC_social_2_coords.get_distances()
angles2 = DLC_social_2_coords.get_angles()
coords_distances2 = merge_tables(coords2, distances2)
......@@ -258,6 +258,7 @@ input_dict_train = {
random_state=42,
filter="gaussian",
sigma=55,
align=True,
),
"dists": distances1.preprocess(
window_size=11,
......@@ -266,6 +267,7 @@ input_dict_train = {
random_state=42,
filter="gaussian",
sigma=55,
align=True,
),
"angles": angles1.preprocess(
window_size=11,
......@@ -274,6 +276,7 @@ input_dict_train = {
random_state=42,
filter="gaussian",
sigma=55,
align=True,
),
"coords+dist": coords_distances1.preprocess(
window_size=11,
......@@ -282,6 +285,7 @@ input_dict_train = {
random_state=42,
filter="gaussian",
sigma=55,
align=True,
),
"coords+angle": coords_angles1.preprocess(
window_size=11,
......@@ -290,6 +294,7 @@ input_dict_train = {
random_state=42,
filter="gaussian",
sigma=55,
align=True,
),
"dists+angle": dists_angles1.preprocess(
window_size=11,
......@@ -298,6 +303,7 @@ input_dict_train = {
random_state=42,
filter="gaussian",
sigma=55,
align=True,
),
"coords+dist+angle": coords_dist_angles1.preprocess(
window_size=11,
......@@ -306,54 +312,60 @@ input_dict_train = {
random_state=42,
filter="gaussian",
sigma=55,
align=True,
),
}
input_dict_val = {
"coords": coords2.preprocess(
window_size=11,
window_step=1,
window_step=10,
scale=True,
random_state=42,
filter="gaussian",
sigma=55,
shuffle=True,
align=True,
),
"dists": distances2.preprocess(
window_size=11,
window_step=1,
window_step=10,
scale=True,
random_state=42,
filter="gaussian",
sigma=55,
shuffle=True,
align=True,
),
"angles": angles2.preprocess(
window_size=11,
window_step=1,
window_step=10,
scale=True,
random_state=42,
filter="gaussian",
sigma=55,
shuffle=True,
align=True,
),
"coords+dist": coords_distances2.preprocess(
window_size=11,
window_step=1,
window_step=10,
scale=True,
random_state=42,
filter="gaussian",
sigma=55,
shuffle=True,
align=True,
),
"coords+angle": coords_angles2.preprocess(
window_size=11,
window_step=1,
window_step=10,
scale=True,
random_state=42,
filter="gaussian",
sigma=55,
shuffle=True,
align=True,
),
"dists+angle": dists_angles2.preprocess(
window_size=11,
......@@ -363,15 +375,17 @@ input_dict_val = {
filter="gaussian",
sigma=55,
shuffle=True,
align=True,
),
"coords+dist+angle": coords_dist_angles2.preprocess(
window_size=11,
window_step=1,
window_step=10,
scale=True,
random_state=42,
filter="gaussian",
sigma=55,
shuffle=True,
align=True,
),
}
......
......@@ -309,7 +309,9 @@ class coordinates:
else:
return "DLC analysis of {} videos".format(len(self._videos))
def get_coords(self, center="arena", polar=False, speed=0, length=None):
def get_coords(
self, center="arena", polar=False, speed=0, length=None, align=False
):
tabs = deepcopy(self._tables)
if polar:
......@@ -384,6 +386,21 @@ class coordinates:
"00:00:00", length, periods=tab.shape[0] + 1, closed="left"
)
if align:
assert (
align in list(tabs.values())[0].columns.levels[0]
), "align must be set to the name of a bodypart"
for key, tab in tabs.items():
# Bring forward the column to align
columns = [i for i in tab.columns if align not in i]
columns = [
(align, ("phi" if polar else "x")),
(align, ("rho" if polar else "y")),
] + columns
tabs[key] = tab[columns]
return table_dict(
tabs,
"coords",
......@@ -524,17 +541,17 @@ class table_dict(dict):
self,
window_size=1,
window_step=1,
scale=True,
scale="standard",
test_proportion=0,
random_state=None,
verbose=False,
filter=None,
sigma=None,
shift=0,
standard_scaler=True,
shuffle=False,
align=False,
):
"""Builds a sliding window. If desired, splits train and test and
"""Builds a sliding window. If specified, splits train and test and
Z-scores the data using sklearn's standard scaler"""
X_train = self.get_training_set()
......@@ -550,16 +567,20 @@ class table_dict(dict):
if verbose:
print("Scaling data...")
if standard_scaler:
if scale == "standard":
scaler = StandardScaler()
else:
elif scale == "minmax":
scaler = MinMaxScaler()
else:
raise ValueError(
"Invalid scaler. Select one of standard, minmax or None"
)
X_train = scaler.fit_transform(
X_train.reshape(-1, X_train.shape[-1])
).reshape(X_train.shape)
if standard_scaler:
if scale == "standard":
assert np.allclose(np.mean(X_train), 0)
assert np.allclose(np.std(X_train), 1)
......@@ -573,6 +594,9 @@ class table_dict(dict):
X_train = rolling_window(X_train, window_size, window_step)
if align:
X_train = align_trajectories(X_train)
if filter == "gaussian":
r = range(-int(window_size / 2), int(window_size / 2) + 1)
r = [i - shift for i in r]
......
......@@ -11,6 +11,7 @@ import pims
import re
import scipy
import seaborn as sns
from copy import deepcopy
from itertools import cycle, combinations, product
from joblib import Parallel, delayed
from numba import jit
......@@ -87,6 +88,32 @@ def angle_trio(array, degrees=False):
return np.array([angle(a, b, c), angle(a, c, b), angle(b, a, c),])
def rotate(p, angles, origin=np.array([0, 0])):
R = np.array([[np.cos(angles), -np.sin(angles)], [np.sin(angles), np.cos(angles)]])
o = np.atleast_2d(origin)
p = np.atleast_2d(p)
return np.squeeze((R @ (p.T - o.T) + o.T).T)
def align_trajectories(data):
data = deepcopy(data)
center_time = (data.shape[1] - 1) // 2
angles = np.arctan2(data[:, center_time, 0], data[:, center_time, 1])
aligned_trajs = np.zeros(data.shape)
for frame in range(data.shape[0]):
aligned_trajs[frame] = rotate(
data[frame].reshape([data.shape[1] * data.shape[2] // 2, 2]), angles[frame],
).reshape(data.shape[1:])
return aligned_trajs
def smooth_boolean_array(a):
"""Returns a boolean array in which isolated appearances of a feature are smoothened"""
for i in range(1, len(a) - 1):
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment