tracking_tools/setup_config.py

import pandas as pd
import glob
import os
import image_marker as im 
import numpy as np
import tracking_result as tr
from IPython import embed
import matplotlib.pyplot as plt

# wir estellen eine Tabelle, die fuer jeden Tag die Setup Configuration enthält.
# Folgende Spalten:
# Versuchstag, Datum, light_on_area, feeder_positions, feeder_risks, temperature, salinity
x_limits = np.array([0, 1.24])
y_limits = np.array([0, 0.81])

def get_risk(position, light_area_y, x_limits, y_limits, light_risk=1):
    """Calculates the risk associated with a certain position in the arena. 
    
    So far the basic idea is to judge the risk by the distance to the closest wall. That is, the risk is assumed to be maximal
    in the center of the tank. Risk values go from 0 to 1. If the position is in the light-ON area, the risk is increased by 1. 

    Args:
        position (iterable): two-element vector of position i.e. (x,y)
        light_area_y (iterable): two element vactor with the start and stop y-coordinates of light area
        light_risk (float, optional): if position is on the bright side, a malus is added. Defaults to 1.
        x_limits : extent of the tank on x axis in cm
        y_limits : extent of the tank on y axis in cm 

    Returns:
        float: the risk for this position
    """
    min_wall_dist_x = min(np.abs(position[0] - x_limits))
    min_wall_dist_y = min(np.abs(position[1] - y_limits))

    risk_x = 1/(max(x_limits)/2) * min_wall_dist_x
    risk_y = 1/(max(y_limits)/2) * min_wall_dist_y
    total_risk = min(risk_x, risk_y)

    is_position_on_the_bright_side = position[1] >= light_area_y[0] and position[1] < light_area_y[1]
    if is_position_on_the_bright_side:
        total_risk = total_risk + light_risk
    return np.round(total_risk, 3)


def get_feeder_risks(feeder_positions, light_on_area):
    """Calculates the risk associated with each feeder position.

    Args:
        feeder_positions (dict): the feeder positions.
        light_on_area (list): start and stop positions of the light on area e.g. [0, 0.405] m

    Returns:
        dict: the risk for each feeder
    """
    feeder_risks = {}
    if light_on_area['light_center'][1] < light_on_area['left'][1]:
        light_area = [0, light_on_area['left'][1]]
    else:
        light_area = [light_on_area['left'][1], y_limits[1]]
    for k in feeder_positions.keys():
        feeder_risks[k] = get_risk(feeder_positions[k], light_area, x_limits, y_limits)
    return feeder_risks


def get_feeder_positions(vid, num_feeders=8):
    """Ask the user to set markers for the feeders

    Args:
        vid (str): the full path to the video file
        num_feeders (int): the number of feeders that should be marked. Default is 8

    Returns:
        dict: the feeder positions. Already transformed to m 
    """
    feeder_task = im.MarkerTask("Feeder positions", list(map(str, range(1, num_feeders+1))), "Mark feeder positions")
    tasks = [feeder_task]
    image_marker = im.ImageMarker(tasks)
    feeder_positions = image_marker.mark_movie(vid, 100)
    for k in feeder_positions[0].keys(): 
        pos = feeder_positions[0][k] 
        new_pos = tr.coordinate_transformation(pos) 
        feeder_positions[0][k] = np.round(new_pos, 3) 
    return feeder_positions[0]


def get_light_on_area(vid):
    """Ask the user to mark the light ON area. Asking for the left and right ends of the light-dark border and the center position of the ligh on area.

    Args:
        vid (str): The full path to a video.

    Returns:
        dict: The positions already transformed to m.
    """
    light_on_task = im.MarkerTask("Light ON", ["left", "right", "light_center"], "Mark light on area")
    tasks = [light_on_task]
    image_marker = im.ImageMarker(tasks)
    marker_positions = image_marker.mark_movie(vid, 100)
    for k in marker_positions[0].keys(): 
        pos = marker_positions[0][k] 
        new_pos = tr.coordinate_transformation(pos) 
        marker_positions[0][k] = np.round(new_pos, 3)
    return marker_positions[0]
    

def check_day(path):
    """For each day we need the setup configuration, i.e. the light on area and the feeder positions.

    Args:
        path (str): the path to the folder containing the labeled videos and tracking results of that recording day.
    
    Returns:
        dict: the setup config as a dictionary.
    """
    vids = sorted(glob.glob(os.path.join(path, '*.mp4')))
    if len(vids) < 1:
        return
    vid = vids[0]
    
    day = path.split(os.sep)[-1]
    date = vid.split(os.sep)[-1].split("_")[0]

    la = get_light_on_area(vid)
    fp = get_feeder_positions(vid)
    fr = get_feeder_risks(fp, la)
    
    day_config = {"date": date, "day": day, "light_on_left": la["left"], "light_on_right":la["right"], "light_on_center": la["light_center"]} 
    for pos in fp.keys():
        day_config["feeder_pos_%s" % pos] = fp[pos]
    for pos in fr.keys():
        day_config["feeder_risk_%s" % pos] = fr[pos]

    return day_config


if __name__ == '__main__':
    folder = '/mnt/movies/merle_verena/boldness/labeled_videos'
    # folder = "."
    days = sorted(glob.glob(os.path.join(folder, 'day*')))
    configs= []
    for d in days:
        day_config = check_day(d)
        configs.append(day_config)
    
    data_frame = pd.DataFrame(configs)
    data_frame.to_csv("setup_configurations.csv", sep=";")
    """
    x_range = np.arange(0.0, 1.24, .01)
    y_range = np.arange(0.0, .82, .01)
    risk_matrix = np.zeros((len(x_range), len(y_range)))
    light_on = [0.5687952439426905, 0.82]
    #light_on = [0, 0.5687952439426905]
    # light_on = [45.5, 81]

    for i, x in enumerate(x_range):
        for j, y in enumerate(y_range):
            risk_matrix[i, j] = get_risk([x, y], light_on, x_limits=x_limits, y_limits=y_limits )
    plt.imshow(risk_matrix.T, origin='lower')
    plt.show()
    exit()
    """