|
|
@@ -0,0 +1,125 @@
|
|
|
+{
|
|
|
+ "cells": [
|
|
|
+ {
|
|
|
+ "cell_type": "code",
|
|
|
+ "execution_count": 9,
|
|
|
+ "id": "cad57710",
|
|
|
+ "metadata": {},
|
|
|
+ "outputs": [],
|
|
|
+ "source": [
|
|
|
+ "import cv2\n",
|
|
|
+ "import numpy as np"
|
|
|
+ ]
|
|
|
+ },
|
|
|
+ {
|
|
|
+ "cell_type": "code",
|
|
|
+ "execution_count": null,
|
|
|
+ "id": "91d021d9",
|
|
|
+ "metadata": {},
|
|
|
+ "outputs": [],
|
|
|
+ "source": [
|
|
|
+ "class SITDisplay:\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), 'pink': (255, 0, 255), 'tortoise': (0, 255, 255)\n",
|
|
|
+ " }\n",
|
|
|
+ "\n",
|
|
|
+ " island_colors = {2: 'green', 3: 'blue', 4: 'yellow', 5: 'black'}\n",
|
|
|
+ " animal_colors = {1: 'tortoise', 2: 'pink', 3: 'white'}\n",
|
|
|
+ "\n",
|
|
|
+ " def __init__(self, pt): # injected dependency - position controller\n",
|
|
|
+ " self.pt = pt\n",
|
|
|
+ " self.masked_frame = None\n",
|
|
|
+ " \n",
|
|
|
+ " def render(self, input_frame, status, islands=[], angles=[], text_infos=[]):\n",
|
|
|
+ " COLORS = SITDisplay.COLORS\n",
|
|
|
+ " FONT = cv2.FONT_HERSHEY_DUPLEX\n",
|
|
|
+ " status_color = COLORS['green'] if status.value == 1 else COLORS['red']\n",
|
|
|
+ " texts = list(text_infos)\n",
|
|
|
+ " \n",
|
|
|
+ " # mask unused areas\n",
|
|
|
+ " frame = cv2.bitwise_and(src1=input_frame, src2=self.pt.mask)\n",
|
|
|
+ " self.masked_frame = frame.copy()\n",
|
|
|
+ "\n",
|
|
|
+ " # size of the arena and status indicator\n",
|
|
|
+ " cv2.circle(frame, (self.pt.cfg['arena_x'], self.pt.cfg['arena_y']), self.pt.cfg['arena_radius'], COLORS['red'], 2)\n",
|
|
|
+ " cv2.circle(frame, (self.pt.cfg['arena_x'], self.pt.cfg['arena_y']), self.pt.cfg['floor_radius'], COLORS['red'], 2)\n",
|
|
|
+ " cv2.circle(frame, (20,20), 10, status_color, -6)\n",
|
|
|
+ " \n",
|
|
|
+ " # draw islands\n",
|
|
|
+ " if len(islands) > 0:\n",
|
|
|
+ " for island in islands:\n",
|
|
|
+ " clr = COLORS[SITDisplay.island_colors[island.sound_id]]\n",
|
|
|
+ " x, y = self.pt.meters_to_px(island.x, island.y)\n",
|
|
|
+ " cv2.circle(frame, (x, y), int(island.r/self.pt.pixel_size), clr, 2)\n",
|
|
|
+ " \n",
|
|
|
+ " # draw anglular indicators\n",
|
|
|
+ " if len(angles) > 0: # should not be more than 4\n",
|
|
|
+ " p_rho = self.pt.cfg['floor_radius'] * self.pt.pixel_size\n",
|
|
|
+ " for p_phi, clr in zip(angles, ['blue', 'yellow', 'pink', 'tortoise'][:len(angles)]):\n",
|
|
|
+ " x_init, y_init = self.pt.meters_to_px(p_rho * np.sin(p_phi), p_rho * np.cos(p_phi))\n",
|
|
|
+ " cv2.circle(frame, (x_init, y_init), 10, COLORS[clr], -1)\n",
|
|
|
+ "\n",
|
|
|
+ " # light or dark period\n",
|
|
|
+ " texts.append('Light' if self.pt.is_light else 'Dark')\n",
|
|
|
+ "\n",
|
|
|
+ " # draw animal position center(s) and contours\n",
|
|
|
+ " if self.pt.positions_in_px is not None:\n",
|
|
|
+ " pos_in_px = self.pt.positions_in_px\n",
|
|
|
+ " pos_in_m = self.pt.positions_in_m\n",
|
|
|
+ " contours = self.pt.contours\n",
|
|
|
+ " \n",
|
|
|
+ " for i in range(len(pos_in_px)):\n",
|
|
|
+ " clr = COLORS[SITDisplay.animal_colors[i+1]]\n",
|
|
|
+ " x_px, y_px, alpha = pos_in_px[i][0], pos_in_px[i][1], self.pt.hds[0]\n",
|
|
|
+ " \n",
|
|
|
+ " # draw animal center and contours\n",
|
|
|
+ " cv2.circle(frame, (x_px, y_px), 5, clr, -1)\n",
|
|
|
+ " cv2.drawContours(frame, [contours[i]], 0, clr, 1, cv2.LINE_AA)\n",
|
|
|
+ " \n",
|
|
|
+ " # draw instant position values and HD vector\n",
|
|
|
+ " cv2.putText(frame, '%.2f %.2f' % (pos_in_m[i][0], pos_in_m[i][1]), (x_px + 10, y_px + 10), FONT, .3, clr)\n",
|
|
|
+ " x1, y1 = int(x_px + 40*np.cos(np.deg2rad(alpha))), int(y_px - 40*np.sin(np.deg2rad(alpha)))\n",
|
|
|
+ " cv2.line(frame, (x_px, y_px), (x1, y1), clr, 1)\n",
|
|
|
+ " \n",
|
|
|
+ " texts.append('Animal %d: %.3f %.3f' % (i+1, pos_in_m[i][0], pos_in_m[i][1]))\n",
|
|
|
+ " texts.append('Speed/HD: %.2f %.2f' % (self.pt.speeds[0], self.pt.hds[0]))\n",
|
|
|
+ " \n",
|
|
|
+ " # draw text infos\n",
|
|
|
+ " for i, text in enumerate(texts):\n",
|
|
|
+ " cv2.putText(frame, text, (10, 30 + 20*(i+1)), FONT, .5, COLORS['white'])\n",
|
|
|
+ "\n",
|
|
|
+ " return frame"
|
|
|
+ ]
|
|
|
+ },
|
|
|
+ {
|
|
|
+ "cell_type": "code",
|
|
|
+ "execution_count": null,
|
|
|
+ "id": "5e7f5cbd",
|
|
|
+ "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": 5
|
|
|
+}
|
