bigmaze/src/data/bonsai_original_smoothing.py

##
import os
import pickle

import numpy as np


"""Reproduce the smoothed centroids that were calculated online in Bonsai and
used to find the center of the cropped image.
"""

smoothed_centroid_dir = ('/media/hdd_data1/michael/deeplabcut/smoothed_centroid_coords')

# TODO: integrate this code with the mx.data.BonsaiData class
crop_length = 300
framerate = 30
interframe_interval = 1000./framerate #  ms

fix_centroid_coord = lambda x: max(min(x, 2048 - crop_length/2 - 1), crop_length/2 + 1)
fix_centroid_coord_vfunc = np.vectorize(fix_centroid_coord)

tracking_data_dir = '/home/michael/data/122014/cv'

all_days = ['15122014',
            '16122014',
            '17122014',
            '18122014',
            '19122014',
            '20122014',
            '21122014',
            '22122014',
            '23122014']

all_mice = ['4I',
            '4II',
            '4III',
            '4IV',
            '5I',
            '5II',
            '5III',
            '5IV']

days = all_days
# file containing start and end times for each mice in the session video in ms
time_shifts_file_path = '/home/michael/code/mousemaze/cache/time_shifts.pickle'

with open(time_shifts_file_path, 'rb') as f:
    times = pickle.load(f)

# TODO: this should be substituted by a for loop for all days
for day in days:
    # First we load the coordinates and smooth it like it was done online in
    # Bonsai. That way we know the position of the cropped frame within the
    # whole frame as it were these smoothed coordinates that were used to
    # decide where to save the crop video.
    tracking_file_path = os.path.join(tracking_data_dir, day, 'tracking.csv')
    tracking_data = np.loadtxt(tracking_file_path, skiprows=1)

    xy_coord = tracking_data[:, 0:2]

    smooth_xy = np.empty_like(xy_coord)

    smooth_xy[0, :] = xy_coord[0, :]

    alpha = 0.3

    for i, xy in enumerate(xy_coord[1:, :]):
        # x and y are smoothed separately
        smooth_xy[i+1] = fix_centroid_coord_vfunc(
            smooth_xy[i] + alpha * (xy - smooth_xy[i]))

    # smooth_xy is now the smoothed coordinates for the whole day. We need
    # to now chop it into a chunk per mouse.
    # We obtain the frame times by multiplying the frame number by the inter
    # frame interval.
    metadata_file_path = os.path.join(tracking_data_dir, day, 'metadata.csv')
    metadata_data = np.loadtxt(metadata_file_path, skiprows=1)
    frame_numbers = metadata_data[:, 1]
    # these don't start at zero so we have to subtract the value of the first
    frame_numbers -= frame_numbers[0]
    frame_times = frame_numbers * interframe_interval

    for mouse in all_mice:
        t_start, t_stop = times[day][mouse]  # in ms

        # we now calculate which index of the data table corresponds to these
        # times by finding the closest
        i_frame_start = np.argmin(np.abs(frame_times - t_start))
        i_frame_stop = np.argmin(np.abs(frame_times - t_stop))

        mouse_frame_times = frame_times[i_frame_start: i_frame_stop + 1]
        # we need to reshape it to add an empty dimension just so we can hstack
        # it later
        mouse_frame_times = np.reshape(
            mouse_frame_times, (mouse_frame_times.size, 1))
        mouse_smoothed_centroid = smooth_xy[i_frame_start: i_frame_stop + 1, :]

        mouse_frame_times_smoothed_centroid = np.hstack(
            (mouse_frame_times, mouse_smoothed_centroid))

        mouse_smoothed_centroid_fname = mouse + '_bonsai_smoothed_coords.pickle'
        mouse_smoothed_centroid_path = os.path.join(
            smoothed_centroid_dir, day, mouse_smoothed_centroid_fname)

        with open(mouse_smoothed_centroid_path, 'wb') as f:
            pickle.dump(mouse_frame_times_smoothed_centroid, f)


##