{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import threading\n",
    "import time, os, json\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import multiprocessing as mp\n",
    "\n",
    "from scipy import stats\n",
    "from situtils import FPSTimes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionTracker(FPSTimes):\n",
    "\n",
    "    default_cfg = {\n",
    "        \"background_light\": \"background_light.png\",\n",
    "        \"background_dark\": \"background_dark.png\",\n",
    "        \"threshold_light\": 40,\n",
    "        \"threshold_dark\": 20,\n",
    "        \"arena_x\": 522,\n",
    "        \"arena_y\": 372,\n",
    "        \"arena_radius\": 330,\n",
    "        \"floor_radius\": 287,\n",
    "        \"max_fps\": 50,\n",
    "        \"file_path\": \"positions.csv\",\n",
    "        \"floor_r_in_meters\": 0.46,\n",
    "        \"angle_compensation\": -90,\n",
    "        \"flip_x\": True,\n",
    "        \"flip_y\": False\n",
    "    }\n",
    "    \n",
    "    def __init__(self, status, video_stream, cfg):\n",
    "        super(PositionTracker, self).__init__()\n",
    "        \n",
    "        self.status = status\n",
    "        self.cfg = cfg\n",
    "        self.video_stream = video_stream\n",
    "        self.bg_light = cv2.imread(cfg['background_light'], 1)\n",
    "        self.bg_dark  = cv2.imread(cfg['background_dark'], 1)\n",
    "        self.background = self.bg_light  # light by default\n",
    "        self.is_light = True\n",
    "        self._x, self._y = None, None\n",
    "        self.pixel_size = cfg['floor_r_in_meters'] / float(cfg['floor_radius'])\n",
    "        self.contour = []\n",
    "        self.lr = None  # linear regression of the contour\n",
    "        self.stopped = False\n",
    "\n",
    "        self.mask = np.zeros(shape=self.background.shape, dtype=\"uint8\")\n",
    "        cv2.circle(self.mask, (cfg['arena_x'], cfg['arena_y']), cfg['arena_radius'], (255,255,255), -1)\n",
    "\n",
    "        with open(cfg['file_path'], 'w') as f:\n",
    "            f.write(\"time,x,y\\n\")\n",
    "\n",
    "    def px_to_meters(self, x, y):\n",
    "        x_m = float(self.cfg['arena_x'] - x) * self.pixel_size * (-1 if self.cfg['flip_x'] else 1)\n",
    "        y_m = float(self.cfg['arena_y'] - y) * self.pixel_size * (-1 if self.cfg['flip_y'] else 1)\n",
    "        return x_m, y_m\n",
    "\n",
    "    def meters_to_px(self, x, y):\n",
    "        x_m = self.cfg['arena_x'] - (x / self.pixel_size) * (-1 if self.cfg['flip_x'] else 1)\n",
    "        y_m = self.cfg['arena_y'] - (y / self.pixel_size) * (-1 if self.cfg['flip_y'] else 1)\n",
    "        return int(x_m), int(y_m)\n",
    "\n",
    "    def correct_angle(self, phi):\n",
    "        return (2*np.pi - phi) + np.deg2rad(self.cfg['angle_compensation'])\n",
    "    \n",
    "    def start(self):\n",
    "        self._th = threading.Thread(target=self.update, args=())\n",
    "        self._th.start()\n",
    "    \n",
    "    def stop(self):\n",
    "        self.stopped = True\n",
    "        self._th.join()\n",
    "        print('Position tracker stopped')\n",
    "        \n",
    "    def update(self):\n",
    "        next_frame = time.time() + 1.0/self.cfg['max_fps']\n",
    "        \n",
    "        while not self.stopped:\n",
    "            frame = self.video_stream.read()\n",
    "            if frame is None:\n",
    "                time.sleep(0.05)\n",
    "                continue\n",
    "                \n",
    "            if time.time() < next_frame:\n",
    "                time.sleep(0.001)\n",
    "                continue\n",
    "                \n",
    "            self.count()  # count FPS\n",
    "            self.detect_position(frame)\n",
    "            next_frame += 1.0/self.cfg['max_fps']\n",
    "\n",
    "            if self.status.value == 2 and self._x is not None:\n",
    "                with open(self.cfg['file_path'], 'a') as f:\n",
    "                    f.write(\",\".join([str(x) for x in (self.frame_times[-1], \\\n",
    "                               round(self.x_in_m, 4), round(self.y_in_m, 4))]) + \"\\n\")         \n",
    "       \n",
    "    def switch_background(self):\n",
    "        self.background = self.bg_dark if self.is_light else self.bg_light \n",
    "        self.is_light = not self.is_light\n",
    "        \n",
    "    def detect_position(self, frame):\n",
    "        masked_frame = cv2.bitwise_and(src1=frame, src2=self.mask)\n",
    "\n",
    "        # Substracts background from current frame\n",
    "        if 'subtract' in self.cfg and self.cfg['subtract'] < 0:\n",
    "            subject = cv2.subtract(masked_frame, self.background)\n",
    "        else:\n",
    "            subject = cv2.subtract(self.background, masked_frame)\n",
    "        subject = cv2.absdiff(masked_frame, self.background)\n",
    "\n",
    "        # Converts subject to grey scale\n",
    "        subject_gray = cv2.cvtColor(subject, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "        # Applies blur and thresholding to the subject\n",
    "        kernel_size = (25,25)\n",
    "        frame_blur = cv2.GaussianBlur(subject_gray, kernel_size, 0)\n",
    "        th = self.cfg['threshold_light'] if self.is_light else self.cfg['threshold_dark']\n",
    "        _, thresh = cv2.threshold(frame_blur, th, 255, cv2.THRESH_BINARY)\n",
    "\n",
    "        # Finds contours and selects the contour with the largest area\n",
    "        contours, hierarchy = cv2.findContours(thresh.copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)\n",
    "\n",
    "        if len(contours) == 0:\n",
    "            return\n",
    "        contour = contours[np.argmax(list(map(cv2.contourArea, contours)))]\n",
    "        M = cv2.moments(contour)\n",
    "        if (M['m00'] == 0):\n",
    "            return\n",
    "\n",
    "        self._x = M['m10'] / M['m00']\n",
    "        self._y = M['m01'] / M['m00']\n",
    "        self.contour = contour  # in pixels\n",
    "        \n",
    "        ctr_in_px = np.array([x for x in zip(contour[:, 0, 0], contour[:, 0, 1])])\n",
    "        if ctr_in_px[:, 0].max() - ctr_in_px[:, 0].min() > ctr_in_px[:, 1].max() - ctr_in_px[:, 1].min():\n",
    "            slope, intercept, r_value, p_value, std_err = stats.linregress(ctr_in_px[:, 0], ctr_in_px[:, 1])\n",
    "            self.lr = {'slope': slope, 'intercept': intercept, 'r_value': r_value, 'p_value': p_value, 'std_err': std_err}\n",
    "        else:\n",
    "            slope, intercept, r_value, p_value, std_err = stats.linregress(ctr_in_px[:, 1], ctr_in_px[:, 0])\n",
    "            self.lr = {'slope': 1 - slope, 'intercept': intercept, 'r_value': r_value, 'p_value': p_value, 'std_err': std_err}\n",
    "            \n",
    "    @property\n",
    "    def x_in_px(self):\n",
    "        return int(self._x) if self._x is not None else None\n",
    "\n",
    "    @property\n",
    "    def y_in_px(self):\n",
    "        return int(self._y) if self._y is not None else None\n",
    "    \n",
    "    @property\n",
    "    def x_in_m(self):\n",
    "        if self._x is None:\n",
    "            return None\n",
    "        return (self.cfg['arena_x'] - self._x) * self.pixel_size * (-1 if self.cfg['flip_x'] else 1)\n",
    "    \n",
    "    @property\n",
    "    def y_in_m(self):\n",
    "        if self._y is None:\n",
    "            return None\n",
    "        return (self.cfg['arena_y'] - self._y) * self.pixel_size * (-1 if self.cfg['flip_y'] else 1)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing position detection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Webcam stream 960.0:720.0 at 20.00 FPS started\n",
      "Camera released\n",
      "Position tracker stopped\n"
     ]
    }
   ],
   "source": [
    "#import nbimporter\n",
    "#from camera import WebcamStream\n",
    "\n",
    "with open(os.path.join('..', 'profiles', 'prodesk-home-aSIT.json')) as json_file:\n",
    "    cfg = json.load(json_file)\n",
    "\n",
    "# controller status: 1 - detecting, 2 - detecting + logging\n",
    "status = mp.Value('i', 1)\n",
    "\n",
    "# let's use a webcam stream\n",
    "vs = WebcamStream(cfg['camera'])\n",
    "vs.start()  # stream runs in a separate thread\n",
    "\n",
    "# init controller\n",
    "pt_cfg = cfg['position']\n",
    "pt_cfg['background_light'] = os.path.join('..', 'assets', 'background_light.png')\n",
    "pt_cfg['background_dark'] = os.path.join('..', 'assets', 'background_dark.png')\n",
    "pt = PositionTracker(status, vs, pt_cfg)\n",
    "pt.start()\n",
    "kernel_size = (25,25)\n",
    "\n",
    "try:\n",
    "    while True:\n",
    "        frame = vs.read()\n",
    "        if frame is not None:\n",
    "            masked_frame = cv2.bitwise_and(src1=frame, src2=pt.mask)\n",
    "            #masked_frame = cv2.subtract(masked_frame, pt.background)\n",
    "            #masked_frame = cv2.cvtColor(masked_frame, cv2.COLOR_BGR2GRAY)\n",
    "            #masked_frame = cv2.GaussianBlur(masked_frame, kernel_size, 0)\n",
    "            \n",
    "            cv2.putText(masked_frame, 'Position: %.2f FPS' % pt.get_avg_fps(), \n",
    "                     (10, 60), cv2.FONT_HERSHEY_DUPLEX, .5, (255, 255, 255))        \n",
    "            \n",
    "            if pt.x_in_px is not None:\n",
    "                color = (127, 255, 0) if status.value == 1 else (0, 0, 255)\n",
    "                cv2.circle(masked_frame, (pt.x_in_px, pt.y_in_px), 2, color, -1)\n",
    "                cv2.drawContours(masked_frame, [pt.contour], 0, color, 1, cv2.LINE_AA)\n",
    "                cv2.putText(masked_frame, 'Animal: %.3f %.3f' % (pt.x_in_m, pt.y_in_m), (10, 40), \\\n",
    "                            cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))\n",
    "            \n",
    "                # regression line\n",
    "                #if pt.lr is not None:\n",
    "                x1, x2 = pt.contour[:, 0, 0].min() - 10, pt.contour[:, 0, 0].max() + 10\n",
    "                y1, y2 = int(x1 * pt.lr['slope'] + pt.lr['intercept']), int(x2 * pt.lr['slope'] + pt.lr['intercept'])\n",
    "                cv2.line(frame, (x1, y1), (x2, y2), color, 1)\n",
    "            \n",
    "            cv2.putText(masked_frame, 'L' if pt.is_light else 'D', (10, 20), cv2.FONT_HERSHEY_DUPLEX, .5, (255, 255, 255))            \n",
    "            cv2.imshow('Webcam', masked_frame)\n",
    "\n",
    "        k = cv2.waitKey(33)\n",
    "        if k == ord('q'):\n",
    "            break\n",
    "            \n",
    "        if k == ord('s'):\n",
    "            status.value = 2 if status.value == 1 else 1\n",
    "\n",
    "        if k == ord('d'):\n",
    "            pt.switch_background()\n",
    "            \n",
    "        if k == ord('r'):\n",
    "            cv2.imwrite('frame.jpg', masked_frame)\n",
    "            \n",
    "finally:\n",
    "    cv2.destroyAllWindows()\n",
    "    vs.stop(), pt.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image sandbox"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = cv2.imread('substract.jpg')\n",
    "#plt.imshow(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "color = ('b','g','r')\n",
    "for channel, col in enumerate(color):\n",
    "    histr = cv2.calcHist([image],[channel],None,[256],[0,256])\n",
    "    plt.plot(histr, color = col)\n",
    "    plt.xlim([0,256])\n",
    "    plt.ylim([0, 1000])"
   ]
  },
  {
   "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
}
