counter.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 탐지된 객체 종류별 그룹화"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from collections import defaultdict\n",
    "from mytool import centroid_tracking, tracking_id, corn2xywh\n",
    "\n",
    "def groupby_detection_to_dic(names, detection):\n",
    "    groupby_dic = defaultdict(lambda:{'ids':[], 'boxes':[], 'centroids':[], })\n",
    "    if len(detection) == 0: return groupby_dic\n",
    "    idxs = detection[:,-1]\n",
    "    boxes = detection[:, :4]\n",
    "    centroids = corn2xywh(boxes)[:,:2]\n",
    "    for idx, box, centroid in zip(idxs, boxes, centroids):\n",
    "        name = names[int(idx)]\n",
    "        groupby_dic[name]['boxes'].append(box)\n",
    "        groupby_dic[name]['centroids'].append(centroid)\n",
    "    \n",
    "    for key in groupby_dic:\n",
    "        groupby_dic[key]['boxes'] = np.stack(groupby_dic[key]['boxes']) if groupby_dic[key]['boxes'] else []\n",
    "        groupby_dic[key]['centroids'] = np.stack(groupby_dic[key]['centroids']) if groupby_dic[key]['centroids'] else []\n",
    "    \n",
    "    return groupby_dic\n",
    "\n",
    "def groupby_dic_updater(next_id_dic, groupby_dic_t0, names, detection): # detection정보로 next_id_dic 객체 카운트?\n",
    "    groupby_dic_t1 = groupby_detection_to_dic(names, detection) # {'car':[[], boxes, centroids], 'bus':[[], boxes, centroids]}\n",
    "    keys = set(groupby_dic_t0.keys()) | set(groupby_dic_t1.keys())\n",
    "    for key in keys:\n",
    "        centroids_t0 = groupby_dic_t0[key]['centroids']\n",
    "        centroids_t1 = groupby_dic_t1[key]['centroids']\n",
    "        pairs = centroid_tracking(centroids_t0, centroids_t1)\n",
    "\n",
    "        ids_t0 = groupby_dic_t0[key]['ids'][:]\n",
    "        len_boxes_t1 = len(groupby_dic_t1[key]['boxes'])\n",
    "        cnt_t0 = next_id_dic[key]\n",
    "        ids_t1, cnt_t1 = tracking_id(ids_t0, pairs, len_boxes_t1, cnt_t0)\n",
    "        next_id_dic[key] = cnt_t1\n",
    "        groupby_dic_t1[key]['ids'] = ids_t1\n",
    "\n",
    "    return groupby_dic_t1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 영역안 탐지된 객체 카운트"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_in_area(centroid, count_area):\n",
    "    x1, y1, x2, y2 = count_area\n",
    "    x, y = centroid\n",
    "    if x1 <= x <= x2 and y1 <= y <= y2: return True\n",
    "    else: return False\n",
    "\n",
    "def counter_dic_updater(counter_dic, groupby_dic, count_area):\n",
    "    for key in groupby_dic:\n",
    "        ids = groupby_dic[key]['ids']\n",
    "        centroids = groupby_dic[key]['centroids']\n",
    "        for id, centroid in zip(ids, centroids):\n",
    "            if is_in_area(centroid, count_area) and not id in counter_dic[key]['set']:\n",
    "                counter_dic[key]['set'].add(id)\n",
    "                counter_dic[key]['cnt'] += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 가장자리에 붙어있는 박스 삭제"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mytool import corn2xywh\n",
    "\n",
    "def delete_edge_boxes(img, detection):\n",
    "    height, width, _ = img.shape\n",
    "    new_detection = []\n",
    "    \n",
    "    for det in detection:\n",
    "        x, y, _, _ = corn2xywh(det[None, :4])[0]\n",
    "        if x < 20 or y < 20: continue\n",
    "        if width-20 < x or height-20 < y: continue\n",
    "        # x1, y1, x2, y2 = det[:4]\n",
    "        # if (x1 < 5 or y1 < 5) and (x < 50 or y < 50): continue\n",
    "        # if (width-5 < x2 or height-5 < y2) and (width-50 < x or height-50 < y): continue\n",
    "        # if x1 == 0 or y1 == 0: continue\n",
    "        # if width-1 == x2 or height-1 == y2: continue\n",
    "        new_detection.append(det)\n",
    "    return np.stack(new_detection) if new_detection else np.array([])\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 카운트 영역과 카운트 수 그리기"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "\n",
    "def draw_area(img, loc):\n",
    "    height, width, _ = img.shape\n",
    "    boundary1 = int(height*loc)\n",
    "    boundary2 = boundary1+30\n",
    "    box = [0, boundary1, width-1, boundary2] # xyxy\n",
    "    img = cv2.rectangle(img, box[:2], box[2:], (0,0,255), 3)\n",
    "    return img, box\n",
    "\n",
    "def draw_count(img, counter_dic):\n",
    "    height, width, _ = img.shape\n",
    "    loc = [0,30]\n",
    "    font = cv2.FONT_HERSHEY_SIMPLEX\n",
    "    for key in counter_dic:\n",
    "        cnt = counter_dic[key]['cnt']\n",
    "        cv2.putText(img, f\"{key} : {cnt}\", loc, font, 1, (255,0,0), 2, cv2.LINE_AA)\n",
    "        loc[1] += 30\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "from yolov5.detect import run\n",
    "from IPython.display import clear_output\n",
    "from collections import defaultdict\n",
    "import time\n",
    "from mytool import show_annotation\n",
    "\n",
    "SOURCE_PATH = \"./temp/traffic_30_crop.mp4\"\n",
    "# SOURCE_PATH = \"http://192.168.0.57:8090/video\"\n",
    "WEIGHTS_PATH = \"./yolov5m.pt\"\n",
    "\n",
    "# 프레임별 탐지 객체 박스 제네레이터\n",
    "gen = run(source=SOURCE_PATH, weights=WEIGHTS_PATH, view_img=False, nosave=True, \\\n",
    "    classes=[2,3,5,7], conf_thres=0.35, iou_thres=0.3, agnostic_nms=True)\n",
    "    # classes=[2,], conf_thres=0.5, iou_thres=0.3, agnostic_nms=True)\n",
    "names = next(gen)\n",
    "# print(\"names : \")\n",
    "# print(names)\n",
    "# print('#')\n",
    "\n",
    "groupby_dic = defaultdict(lambda:{'ids':[], 'boxes':[], 'centroids':[], })\n",
    "next_id_dic = defaultdict(int)\n",
    "counter_dic = defaultdict(lambda:{'set':set(), 'cnt':0})\n",
    "\n",
    "# 카운트 영상 저장\n",
    "vcap = cv2.VideoCapture(SOURCE_PATH)\n",
    "width  = int(vcap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "height = int(vcap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "vcap.release()\n",
    "fourcc = cv2.VideoWriter_fourcc(*'MP4V')\n",
    "out = cv2.VideoWriter('output.mp4', fourcc, 20.0, (width,height))\n",
    "\n",
    "\n",
    "# try:\n",
    "for img, detection in gen:\n",
    "    # print(detection) # xywh, conf, id\n",
    "    detection = detection.numpy()\n",
    "    detection = delete_edge_boxes(img, detection)\n",
    "    groupby_dic = groupby_dic_updater(next_id_dic, groupby_dic, names, detection)\n",
    "    img = show_annotation(img, groupby_dic)\n",
    "    img, count_area = draw_area(img, 1/2)\n",
    "    counter_dic_updater(counter_dic, groupby_dic, count_area)\n",
    "    draw_count(img, counter_dic)\n",
    "    # clear_output()#wait=True)\n",
    "    out.write(img)\n",
    "    cv2.imshow('test', img)\n",
    "    cv2.waitKey(1)\n",
    "    # print(counter_dic)\n",
    "    # print(time.time())\n",
    "# except:\n",
    "#     pass\n",
    "\n",
    "out.release()\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cv2.imshow('test', img)\n",
    "cv2.waitKey(0)\n",
    "cv2.destroyAllWindows()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "57f3a4084d9bbea5e2496975c38f9f0c4e2cb16976f2cda4d07e2a8c664c7c73"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}