run_colors.ipynb

{
 "metadata": {
  "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
  },
  "orig_nbformat": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 说明\n",
    "\n",
    "在开发环境中将目录文件都进行 Sync Obs 同步\n",
    "\n",
    "> model_data/train_mask_rcnn.h5 模型通过训练作业得到，也可以用官方权重\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 查看同步环境文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pwd\n",
    "\n",
    "!ls -lh\n",
    "\n",
    "!ls cailiao -lh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 开始\n",
    "\n",
    "导入一些默认包和初始项目路径\n",
    "\n",
    "**参数配置**\n",
    "\n",
    "- video_file 视频文件\n",
    "- min_score 最小显示检测分\n",
    "- input_size 统一输入图像大小\n",
    "- model_file 模型文件\n",
    "- model_feature 跟踪特征文件\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 参数配置\n",
    "video_file ='test.mp4'\n",
    "min_score =0.3\n",
    "input_size =1024\n",
    "model_file = 'model_data/train_mask_rcnn.h5'\n",
    "# model_feature = 'model_data/market1501.pb'\n",
    "model_feature = 'model_data/mars-small128.pb'\n",
    "\n",
    "box_size = 2        # 边框大小\n",
    "font_scale = 0.4    # 字体比例大小\n",
    "\n",
    "import time\n",
    "import json\n",
    "import sys, os\n",
    "import cv2\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from moxing.framework import file\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# 执行所在路径\n",
    "print(os.getcwd())\n",
    "sys.path.append(os.getcwd())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 导入所需模型依赖和配置\n",
    "\n",
    "设置存储的位置和临时变量\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# deep-sort跟踪\n",
    "from deep_sort import preprocessing\n",
    "from deep_sort import nn_matching\n",
    "from deep_sort.detection import Detection\n",
    "from deep_sort.tracker import Tracker\n",
    "from deep_sort import generate_detections\n",
    "from collections import deque\n",
    "\n",
    "# mask-rcnn模型 - 根据识别数量随机类型颜色\n",
    "from mrcnn.mrcnn_colors import MRCNN, isInSide\n",
    "\n",
    "# obs桶路径\n",
    "obs_path = \"obs://puddings/deep-sort-mask-rcnn/cailiao\"\n",
    "\n",
    "# 输出目录\n",
    "out_path = \"cailiao\"\n",
    "\n",
    "# 输出目录存在需要删除里边的内容\n",
    "if os.path.exists(out_path):\n",
    "    file.remove(out_path, recursive=True)\n",
    "os.makedirs(out_path)\n",
    "\n",
    "# 运动轨迹\n",
    "pts = [deque(maxlen=30) for _ in range(9999)]\n",
    "\n",
    "# 跟踪统计\n",
    "track_total = []\n",
    "\n",
    "# 跟踪类型总数量\n",
    "total_count = {}\n",
    "\n",
    "# 帧数，用于通过帧数取图\n",
    "frameNum = 0\n",
    "\n",
    "# Deep SORT 跟踪器\n",
    "encoder = generate_detections.create_box_encoder(model_feature, batch_size=1)\n",
    "metric = nn_matching.NearestNeighborDistanceMetric(\"cosine\", min_score, None)\n",
    "tracker = Tracker(metric)\n",
    "\n",
    "# 载入模型\n",
    "mrcnn = MRCNN(model_file, input_size, min_score)\n",
    "\n",
    "# 读取视频\n",
    "video = cv2.VideoCapture(video_file)\n",
    "\n",
    "# 输出保存视频\n",
    "fourcc = cv2.VideoWriter_fourcc(*'XVID')\n",
    "fps = video.get(cv2.CAP_PROP_FPS)\n",
    "size = (int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)))\n",
    "video_out = cv2.VideoWriter(out_path + \"/outputVideo.mp4\", fourcc, fps, size)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 进行视频帧识别\n",
    "\n",
    "对模型识别得到的类别特征进行目标跟踪统计\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 视频是否可以打开，进行逐帧识别绘制\n",
    "while video.isOpened:\n",
    "    # 视频读取图片帧\n",
    "    retval, frame = video.read()\n",
    "    if retval:\n",
    "        frame_orig = frame.copy()\n",
    "    else:\n",
    "        print(\"没有图像！尝试使用其他视频\")\n",
    "        break\n",
    "\n",
    "    prev_time = time.time()\n",
    "\n",
    "    # 识别结果\n",
    "    boxes, scores, classes, masks, colors = mrcnn.detect_result(frame, min_score)\n",
    "\n",
    "    # 特征提取和检测对象列表\n",
    "    features = encoder(frame, boxes)\n",
    "    detections = []\n",
    "    for bbox, score, classe, mask, color, feature in zip(boxes, scores, classes, masks, colors, features):\n",
    "        detections.append(Detection(bbox, score, classe, mask, color, feature))\n",
    "\n",
    "    # 运行非最大值抑制\n",
    "    boxes = np.array([d.tlwh for d in detections])\n",
    "    scores = np.array([d.score for d in detections])\n",
    "    indices = preprocessing.non_max_suppression(boxes, 1.0, scores)\n",
    "    detections = [detections[i] for i in indices]\n",
    "\n",
    "    # 遍历绘制检测对象信息\n",
    "    detect_count = {}\n",
    "    detect_temp = []\n",
    "    for det in detections:\n",
    "        y1, x1, y2, x2 = np.array(det.to_tlbr(), dtype=np.int32)\n",
    "        caption = '{} {:.2f}'.format(det.classe, det.score) if det.classe else det.score\n",
    "        \n",
    "        frame = mrcnn.apply_mask(frame, det.mask, det.color, 0.3)         # 类别掩膜颜色透明度\n",
    "        for c in range(3): det.color += (int(det.color[c]*255),)          # 颜色\n",
    "        cv2.rectangle(frame, (y1, x1), (y2, x2), det.color[3:], box_size) # 绘制类别边框\n",
    "\n",
    "        # 中心点\n",
    "        point = (int((y1+y2)/2),int((x1+x2)/2))\n",
    "        # cv2.circle(frame, point, 1, det.color[3:], box_size)\n",
    "        \n",
    "        # 类别文字显示\n",
    "        cv2.putText(\n",
    "            frame,\n",
    "            caption,\n",
    "            (y1, x1 - 5),\n",
    "            cv2.FONT_HERSHEY_SIMPLEX,\n",
    "            font_scale, det.color[3:],\n",
    "            box_size//2,\n",
    "            lineType=cv2.LINE_AA\n",
    "        )\n",
    "        # 统计物体数\n",
    "        if det.classe not in detect_count: detect_count[det.classe] = 0\n",
    "        detect_count[det.classe] += 1\n",
    "        detect_temp.append([det.classe, det.color[3:], point])\n",
    "        \n",
    "    # 追踪器刷新\n",
    "    tracker.predict()\n",
    "    tracker.update(detections)\n",
    "\n",
    "    # 遍历绘制跟踪信息\n",
    "    track_count = 0\n",
    "    for track in tracker.tracks:\n",
    "        if not track.is_confirmed() or track.time_since_update > 1: continue\n",
    "        y1, x1, y2, x2 = np.array(track.to_tlbr(), dtype=np.int32)\n",
    "        # cv2.rectangle(frame, (y1, x1), (y2, x2), (255, 255, 255), box_size//4)\n",
    "\n",
    "        # 跟踪统计数量\n",
    "        track_total.append(track.track_id)\n",
    "        track_count += 1\n",
    "\n",
    "        # 运动点轨迹\n",
    "        point = (int((y1+y2)/2),int((x1+x2)/2))\n",
    "        # cv2.circle(frame, point, 1, (255, 255, 255), box_size)\n",
    "        pts[track.track_id].append(point)\n",
    "        # 在识别类中标记跟踪 [ classe, color , point ]\n",
    "        for d in range(len(detect_temp)):\n",
    "            # 非标记目标跳过\n",
    "            if not isInSide(detect_temp[d][2], track.to_tlbr()): continue\n",
    "            \n",
    "            # 总统计数量\n",
    "            if detect_temp[d][0] not in total_count: total_count[detect_temp[d][0]] = [0, []]\n",
    "            if track.track_id not in total_count[detect_temp[d][0]][1]:\n",
    "                total_count[detect_temp[d][0]][0] += 1\n",
    "                total_count[detect_temp[d][0]][1].append(track.track_id)\n",
    "                # 输出小图目录,不存目录需要创建\n",
    "                label_path = os.path.join(out_path, \"{0}/{1}\".format('imageSeg', detect_temp[d][0]))\n",
    "                if not os.path.exists(label_path): os.makedirs(label_path)\n",
    "                cv2.imwrite(\"{0}/{1}.jpg\".format(label_path, track.track_id), frame_orig[x1:x2, y1:y2])\n",
    "                \n",
    "            # 跟踪标记号码\n",
    "            cv2.putText(\n",
    "                frame, \n",
    "                \"No. \" + str(track.track_id),\n",
    "                (y1, x1 - 15),\n",
    "                cv2.FONT_HERSHEY_SIMPLEX,\n",
    "                font_scale, (255, 255, 255),\n",
    "                box_size//2,\n",
    "                lineType=cv2.LINE_AA\n",
    "            )\n",
    "            \n",
    "            # 绘制运动路径\n",
    "            for j in range(1, len(pts[track.track_id])):\n",
    "                if pts[track.track_id][j - 1] is None or pts[track.track_id][j] is None: continue\n",
    "                thickness = int(np.sqrt(64 / float(j + 1)) * 2)\n",
    "                cv2.line(frame, (pts[track.track_id][j-1]), (pts[track.track_id][j]), detect_temp[d][1], thickness)\n",
    "\n",
    "    # 跟踪统计\n",
    "    trackTotalStr = 'Track Total: %s' % str(len(set(track_total)))\n",
    "    cv2.putText(frame, trackTotalStr, (20,20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (244, 67, 54), 1, cv2.LINE_AA)\n",
    "\n",
    "    # 跟踪数量\n",
    "    trackCountStr = 'Track Count: %s' % str(track_count)\n",
    "    cv2.putText(frame, trackCountStr, (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 193, 7), 1, cv2.LINE_AA)\n",
    "\n",
    "    # 识别类数统计\n",
    "    totalStr = \"\"\n",
    "    for k in detect_count.keys(): totalStr += '%s: %d    ' % (k, detect_count[k])\n",
    "    cv2.putText(frame, totalStr, (20, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (50, 0, 255), 1, cv2.LINE_AA)\n",
    "    \n",
    "    for i, label in enumerate(total_count):\n",
    "        labelTotal = '%s: %d ' % (label, total_count[label][0])\n",
    "        cv2.putText(frame, labelTotal, (20, 80 + 20 * i), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 87, 34), 1, cv2.LINE_AA)\n",
    "    \n",
    "    \n",
    "    # 绘制时间\n",
    "    curr_time = time.time()\n",
    "    exec_time = curr_time - prev_time\n",
    "    print(\"识别帧：{:.0f}/{:.0f} , 识别耗时: {:.2f} ms\".format(frameNum, video.get(7), 1000*exec_time))\n",
    "    \n",
    "    frameNum += 1\n",
    "    # 视频输出逐帧保存\n",
    "    video_out.write(frame)\n",
    "    # 绘制结果ipynb显示\n",
    "    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "    plt.figure(figsize=(10,10))\n",
    "    plt.imshow(frame)\n",
    "    plt.axis('on')\n",
    "    plt.show()\n",
    "\n",
    "# 任务完成后释放所有内容\n",
    "video.release()\n",
    "video_out.release()\n",
    "\n",
    "# 打开文件统计后遍历物体结果数据\n",
    "totalFile = open(out_path + \"/totalCount.txt\",\"w\")\n",
    "# 统计数量写入文件txt\n",
    "for label in total_count.keys():\n",
    "    labelTotal = \"{0}：{1} \\n\".format(label, total_count[label][0])\n",
    "    totalFile.write(labelTotal)\n",
    "# 关闭文件统计        \n",
    "totalFile.close()\n",
    "\n",
    "# 统计写入文件josn\n",
    "with open(out_path + \"/totalCount.json\", 'w') as tc:\n",
    "    json.dump(total_count, tc)\n",
    "\n",
    "# 复制保存到桶\n",
    "file.copy_parallel(out_path, obs_path)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 测试py文件\n",
    "\n",
    "py文件可以在 `Terminal` 中使用命令 `source /home/ma-user/anaconda3/bin/activate TensorFlow-1.13.1` 可以切换到 `TensorFlow-1.13.1` 的环境中运行\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python detect_video_tracker_colors.py --video_file test.mp4 --min_score 0.3 --input_size 1024 --model_file model_data/train_mask_rcnn.h5 --model_feature model_data/mars-small128.pb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 复制已经训练的模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import moxing as mox\n",
    "\n",
    "# 复制文件夹\n",
    "# mox.file.copy_parallel('obs://self-ma/ma-mask-rcnn/model', 'obs://self-ma/deep-sort-mask-rcnn/model')\n",
    "\n",
    "# 复制文件\n",
    "# mox.file.copy('s3://self-ma/notebook/out/video/outputVideo.mp4', 's3://dsjga/video/队名cailiao')"
   ]
  }
 ]
}