Treadmill_Paper/organize_data.py at main · leesuyee/Treadmill_Paper · 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
import numpy as np

def organize_data(fly_dataframes):

    # calib = 0.026 - old calibration, updated 03/09/23
    calib_x = 0.0222 # updated calibration 03/09/23
    calib_y = 0.0246


    frame_idx = fly_dataframes['frame_idx'].values.astype(float)

    # thorax
    tx=fly_dataframes['thorax.x'].values.astype(float)*calib_x
    ty=fly_dataframes['thorax.y'].values.astype(float)*calib_y

    # head
    hx=fly_dataframes['head.x'].values.astype(float)*calib_x#-tx
    hy=fly_dataframes['head.y'].values.astype(float)*calib_y#-ty

    # abdomen
    ax=fly_dataframes['abdomen.x'].values.astype(float)*calib_x#-tx
    ay=fly_dataframes['abdomen.y'].values.astype(float)*calib_y#-ty

    # l1
    l1_x=fly_dataframes['L1.x'].values.astype(float)*calib_x#-tx
    l1_y=fly_dataframes['L1.y'].values.astype(float)*calib_y#-ty

    # l2
    l2_x=fly_dataframes['L2.x'].values.astype(float)*calib_x#-tx
    l2_y=fly_dataframes['L2.y'].values.astype(float)*calib_y#-ty

    # l3
    l3_x=fly_dataframes['L3.x'].values.astype(float)*calib_x#-tx
    l3_y=fly_dataframes['L3.y'].values.astype(float)*calib_y#-ty

    # r1
    r1_x=fly_dataframes['R1.x'].values.astype(float)*calib_x#-tx
    r1_y=fly_dataframes['R1.y'].values.astype(float)*calib_y#-ty

    # r2
    r2_x=fly_dataframes['R2.x'].values.astype(float)*calib_x#-tx
    r2_y=fly_dataframes['R2.y'].values.astype(float)*calib_y#-ty

    # r3
    r3_x=fly_dataframes['R3.x'].values.astype(float)*calib_x#-tx
    r3_y=fly_dataframes['R3.y'].values.astype(float)*calib_y#-ty

#     raw_positions = -tx, -ty, -hx, -hy, -ax, -ay, -r1_x, -r1_y, -r2_x, -r2_y, -r3_x, -r3_y, -l1_x, -l1_y, -l2_x, -l2_y, -l3_x,-l3_y
    raw_positions = tx, ty, hx, hy, ax, ay, r1_x, r1_y, r2_x, r2_y, r3_x, r3_y, l1_x, l1_y, l2_x, l2_y, l3_x,l3_y

    return frame_idx, raw_positions