Gait_Analysis_DeepLabCut/analysis.py at main · farhanaugustine/Gait_Analysis_DeepLabCut · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
# analysis.py
import pandas as pd
import numpy as np
import cv2
import logging
from collections import defaultdict

import config

logger = logging.getLogger(__name__)

def process_data(df, rois):
    if df.empty: return df, pd.DataFrame()
    df = df.sort_values(by=['track_id', 'frame']).reset_index(drop=True)
    logger.info("Calculating pose metrics...")
    df = calculate_pose_metrics(df, config)
    logger.info("Assigning ROIs to detections...")
    df['current_roi'] = assign_rois(df, rois)
    logger.info("Performing gait and step analysis (if keypoints are available)...")
    gait_df = perform_gait_analysis(df, config)
    return df, gait_df

def perform_gait_analysis(df, config):
    paw_names = config.GAIT_PAWS
    if not paw_names:
        logger.warning("Gait analysis is disabled because 'GAIT_PAWS' is empty in the config.")
        return pd.DataFrame()
    if not all(f'{p}_x' in df.columns for p in paw_names):
        logger.warning("Required paw keypoints for gait analysis not found in data. Skipping.")
        return pd.DataFrame()

    for paw in paw_names:
        df[f'{paw}_speed'] = np.sqrt(df.groupby('track_id')[f'{paw}_x'].diff()**2 + df.groupby('track_id')[f'{paw}_y'].diff()**2)
        df[f'{paw}_phase'] = np.where(df[f'{paw}_speed'] < config.PAW_SPEED_THRESHOLD_PX_PER_FRAME, 'stance', 'swing')

    gait_events = []
    for track_id, track_df in df.groupby('track_id'):
        for paw in paw_names:
            phase, prev_phase = track_df[f'{paw}_phase'], track_df[f'{paw}_phase'].shift(1)
            toe_offs = track_df[(phase == 'swing') & (prev_phase == 'stance')]
            foot_strikes = track_df[(phase == 'stance') & (prev_phase == 'swing')]
            for _, row in toe_offs.iterrows():
                gait_events.append({'track_id': track_id, 'frame': row['frame'], 'paw': paw, 'event': 'toe_off', 'x': row[f'{paw}_x'], 'y': row[f'{paw}_y']})
            for _, row in foot_strikes.iterrows():
                gait_events.append({'track_id': track_id, 'frame': row['frame'], 'paw': paw, 'event': 'foot_strike', 'x': row[f'{paw}_x'], 'y': row[f'{paw}_y']})

    if not gait_events: return pd.DataFrame()
    events_df = pd.DataFrame(gait_events).sort_values(by=['track_id', 'frame'])
    all_cycles_df = calculate_all_gait_metrics(events_df, df, config)
    return all_cycles_df

def calculate_all_gait_metrics(events_df, full_df, config):
    all_cycles_data = []
    ref_paw = config.STRIDE_REFERENCE_PAW
    other_paws = [p for p in config.GAIT_PAWS if p != ref_paw]

    # FIX: Make opposing paw selection more specific to avoid errors
    if 'Hip' in ref_paw:
        opposing_paw_name = next((p for p in other_paws if 'Right' in p and 'Hip' in p), None)
    else: # Assumes Shoulder
        opposing_paw_name = next((p for p in other_paws if 'Right' in p and 'Shoulder' in p), None)

    if 'Nose_x' in full_df.columns:
        full_df['body_speed'] = np.sqrt(full_df.groupby('track_id')['Nose_x'].diff()**2 + full_df.groupby('track_id')['Nose_y'].diff()**2)
    else:
        full_df['body_speed'] = 0

    ref_paw_events = events_df[events_df['paw'] == ref_paw]
    for track_id, track_events in ref_paw_events.groupby('track_id'):
        ref_foot_strikes = track_events[track_events['event'] == 'foot_strike'].sort_values('frame')
        for i in range(len(ref_foot_strikes) - 1):
            start_strike, end_strike = ref_foot_strikes.iloc[i], ref_foot_strikes.iloc[i+1]
            stride_length = np.linalg.norm([start_strike['x'] - end_strike['x'], start_strike['y'] - end_strike['y']])
            stride_frames = full_df[(full_df['track_id'] == track_id) & (full_df['frame'] >= start_strike['frame']) & (full_df['frame'] <= end_strike['frame'])]
            stride_speed = stride_frames['body_speed'].mean()
            step_length, step_width = np.nan, np.nan

            if opposing_paw_name:
                opposing_strikes = events_df[(events_df['track_id'] == track_id) & (events_df['paw'] == opposing_paw_name) & (events_df['event'] == 'foot_strike')]
                opposing_strike_df = opposing_strikes[(opposing_strikes['frame'] > start_strike['frame']) & (opposing_strikes['frame'] < end_strike['frame'])]
                if not opposing_strike_df.empty:
                    opposing_strike = opposing_strike_df.iloc[0]
                    step_length = np.linalg.norm([opposing_strike['x'] - start_strike['x'], opposing_strike['y'] - start_strike['y']])
                    p1, p2, p3 = np.array([start_strike['x'], start_strike['y']]), np.array([end_strike['x'], end_strike['y']]), np.array([opposing_strike['x'], opposing_strike['y']])
                    if np.linalg.norm(p2-p1) > 0:
                        step_width = np.abs(np.cross(p2-p1, p1-p3)) / np.linalg.norm(p2-p1)

            all_cycles_data.append({
                'track_id': track_id, 'paw': ref_paw, 'start_frame': start_strike['frame'], 'end_frame': end_strike['frame'],
                'stride_length': stride_length, 'stride_speed': stride_speed, 'step_length': step_length, 'step_width': step_width
            })
    return pd.DataFrame(all_cycles_data)

def calculate_pose_metrics(df, config):
    df['dx'] = df.groupby('track_id')['center_x'].diff()
    df['dy'] = df.groupby('track_id')['center_y'].diff()
    df['speed'] = np.sqrt(df['dx']**2 + df['dy']**2)
    p1_elong, p2_elong = config.ELONGATION_CONNECTION
    if f'{p1_elong}_x' in df.columns and f'{p2_elong}_x' in df.columns:
        df['elongation'] = np.linalg.norm(df[[f'{p1_elong}_x', f'{p1_elong}_y']].values - df[[f'{p2_elong}_x', f'{p2_elong}_y']].values, axis=1)
        df['posture_variability'] = df.groupby('track_id')['elongation'].transform(lambda x: x.rolling(window=30, min_periods=1).std())
    else: df[['elongation', 'posture_variability']] = np.nan
    p1_angle, p2_angle = config.BODY_ANGLE_CONNECTION
    if f'{p1_angle}_x' in df.columns and f'{p2_angle}_x' in df.columns:
        vec = df[[f'{p2_angle}_x', f'{p2_angle}_y']].values - df[[f'{p1_angle}_x', f'{p1_angle}_y']].values
        rad = np.arctan2(vec[:, 1], vec[:, 0])
        df['body_angle_rad'] = rad
        angle_diff = df.groupby('track_id')['body_angle_rad'].diff()
        angle_diff_wrapped = np.arctan2(np.sin(angle_diff), np.cos(angle_diff))
        df['body_angle_deg'] = np.degrees(rad)
        df['turning_speed_rad_per_frame'] = angle_diff_wrapped
        df['turning_speed_deg_per_frame'] = np.degrees(angle_diff_wrapped)
    else: df[['body_angle_rad', 'body_angle_deg', 'turning_speed_rad_per_frame', 'turning_speed_deg_per_frame']] = np.nan
    return df

def assign_rois(df, rois):
    if not rois: return 'None'
    def get_roi(row):
        point = (row['center_x'], row['center_y'])
        if pd.isna(point[0]): return 'None'
        for roi in rois:
            if cv2.pointPolygonTest(roi['coords'], point, False) >= 0: return roi['name']
        return 'None'
    return df.apply(get_roi, axis=1)

def calculate_roi_event_timeline(df):
    timeline = defaultdict(list)
    for _, track_df in df.groupby('track_id'):
        last, current = track_df['current_roi'].shift(1).fillna('None'), track_df['current_roi'].fillna('None')
        for idx, row in track_df[last != current].iterrows():
            frame, lr, cr = int(row['frame']), last.loc[idx], current.loc[idx]
            if lr != 'None': timeline[frame].append({'type': 'exit', 'roi_name': lr})
            if cr != 'None': timeline[frame].append({'type': 'entry', 'roi_name': cr})
    return timeline