{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import os, json\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "import nbimporter\n",
    "from controllers.situtils import FPSTimes\n",
    "from controllers.serial import MCSArduino"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Capture a background image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "is_light = True\n",
    "#settings_filename = os.path.join('profiles', 'kate_postrack.json')\n",
    "settings_filename = os.path.join('profiles', 'kate_aSIT_single.json')\n",
    "#settings_filename = os.path.join('profiles', 'andrey_hippoSIT_prod.json')\n",
    "\n",
    "with open(settings_filename) as json_file:\n",
    "    cfg = json.load(json_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "board = MCSArduino(cfg['experiment']['MCSArduinoPort'])\n",
    "\n",
    "# Define BGR colors\n",
    "BGR_COLOR = {\n",
    "    'red': (0,0,255),\n",
    "    'green': (127,255,0),\n",
    "    'blue': (255,127,0),\n",
    "    'yellow': (0,127,255),\n",
    "    'black': (0,0,0),\n",
    "    'white': (255,255,255)\n",
    "}\n",
    "\n",
    "cfg_exp = cfg['experiment']\n",
    "cfg_cam = cfg['camera']\n",
    "cfg_pos = cfg['position']\n",
    "\n",
    "cap = cv2.VideoCapture(cfg_cam['source'], cfg_cam['api']) if cfg_cam['api'] else cv2.VideoCapture(cfg_cam['source'])\n",
    "cap.set(cv2.CAP_PROP_FPS, cfg_cam['fps'])\n",
    "time.sleep(1)  # this helps to keep the FPS stable\n",
    "cap.set(cv2.CAP_PROP_FRAME_WIDTH, cfg_cam['frame_width'])\n",
    "cap.set(cv2.CAP_PROP_FRAME_HEIGHT, cfg_cam['frame_height'])\n",
    "\n",
    "cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'))    \n",
    "\n",
    "fps = FPSTimes()\n",
    "\n",
    "while(True):\n",
    "    ret, frame = cap.read()\n",
    "    fps.count()\n",
    "\n",
    "    # Draw the arena area\n",
    "    cv2.circle(frame, (cfg_pos['arena_x'], cfg_pos['arena_y']), cfg_pos['arena_radius'], BGR_COLOR['red'], 2)\n",
    "\n",
    "    # Mask the space outside the arena\n",
    "    mask = np.zeros(shape=frame.shape, dtype=\"uint8\")\n",
    "    cv2.circle(mask, (cfg_pos['arena_x'], cfg_pos['arena_y']), cfg_pos['arena_radius'], BGR_COLOR['white'], -1)\n",
    "    \n",
    "    frame = cv2.bitwise_and(src1=frame, src2=mask)\n",
    "    \n",
    "    k = cv2.waitKey(33)\n",
    "    if k == ord('c'):\n",
    "        f_name = cfg_pos['background_light'] if is_light else cfg_pos['background_dark']\n",
    "        cv2.imwrite(os.path.join('assets', f_name), frame)\n",
    "    if k == ord('l'):\n",
    "        board.switch_light()\n",
    "        is_light = not is_light\n",
    "    if k == ord('q'):  # full quit\n",
    "        break\n",
    "        \n",
    "    cv2.putText(frame, 'LIGHT' if is_light else 'DARK', (10, 20), cv2.FONT_HERSHEY_DUPLEX, .5, BGR_COLOR['white'])\n",
    "    cv2.putText(frame, '%.2f FPS' % fps.get_avg_fps(), (10, 40), cv2.FONT_HERSHEY_DUPLEX, .5, BGR_COLOR['white'])    \n",
    "    cv2.imshow('Press \"c\" to capture, \"q\" to quit', frame)\n",
    "\n",
    "# When the background image is captured, release the capture\n",
    "if board.is_light_off:\n",
    "    board.switch_light()  # turn light back on\n",
    "    time.sleep(0.1)\n",
    "board.exit()\n",
    "cap.release()\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "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": 4
}