added postrack script

postrack.ipynb

@@ -0,0 +1,329 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# include controllers to the path\n",
+    "import sys, os\n",
+    "sys.path.append(os.getcwd())\n",
+    "sys.path.append(os.path.join(os.getcwd(), 'controllers'))\n",
+    "\n",
+    "import cv2\n",
+    "import threading\n",
+    "import math\n",
+    "import time\n",
+    "import random\n",
+    "import json\n",
+    "import datetime\n",
+    "import os, shutil\n",
+    "import numpy as np\n",
+    "import multiprocess as mp\n",
+    "\n",
+    "# controllers\n",
+    "import nbimporter\n",
+    "from controllers.situtils import FPSTimes\n",
+    "from controllers.camera import WebcamStream\n",
+    "from controllers.video import VideoWriter\n",
+    "from controllers.position import PositionTracker\n",
+    "from controllers.sound import SoundController\n",
+    "from controllers.serial import MCSArduino, FakeArduino, Feeder\n",
+    "from postprocessing import pack"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load experiment settings\n",
+    "\n",
+    "For every experimental cofiguration you can copy the original 'settings.json' file, build your own specific experimental preset, save it in this folder as e.g. 'settings_elena.json' and load it here instead of 'settings.json'."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cfg_filename = os.path.join('profiles', 'andrey_postrack_test.json')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "with open(cfg_filename) as json_file:\n",
+    "    cfg = json.load(json_file)\n",
+    "cfg['experiment']['experiment_date'] = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Initialize session folder\n",
+    "\n",
+    "Run the upcoming cell, to create a session folder and to save the chosen experimetal parameters to a JSON-file (\"experiment_id_parameters.json\"). The session folder will be created here where this notebook is located."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# This session's protocols will be saved to this folder\n",
+    "cfg_exp = cfg['experiment']\n",
+    "experiment_id = \"%s_%s_%s\" % (cfg_exp['subject'], cfg_exp['experiment_type'], cfg_exp['experiment_date'])\n",
+    "save_to = os.path.join('sessions', experiment_id)\n",
+    "             \n",
+    "if not os.path.exists(save_to):\n",
+    "    os.makedirs(save_to)\n",
+    "\n",
+    "# update paths (assuming this paths are relative to this notebook)\n",
+    "cfg['video']['file_path'] = os.path.join(save_to, 'video.avi')\n",
+    "cfg['position']['file_path'] = os.path.join(save_to, 'positions.csv')\n",
+    "cfg['experiment']['file_path'] = os.path.join(save_to, 'events.csv')\n",
+    "cfg['sound']['file_path'] = os.path.join(save_to, 'sounds.csv')\n",
+    "cfg['position']['background_light'] = os.path.join('assets', cfg['position']['background_light'])\n",
+    "cfg['position']['background_dark'] = os.path.join('assets', cfg['position']['background_dark'])\n",
+    "    \n",
+    "# Saves all parameters to a JSON file with the user-defined \"Experiment ID\" as filename\n",
+    "with open(os.path.join(save_to, experiment_id + '.json'), 'w') as f:\n",
+    "    json.dump(cfg, f, indent=4)\n",
+    "    \n",
+    "with open(cfg['experiment']['file_path'], 'w') as f:\n",
+    "    # state: 0 - trial start, 1 - trial success, 2 - trial fail\n",
+    "    f.write('time, target_x, target_y, target_r, trial, state\\n')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def timeout(t_start):\n",
+    "    return time.time() - t_start > cfg_exp['session_duration'] if t_start is not None else False"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def log_event(*args):  # log start / end of a trial\n",
+    "    with open(cfg_exp['file_path'], 'a') as f:\n",
+    "        f.write(\",\".join([str(x) for x in args]) + \"\\n\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Start the experiment\n",
+    "\n",
+    "This cell contains code for animal tracking. We hope that the comments provided in the code suffice to understand the individual steps and to adjust them to your own setup and needs, if necessary.\n",
+    "\n",
+    "- press 's' to start recording\n",
+    "- press 's' again to stop recording\n",
+    "- press 'q' to quit\n",
+    "\n",
+    "The experiment will stop automatically if the pre-defined session duration is reached."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Webcam stream 1024.0:768.0 at 20.00 FPS started\n"
+     ]
+    }
+   ],
+   "source": [
+    "# actual sound selector: 0 - silence, 1 - foraging, 2 - target, 3 - distractor\n",
+    "sound = mp.Value('i', 1)\n",
+    "\n",
+    "# experiment status: 1 - idle, 2 - running (recording, logging), 0 - stopped\n",
+    "status = mp.Value('i', 1)\n",
+    "\n",
+    "# init the sync with the acquisition system via Arduino\n",
+    "if cfg['experiment']['MCSArduinoPort'] == 'fake':\n",
+    "    board = FakeArduino()\n",
+    "else:\n",
+    "    board = MCSArduino(cfg['experiment']['MCSArduinoPort'])\n",
+    "\n",
+    "# init the feeder\n",
+    "feeder = Feeder(cfg['experiment']['feeder_port'])\n",
+    "    \n",
+    "# start the camera stream\n",
+    "vs = WebcamStream(cfg['camera'])\n",
+    "vs.start()\n",
+    "\n",
+    "# init video recorder\n",
+    "vw = VideoWriter(status, vs, cfg['video'])\n",
+    "vw.start()\n",
+    "\n",
+    "# start position tracking\n",
+    "pt = PositionTracker(status, vs, cfg['position'])\n",
+    "pt.start()\n",
+    "\n",
+    "# playing sound in a separate process for performance\n",
+    "#sc = mp.Process(target=SoundController.run, args=(sound, status, cfg['sound']))\n",
+    "#sc.start()\n",
+    "\n",
+    "timers = []\n",
+    "fps = FPSTimes()\n",
+    "names = ['camera', 'video', 'position', 'main']\n",
+    "trial = 0\n",
+    "rewards = 0\n",
+    "t_start = None\n",
+    "target_since = None\n",
+    "punishment_since = None\n",
+    "trial_start = time.time()\n",
+    "phase = 0  # 0 - idle, 1 - foraging, 2 - inter-trial interval\n",
+    "cfg_exp = cfg['experiment']\n",
+    "cfg_pos = cfg['position']\n",
+    "COLORS = {\n",
+    "    'red': (0,0,255), 'green': (127,255,0), 'blue': (255,127,0), 'yellow': (0,127,255), \\\n",
+    "    'black': (0,0,0), 'white': (255,255,255)\n",
+    "}\n",
+    "islands = []\n",
+    "contour = None\n",
+    "\n",
+    "try:\n",
+    "    while trial <= cfg_exp['trial_number'] and not timeout(t_start):\n",
+    "        frame = vs.read()\n",
+    "        if frame is None:\n",
+    "            time.sleep(0.1)\n",
+    "            continue # wait for the stream\n",
+    "            \n",
+    "        c_time = time.time()\n",
+    "        fps.count()\n",
+    "        status_color = COLORS['green'] if status.value == 1 else COLORS['red']\n",
+    "\n",
+    "        # -------- prepare the video frame ---------------\n",
+    "        \n",
+    "        # mask space outside arena\n",
+    "        frame = cv2.bitwise_and(src1=frame, src2=pt.mask)\n",
+    "        #frame = cv2.subtract(frame, pt.background)\n",
+    "        #frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)\n",
+    "\n",
+    "        # draw position center and contours\n",
+    "        if pt.x_in_px is not None:\n",
+    "            cv2.circle(frame, (pt.x_in_px, pt.y_in_px), 7, status_color, -1)\n",
+    "            cv2.drawContours(frame, [pt.contour], 0, status_color, 1, cv2.LINE_AA)  \n",
+    "\n",
+    "        # add FPS indicators\n",
+    "        for i, ctrl in enumerate([vs, vw, pt, fps]):\n",
+    "            cv2.putText(frame, '%s: %.2f FPS' % (names[i], ctrl.get_avg_fps()), \n",
+    "                     (10, 30 + 20*(i+1)), cv2.FONT_HERSHEY_DUPLEX, .5, COLORS['white'])\n",
+    "\n",
+    "        # size of the arena and status indicator\n",
+    "        cv2.circle(frame, (cfg_pos['arena_x'], cfg_pos['arena_y']), cfg_pos['arena_radius'], COLORS['red'], 2)\n",
+    "        cv2.circle(frame, (cfg_pos['arena_x'], cfg_pos['arena_y']), cfg_pos['floor_radius'], COLORS['red'], 2)\n",
+    "        cv2.circle(frame, (20,20), 10, status_color, -6)\n",
+    "\n",
+    "        # positions of animal and target\n",
+    "        if pt.x_in_px is not None:\n",
+    "            cv2.putText(frame, 'Animal: %.3f %.3f' % (pt.x_in_m, pt.y_in_m), (10, 550), \\\n",
+    "                        cv2.FONT_HERSHEY_DUPLEX, .5, COLORS['white'])\n",
+    "            cv2.putText(frame, '%.2f %.2f' % (pt.x_in_m, pt.y_in_m), (pt.x_in_px + 10, pt.y_in_px + 10), \\\n",
+    "                        cv2.FONT_HERSHEY_DUPLEX, .3, COLORS['white'])\n",
+    "            \n",
+    "        # stopwatch\n",
+    "        stopwatch = 'Time: %.2f' % float(c_time - t_start) if t_start is not None else 'Time: Idle'\n",
+    "        cv2.putText(frame, stopwatch, (10, 590), cv2.FONT_HERSHEY_DUPLEX, .5, COLORS['white'])\n",
+    "\n",
+    "        # rewards\n",
+    "        cv2.putText(frame, 'Rewards: %s' % rewards, (10, 650), cv2.FONT_HERSHEY_DUPLEX, .5, COLORS['white'])\n",
+    "        \n",
+    "        # assign the frame back to the video stream for other controllers\n",
+    "        vs.frame_with_infos = frame\n",
+    "        \n",
+    "        cv2.imshow('Press (s)-to start/stop, (q)-to end', frame)\n",
+    "\n",
+    "        # -------- key press events ---------------\n",
+    "        \n",
+    "        k = cv2.waitKey(33)\n",
+    "        if k == ord('q'):\n",
+    "            break\n",
+    "\n",
+    "        if k == ord('s'):\n",
+    "            board.start_or_stop()  # start/stop data acquisition\n",
+    "            c_time = time.time()   # more accurate time\n",
+    "            \n",
+    "            if status.value == 1: # start the session\n",
+    "                if t_start is None:\n",
+    "                    t_start = c_time\n",
+    "                trial += 1\n",
+    "                log_event(c_time, 0, 0, 0, trial, 0)\n",
+    "                status.value = 2\n",
+    "                \n",
+    "            elif status.value == 2:  # pause the session\n",
+    "                log_event(c_time, 0, 0, 0, trial, -1)\n",
+    "                status.value = 1\n",
+    "                    \n",
+    "        if k == ord('a'):\n",
+    "            sound.value = 0 if sound.value == 1 else 1\n",
+    "\n",
+    "finally:\n",
+    "    if status.value == 2:  # stop data acquisition, in case exited via timeout/trial number/quit\n",
+    "        board.start_or_stop()\n",
+    "        log_event(c_time, 0, 0, 0, trial, -1)\n",
+    "        \n",
+    "    status.value = 0\n",
+    "    time.sleep(0.01)\n",
+    "    \n",
+    "    feeder.exit()\n",
+    "    board.exit()\n",
+    "    cv2.destroyAllWindows()\n",
+    "    for ctrl in [pt, vw, vs]:\n",
+    "        ctrl.stop()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}