{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import keras\n",
    "import numpy as np\n",
    "import threading\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "from ultralytics import YOLO\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def object_movie_detection(movie_name, background_substract_method, substractor_threshold, mask_parameters, movie_resize, detection_area,object_area_to_detect):\n",
    "    if background_substract_method == 'KNN':\n",
    "        object_detector = cv2.createBackgroundSubtractorKNN(history=500, dist2Threshold=substractor_threshold)\n",
    "    else :\n",
    "        object_detector = cv2.createBackgroundSubtractorMOG2(history=500, varThreshold=substractor_threshold)\n",
    "    \n",
    "    cap = cv2.VideoCapture(movie_name)\n",
    "    #cap.set(cv2.CAP_PROP_FPS, 60)\n",
    "    frames = []\n",
    "    pbar = tqdm(total = cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "    while True:\n",
    "        _, frame = cap.read()\n",
    "        if(frame is not None):\n",
    "            pbar.update(1)\n",
    "            detection_field = frame[detection_area[0][0]: detection_area[0][1],detection_area[1][0]:detection_area[1][1]]\n",
    "            mask = object_detector.apply(detection_field)\n",
    "            _, mask = cv2.threshold(mask,mask_parameters[0],mask_parameters[1], cv2.THRESH_BINARY)\n",
    "            conturs, _ =cv2.findContours(mask, cv2.RETR_LIST , cv2.CHAIN_APPROX_TC89_L1)\n",
    "\n",
    "            images = []\n",
    "            counter = 0\n",
    "            for cnt in conturs:\n",
    "                area = cv2.contourArea(cnt)\n",
    "                if area > object_area_to_detect:\n",
    "                    counter += 1\n",
    "                    x,y,w,h = cv2.boundingRect(cnt)\n",
    "                    rectangle = cv2.rectangle(detection_field,(x,y),(x+w,y+h),(0,255,0),3)\n",
    "            rectangle = cv2.rectangle(frame,(330,100),(1900,900),(255,0,0),3)\n",
    "            cv2.putText(rectangle,f'DETECTION AREA    DETECTED OBJECTS:  {counter}',(330,80),cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,0,255), 4)\n",
    "            frames.append(frame)\n",
    "            frame = cv2.resize(frame, (movie_resize[0], movie_resize[1])) \n",
    "            #cv2.imshow(\"roi\", frame)\n",
    "\n",
    "            key = cv2.waitKey(30)\n",
    "            if key == 27:\n",
    "                cv2.destroyAllWindows()\n",
    "                pbar.close()\n",
    "                return frames\n",
    "        else:\n",
    "            cv2.destroyAllWindows()\n",
    "            pbar.close()\n",
    "            return frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_video(video_name, frames, frame_sizes,fps):\n",
    "    out = cv2.VideoWriter(f'{video_name}.avi', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps, (frame_sizes[1], frame_sizes[0]))\n",
    "    for frame in frames:\n",
    "        out.write(frame)\n",
    "    out.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 660/660.0 [01:12<00:00,  9.14it/s]\n"
     ]
    }
   ],
   "source": [
    "movie_detection_1 = object_movie_detection(\n",
    "    '../examples_movie/rybki.mp4',\n",
    "    background_substract_method='KNN',\n",
    "    substractor_threshold=100,\n",
    "    mask_parameters=[254,255],\n",
    "    movie_resize=[960,540],\n",
    "    detection_area=[(100,900),(330,1900)],\n",
    "    object_area_to_detect=1000\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_video('saved_movies/example_detection_1',movie_detection_1,movie_detection_1[0].shape,60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1037/1037.0 [03:21<00:00,  5.15it/s]\n"
     ]
    }
   ],
   "source": [
    "movie_detection_2 = object_movie_detection(\n",
    "    '../examples_movie/rybki2.mp4',\n",
    "    background_substract_method='KNN',\n",
    "    substractor_threshold=100,\n",
    "    mask_parameters=[254,255],\n",
    "    movie_resize=[960,540],\n",
    "    detection_area=[(100,900),(330,1900)],\n",
    "    object_area_to_detect=1000\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_video('saved_movies/example_detection_2',movie_detection_2,movie_detection_2[0].shape,60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████▉| 962/963.0 [01:20<00:00, 11.98it/s]\n"
     ]
    }
   ],
   "source": [
    "movie_detection_3 = object_movie_detection(\n",
    "    '../examples_movie/rybki3.mp4',\n",
    "    background_substract_method='KNN',\n",
    "    substractor_threshold=100,\n",
    "    mask_parameters=[254,255],\n",
    "    movie_resize=[960,540],\n",
    "    detection_area=[(100,900),(330,1900)],\n",
    "    object_area_to_detect=1000\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_video('saved_movies/example_detection_3',movie_detection_3,movie_detection_3[0].shape,60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████▉| 857/858.0 [01:55<00:00,  7.43it/s]\n"
     ]
    }
   ],
   "source": [
    "movie_detection_4 = object_movie_detection(\n",
    "    '../examples_movie/rybki4.mp4',\n",
    "    background_substract_method='KNN',\n",
    "    substractor_threshold=100,\n",
    "    mask_parameters=[254,255],\n",
    "    movie_resize=[960,540],\n",
    "    detection_area=[(100,900),(330,1900)],\n",
    "    object_area_to_detect=1000\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_video('saved_movies/example_detection_4',movie_detection_4,movie_detection_4[0].shape,60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████▉| 779/780.0 [01:01<00:00, 12.64it/s]\n"
     ]
    }
   ],
   "source": [
    "movie_detection_5 = object_movie_detection(\n",
    "    '../examples_movie/rybki5.mp4',\n",
    "    background_substract_method='KNN',\n",
    "    substractor_threshold=100,\n",
    "    mask_parameters=[254,255],\n",
    "    movie_resize=[960,540],\n",
    "    detection_area=[(100,900),(330,1900)],\n",
    "    object_area_to_detect=1000\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_video('saved_movies/example_detection_5',movie_detection_5,movie_detection_5[0].shape,60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 388/388.0 [00:26<00:00, 14.55it/s]\n"
     ]
    }
   ],
   "source": [
    "movie_detection_6 = object_movie_detection(\n",
    "    '../examples_movie/jelly2.mp4',\n",
    "    background_substract_method='KNN',\n",
    "    substractor_threshold=100,\n",
    "    mask_parameters=[254,255],\n",
    "    movie_resize=[960,540],\n",
    "    detection_area=[(100,900),(330,1900)],\n",
    "    object_area_to_detect=1000\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_video('saved_movies/example_detection_6',movie_detection_6,movie_detection_6[0].shape,60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 388/388.0 [00:37<00:00, 10.38it/s]\n"
     ]
    }
   ],
   "source": [
    "movie_detection_7 = object_movie_detection(\n",
    "    '../examples_movie/jelly3.mp4',\n",
    "    background_substract_method='KNN',\n",
    "    substractor_threshold=100,\n",
    "    mask_parameters=[254,255],\n",
    "    movie_resize=[960,540],\n",
    "    detection_area=[(100,900),(330,1900)],\n",
    "    object_area_to_detect=1000\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_video('saved_movies/example_detection_7',movie_detection_7,movie_detection_7[0].shape,60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames = object_movie_detection(\n",
    "    '../examples_movie/rybki.mp4',\n",
    "    background_substract_method='KNN',\n",
    "    substractor_threshold=100,\n",
    "    mask_parameters=[254,255],\n",
    "    movie_resize=[960,540],\n",
    "    detection_area=[(100,900),(330,1900)],\n",
    "    object_area_to_detect=1000\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def object_movie_yolo_detection(movie_name,model,detection_time):\n",
    "\n",
    "    cap = cv2.VideoCapture(movie_name)\n",
    "    cap.set(cv2.CAP_PROP_FPS, 60)\n",
    "\n",
    "    frames = []\n",
    "    pbar = tqdm(total = detection_time * 60)\n",
    "    i = 0\n",
    "\n",
    "    while True:\n",
    "        i += 1\n",
    "        ret, frame = cap.read()\n",
    "        if(frame is not None):\n",
    "            pbar.update(1)\n",
    "            model.predict(frame,save_crop=True)\n",
    "            frames.append(frame)\n",
    "            key = cv2.waitKey(30)\n",
    "            if key == 27:\n",
    "                cv2.destroyAllWindows()\n",
    "                return frames\n",
    "            if i >= detection_time * 60:\n",
    "                pbar.close\n",
    "                cv2.destroyAllWindows()\n",
    "                return frames\n",
    "        else:\n",
    "            cv2.destroyAllWindows()\n",
    "            pbar.close\n",
    "            return frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = YOLO(\"../yolo_model/best.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jelly2 = object_movie_yolo_detection('../examples_movie/jelly2.mp4',model)\n",
    "jelly3 = object_movie_yolo_detection('../examples_movie/jelly3.mp4',model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_video('saved_movies/yolo_jelly1', jelly2,jelly2[0].shape,60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_video('saved_movies/yolo_jelly2', jelly3,jelly3[0].shape,60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yolo_detect_1 = object_movie_yolo_detection('../examples_movie/rybki3.mp4',model,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_video('saved_movies/example_detection_8', yolo_detect_1,yolo_detect_1[0].shape,60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def object_detection_with_neural_pred(\n",
    "        movie_name, \n",
    "        background_substract_method, \n",
    "        substractor_threshold, \n",
    "        mask_parameters, \n",
    "        movie_resize, \n",
    "        detection_area,\n",
    "        object_area_to_detect,\n",
    "        model,\n",
    "        class_names,\n",
    "        time_to_capture):\n",
    "    if background_substract_method == 'KNN':\n",
    "        object_detector = cv2.createBackgroundSubtractorKNN(history=500, dist2Threshold=substractor_threshold)\n",
    "    else :\n",
    "        object_detector = cv2.createBackgroundSubtractorMOG2(history=500, varThreshold=substractor_threshold)\n",
    "    \n",
    "    cap = cv2.VideoCapture(movie_name)\n",
    "    cap.set(cv2.CAP_PROP_FPS, 60)\n",
    "\n",
    "    pbar = tqdm(total = time_to_capture * 60)\n",
    "    i = 0\n",
    "    frames = []\n",
    "    while True:\n",
    "        ret, frame = cap.read()\n",
    "        if(frame is not None):\n",
    "            pbar.update(1)\n",
    "            i += 1\n",
    "            detection_field = frame[detection_area[0][0]: detection_area[0][1],detection_area[1][0]:detection_area[1][1]]\n",
    "            mask = object_detector.apply(detection_field)\n",
    "            _, mask = cv2.threshold(mask,mask_parameters[0],mask_parameters[1], cv2.THRESH_BINARY)\n",
    "            conturs, _ =cv2.findContours(mask, cv2.RETR_LIST , cv2.CHAIN_APPROX_TC89_L1)\n",
    "\n",
    "            images = []\n",
    "            counter = 0\n",
    "            for cnt in conturs:\n",
    "                area = cv2.contourArea(cnt)\n",
    "                if area > object_area_to_detect:\n",
    "                    counter += 1\n",
    "                    x,y,w,h = cv2.boundingRect(cnt)\n",
    "                    rectangle = cv2.rectangle(detection_field,(x,y),(x+w,y+h),(0,255,0),3)\n",
    "                    image_to_predict = detection_field[y:y+h,x:x+w]\n",
    "                    image_to_predict = cv2.resize(image_to_predict,(224,224))\n",
    "                    images.append((x,y,rectangle,image_to_predict))\n",
    "            if images:\n",
    "                for i,image in enumerate(images):\n",
    "                    pred = model(x=tf.convert_to_tensor(image[3][None, :], dtype='float32'))\n",
    "                    labels = [class_names[np.argmax(pre)] for pre in pred]\n",
    "                    cv2.putText(image[2], labels[0], (image[0], image[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 1)\n",
    "            rectangle = cv2.rectangle(frame,(detection_area[1][0],detection_area[0][0]),(detection_area[1][1],detection_area[0][1]),(255,0,0),3)\n",
    "            cv2.putText(rectangle,f'DETECTION AREA    DETECTED OBJECTS:  {counter}',(330,80),cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,0,255), 4)\n",
    "            frames.append(frame)\n",
    "            if i >= time_to_capture * 60:\n",
    "                cv2.destroyAllWindows()\n",
    "                pbar.close()\n",
    "                return frames\n",
    "            #frame = cv2.resize(frame, (movie_resize[0], movie_resize[1])) \n",
    "            #cv2.imshow(\"roi\", frame)\n",
    "\n",
    "            key = cv2.waitKey(30)\n",
    "            if key == 27:\n",
    "                cv2.destroyAllWindows()\n",
    "                pbar.close()\n",
    "                return frames\n",
    "        else:\n",
    "            cv2.destroyAllWindows()\n",
    "            pbar.close()\n",
    "            return frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def wrap_frozen_graph(graph_def, inputs, outputs, print_graph=False):\n",
    "    def _imports_graph_def():\n",
    "        tf.compat.v1.import_graph_def(graph_def, name=\"\")\n",
    "\n",
    "    wrapped_import = tf.compat.v1.wrap_function(_imports_graph_def, [])\n",
    "    import_graph = wrapped_import.graph\n",
    "\n",
    "    if print_graph == True:\n",
    "        print(\"-\" * 50)\n",
    "        print(\"Frozen model layers: \")\n",
    "        layers = [op.name for op in import_graph.get_operations()]\n",
    "        for layer in layers:\n",
    "            print(layer)\n",
    "        print(\"-\" * 50)\n",
    "\n",
    "    return wrapped_import.prune(\n",
    "        tf.nest.map_structure(import_graph.as_graph_element, inputs),\n",
    "        tf.nest.map_structure(import_graph.as_graph_element, outputs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.io.gfile.GFile(\"../frozen_models/frozen_graph_best_vgg.pb\", \"rb\") as f:\n",
    "    graph_def = tf.compat.v1.GraphDef()\n",
    "    loaded = graph_def.ParseFromString(f.read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "frozen_func = wrap_frozen_graph(graph_def=graph_def,\n",
    "                                inputs=[\"x:0\"],\n",
    "                                outputs=[\"Identity:0\"],\n",
    "                                print_graph=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class_names=sorted(['Fish', \"Jellyfish\", 'Lionfish', 'Shark', 'Stingray', 'Turtle'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1037it [27:07,  1.57s/it]                        \n"
     ]
    }
   ],
   "source": [
    "frames = object_detection_with_neural_pred(\n",
    "    '../examples_movie/rybki2.mp4',\n",
    "    background_substract_method='KNN',\n",
    "    substractor_threshold=100,\n",
    "    mask_parameters=[254,255],\n",
    "    movie_resize=[960,540],\n",
    "    detection_area=[(500,900),(800,1900)],\n",
    "    object_area_to_detect=1000,\n",
    "    model=frozen_func,\n",
    "    class_names=class_names,\n",
    "    time_to_capture = 5\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_video('saved_movies/example_detection_9', frames,frames[0].shape,60)"
   ]
  }
 ],
 "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.9.2 (tags/v3.9.2:1a79785, Feb 19 2021, 13:44:55) [MSC v.1928 64 bit (AMD64)]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "393784674bcf6e74f2fc9b1b5fb3713f9bd5fc6f8172c594e5cfa8e8c12849bc"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}