[setup_config] docstrings, rounding of positions, data_frame output

setup_config.py

@@ -14,7 +14,10 @@ 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
+    """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)
@@ -36,10 +39,19 @@ def get_risk(position, light_area_y, x_limits, y_limits, light_risk=1):
     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 total_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]]
@@ -47,22 +59,39 @@ def get_feeder_risks(feeder_positions, light_on_area):
         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)
-    embed()
+    return feeder_risks
 
 
-def get_feeder_positions(vid):
-    feeder_task = im.MarkerTask("Feeder positions", list(map(str, range(1, 9))), "Mark feeder positions")
+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] = new_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)
@@ -70,31 +99,55 @@ def get_light_on_area(vid):
     for k in marker_positions[0].keys(): 
         pos = marker_positions[0][k] 
         new_pos = tr.coordinate_transformation(pos) 
-        marker_positions[0][k] = new_pos
-    print(marker_positions) 
+        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)
     
-    pass
+    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]
@@ -106,7 +159,4 @@ if __name__ == '__main__':
     plt.show()
     exit()
     """
-    for d in days:
-        check_day(d)
-        print(d)
-        exit()
+