GP2023_chirp_detection/code/CTCPTC.py
2023-01-25 10:19:19 +01:00

99 lines
2.9 KiB
Python

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from tqdm import tqdm
from IPython import embed
from pandas import read_csv
from modules.logger import makeLogger
from modules.plotstyle import PlotStyle
from modules.datahandling import flatten
from modules.behaviour_handling import Behavior, correct_chasing_events, event_triggered_chirps
from extract_chirps import get_valid_datasets
logger = makeLogger(__name__)
ps = PlotStyle()
def get_chirp_winner_loser(folder_name, Behavior, order_meta_df):
foldername = folder_name.split('/')[-2]
winner_row = order_meta_df[order_meta_df['recording'] == foldername]
winner = winner_row['winner'].values[0].astype(int)
winner_fish1 = winner_row['fish1'].values[0].astype(int)
winner_fish2 = winner_row['fish2'].values[0].astype(int)
if winner > 0:
if winner == winner_fish1:
winner_fish_id = winner_row['rec_id1'].values[0]
loser_fish_id = winner_row['rec_id2'].values[0]
elif winner == winner_fish2:
winner_fish_id = winner_row['rec_id2'].values[0]
loser_fish_id = winner_row['rec_id1'].values[0]
chirp_winner = Behavior.chirps[Behavior.chirps_ids == winner_fish_id]
chirp_loser = Behavior.chirps[Behavior.chirps_ids == loser_fish_id]
return chirp_winner, chirp_loser
else:
return None, None
def main(dataroot):
foldernames, _ = get_valid_datasets(dataroot)
meta_path = (
'/').join(foldernames[0].split('/')[:-2]) + '/order_meta.csv'
meta = pd.read_csv(meta_path)
meta['recording'] = meta['recording'].str[1:-1]
winner_chirps = []
loser_chirps = []
onsets = []
offsets = []
physicals = []
# Iterate over all recordings and save chirp- and event-timestamps
for folder in foldernames:
logger.info('Loading data from folder: {}'.format(folder))
time_before = 30
time_after = 60
dt = 0.1
kernel_width = 2
kde_time = np.arange(-time_before, time_after, dt)
broken_folders = ['../data/mount_data/2020-05-12-10_00/']
if folder in broken_folders:
continue
bh = Behavior(folder)
winner, loser = get_chirp_winner_loser(folder, bh, meta)
if winner is None:
continue
# Chirps are already sorted
winner_chirps.append(bh.chirps)
loser_chirps.append(bh.chirps)
# Correct for doubles in chasing on- and offsets to get the right on-/offset pairs
# Get rid of tracking faults (two onsets or two offsets after another)
category, timestamps = correct_chasing_events(bh.behavior, bh.start_s)
# Split categories
onsets.append(timestamps[category == 0])
offsets.append(timestamps[category == 1])
physicals.append(timestamps[category == 2])
# center chirps around events
if __name__ == '__main__':
main('../data/mount_data/')