{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys, os\n",
    "from math import exp as exp\n",
    "import caffe\n",
    "import matplotlib.pyplot as plt\n",
    "from argparse import ArgumentParser\n",
    "logging.basicConfig(format=\"[ %(levelname)s ] %(message)s\", level=logging.INFO, stream=sys.stdout)\n",
    "log = logging.getLogger()\n",
    "caffe_root = '/data/ssd-caffe/new-yolov3-caffe'\n",
    "img_file = '/data/github_repos/yolov3-tiny-fit-ncs/ncs/standard_size_416_persons.jpg'\n",
    "prototxt = '/data/NCS/OpenVINO/tensorflow-yolo-v3/model_optimizer/20190218182235_TinyYoloV3NCS.prototxt'\n",
    "caffemodel = '/data/NCS/OpenVINO/tensorflow-yolo-v3/model_optimizer/20190218182235_TinyYoloV3NCS.caffemodel'\n",
    "sys.path.insert(0, caffe_root + '/python')\n",
    "\n",
    "#%matplotlib inline\n",
    "caffe.set_mode_cpu()\n",
    "net = caffe.Net(prototxt, caffemodel, caffe.TEST)\n",
    "\n",
    "#img_caffe_io = caffe.io.load_image(img_file)\n",
    "#print(img.shape)\n",
    "#plt.imshow(img)\n",
    "import cv2\n",
    "import numpy as np\n",
    "cv2.__version__\n",
    "\n",
    "class YoloV3Params:\n",
    "    # ------------------------------------------- Extracting layer parameters ------------------------------------------\n",
    "    # Magic numbers are copied from yolo samples\n",
    "    def __init__(self, param, side):\n",
    "        self.num = 3 if 'num' not in param else len(param['mask'].split(',')) if 'mask' in param else int(param['num'])\n",
    "        self.coords = 4 if 'coords' not in param else int(param['coords'])\n",
    "        #self.classes = 80 if 'classes' not in param else int(param['classes'])\n",
    "        self.classes = 6 if 'classes' not in param else int(param['classes'])\n",
    "        #self.anchors = [10.0, 13.0, 16.0, 30.0, 33.0, 23.0, 30.0, 61.0, 62.0, 45.0, 59.0, 119.0, 116.0, 90.0, 156.0, 198.0,\n",
    "                        #373.0, 326.0] if 'anchors' not in param else [float(a) for a in param['anchors'].split(',')]\n",
    "        self.anchors = [10,25,  20,50,  30,75, 50,125,  80,200,  150,150] if 'anchors' not in param else [float(a) for a in param['anchors'].split(',')]\n",
    "        self.side = side\n",
    "        if self.side == 13:\n",
    "            self.anchor_offset = 2 * 3\n",
    "        elif self.side == 26:\n",
    "            self.anchor_offset = 2 * 0\n",
    "        elif self.side == 52:\n",
    "            self.anchor_offset = 2 * 0\n",
    "        else:\n",
    "            assert False, \"Invalid output size. Only 13, 26 and 52 sizes are supported for output spatial dimensions\"\n",
    "\n",
    "    def log_params(self):\n",
    "        params_to_print = {'classes': self.classes, 'num': self.num, 'coords': self.coords, 'anchors': self.anchors}\n",
    "        [log.info(\"         {:8}: {}\".format(param_name, param)) for param_name, param in params_to_print.items()]\n",
    "\n",
    "\n",
    "def entry_index(side, coord, classes, location, entry):\n",
    "    side_power_2 = side ** 2\n",
    "    n = location // side_power_2\n",
    "    loc = location % side_power_2\n",
    "    return int(side_power_2 * (n * (coord + classes + 1) + entry) + loc)\n",
    "\n",
    "\n",
    "def scale_bbox(x, y, h, w, class_id, confidence, h_scale, w_scale):\n",
    "    xmin = int((x - w / 2) * w_scale)\n",
    "    ymin = int((y - h / 2) * h_scale)\n",
    "    xmax = int(xmin + w * w_scale)\n",
    "    ymax = int(ymin + h * h_scale)\n",
    "    return dict(xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax, class_id=class_id, confidence=confidence)\n",
    "\n",
    "\n",
    "def parse_yolo_region(blob, resized_image_shape, original_im_shape, params, threshold):\n",
    "    # ------------------------------------------ Validating output parameters ------------------------------------------\n",
    "    _, _, out_blob_h, out_blob_w = blob.shape\n",
    "    assert out_blob_w == out_blob_h, \"Invalid size of output blob. It sould be in NCHW layout and height should \" \\\n",
    "                                     \"be equal to width. Current height = {}, current width = {}\" \\\n",
    "                                     \"\".format(out_blob_h, out_blob_w)\n",
    "\n",
    "    # ------------------------------------------ Extracting layer parameters -------------------------------------------\n",
    "    orig_im_h, orig_im_w = (416, 416)\n",
    "    print(original_im_shape)\n",
    "    resized_image_h, resized_image_w = (416,416)\n",
    "    print(resized_image_shape)\n",
    "    objects = list()\n",
    "    predictions = blob.flatten()\n",
    "    side_square = params.side * params.side\n",
    "\n",
    "    # ------------------------------------------- Parsing YOLO Region output -------------------------------------------\n",
    "    for i in range(side_square):\n",
    "        row = i // params.side\n",
    "        col = i % params.side\n",
    "        for n in range(params.num):\n",
    "            obj_index = entry_index(params.side, params.coords, params.classes, n * side_square + i, params.coords)\n",
    "            scale = predictions[obj_index]\n",
    "            if scale < threshold:\n",
    "                continue\n",
    "            box_index = entry_index(params.side, params.coords, params.classes, n * side_square + i, 0)\n",
    "            x = (col + predictions[box_index + 0 * side_square]) / params.side * resized_image_w\n",
    "            y = (row + predictions[box_index + 1 * side_square]) / params.side * resized_image_h\n",
    "            # Value for exp is very big number in some cases so following construction is using here\n",
    "            try:\n",
    "                w_exp = exp(predictions[box_index + 2 * side_square])\n",
    "                h_exp = exp(predictions[box_index + 3 * side_square])\n",
    "            except OverflowError:\n",
    "                continue\n",
    "            w = w_exp * params.anchors[params.anchor_offset + 2 * n]\n",
    "            h = h_exp * params.anchors[params.anchor_offset + 2 * n + 1]\n",
    "            for j in range(params.classes):\n",
    "                class_index = entry_index(params.side, params.coords, params.classes, n * side_square + i,\n",
    "                                          params.coords + 1 + j)\n",
    "                confidence = scale * predictions[class_index]\n",
    "                if confidence < threshold:\n",
    "                    continue\n",
    "                objects.append(scale_bbox(x=x, y=y, h=h, w=w, class_id=j, confidence=confidence,\n",
    "                                          h_scale=orig_im_h / resized_image_h, w_scale=orig_im_w / resized_image_w))\n",
    "    return objects\n",
    "\n",
    "\n",
    "def intersection_over_union(box_1, box_2):\n",
    "    width_of_overlap_area = min(box_1['xmax'], box_2['xmax']) - max(box_1['xmin'], box_2['xmin'])\n",
    "    height_of_overlap_area = min(box_1['ymax'], box_2['ymax']) - max(box_1['ymin'], box_2['ymin'])\n",
    "    if width_of_overlap_area < 0 or height_of_overlap_area < 0:\n",
    "        area_of_overlap = 0\n",
    "    else:\n",
    "        area_of_overlap = width_of_overlap_area * height_of_overlap_area\n",
    "    box_1_area = (box_1['ymax'] - box_1['ymin']) * (box_1['xmax'] - box_1['xmin'])\n",
    "    box_2_area = (box_2['ymax'] - box_2['ymin']) * (box_2['xmax'] - box_2['xmin'])\n",
    "    area_of_union = box_1_area + box_2_area - area_of_overlap\n",
    "    if area_of_union == 0:\n",
    "        return 0\n",
    "    return area_of_overlap / area_of_union\n",
    "\n",
    "def build_argparser():\n",
    "    parser = ArgumentParser()\n",
    "    parser.add_argument(\"-pt\", \"--prob_threshold\", help=\"Probability threshold for detections filtering\",\n",
    "                        default=0.5, type=float)\n",
    "    parser.add_argument(\"-iout\", \"--iou_threshold\", help=\"Intersection over union threshold for overlapping detections\"\n",
    "                                                         \" filtering\", default=0.4, type=float)\n",
    "    parser.add_argument(\"-ni\", \"--number_iter\", help=\"Number of inference iterations\", default=1, type=int)\n",
    "    parser.add_argument(\"-pc\", \"--perf_counts\", help=\"Report performance counters\", default=False, action=\"store_true\")\n",
    "    parser.add_argument(\"-i\", \"--input\", help=\"input_video\", default='/video/09_58_09.264', type=str)\n",
    "    return parser\n",
    "\n",
    "\n",
    "objects = list()\n",
    "args = build_argparser().parse_args()\n",
    "video_file = args.input\n",
    "\n",
    "cap = cv2.VideoCapture(video_file)\n",
    "while cap.isOpened():\n",
    "    ret, img = cap.read()\n",
    "    if ret == True:\n",
    "    #img = cv2.imread(img_file) ###BGR\n",
    "        resized_img = cv2.resize(img, (416,416))\n",
    "        resized_img = cv2.cvtColor(resized_img, cv2.COLOR_BGR2RGB) # cv2默认为bgr顺序\n",
    "        img_ = np.transpose(resized_img, (2, 0, 1))\n",
    "        print(img_.shape)\n",
    "        #print(img_)\n",
    "        normerlized_img = img_/255.0\n",
    "        #normerlized_img = img_\n",
    "        #print(normerlized_img)\n",
    "\n",
    "        net.blobs['data'].data[...] = normerlized_img\n",
    "        inputData = net.blobs['data'].data[...]\n",
    "        #inputData\n",
    "        outputs = net.forward()\n",
    "        print(outputs)\n",
    "\n",
    "\n",
    "        for layer_name, out_blob in outputs.items():\n",
    "            print(layer_name)\n",
    "            print(out_blob)\n",
    "            #layer_params = YoloV3Params(net.layers[layer_name].params, out_blob.shape[2])\n",
    "            layer_params = YoloV3Params('', out_blob.shape[2])\n",
    "\n",
    "            log.info(\"Layer {} parameters: \".format(layer_name))\n",
    "            layer_params.log_params()\n",
    "            in_frame = inputData\n",
    "            frame = img\n",
    "            objects += parse_yolo_region(out_blob, in_frame.shape[2:],\n",
    "                                     frame.shape[:-1], layer_params,\n",
    "                                     args.prob_threshold)\n",
    "\n",
    "        # Filtering overlapping boxes with respect to the --iou_threshold CLI parameter\n",
    "        for i in range(len(objects)):\n",
    "            if objects[i]['confidence'] == 0:\n",
    "            continue\n",
    "            for j in range(i + 1, len(objects)):\n",
    "            if intersection_over_union(objects[i], objects[j]) > args.iou_threshold:\n",
    "                objects[j]['confidence'] = 0\n",
    "\n",
    "        # Drawing objects with respect to the --prob_threshold CLI parameter\n",
    "        objects = [obj for obj in objects if obj['confidence'] >= args.prob_threshold]\n",
    "\n",
    "        if len(objects):\n",
    "            log.info(\"\\nDetected boxes for batch {}:\".format(1))\n",
    "            log.info(\" Class ID | Confidence | XMIN | YMIN | XMAX | YMAX | COLOR \")\n",
    "\n",
    "        origin_im_size = frame.shape[:-1]\n",
    "        for obj in objects:\n",
    "            # Validation bbox of detected object\n",
    "            if obj['xmax'] > origin_im_size[0] or obj['ymax'] > origin_im_size[0] or obj['xmin'] < 0 or obj['ymin'] < 0:\n",
    "            continue\n",
    "            color = (int(min(obj['class_id'] * 12.5, 255)),\n",
    "                 min(obj['class_id'] * 7, 255), min(obj['class_id'] * 5, 255))\n",
    "            #det_label = labels_map[obj['class_id']] if labels_map and len(labels_map) >= obj['class_id'] else str(obj['class_id'])\n",
    "            log.info(\"{:10f} | {:4} | {:4} | {:4} | {:4} | {} \".format(obj['confidence'], obj['xmin'],\n",
    "                                                                           obj['ymin'], obj['xmax'], obj['ymax'],\n",
    "                                                                           color))\n",
    "\n",
    "            cv2.rectangle(frame, (obj['xmin'], obj['ymin']), (obj['xmax'], obj['ymax']), color, 2)\n",
    "            cv2.putText(frame,\n",
    "                    \"#\" + ' ' + str(round(obj['confidence'] * 100, 1)) + ' %',\n",
    "                    (obj['xmin'], obj['ymin'] - 7), cv2.FONT_HERSHEY_COMPLEX, 0.6, color, 1)\n",
    "\n",
    "        # Draw performance stats over frame\n",
    "        cv2.putText(frame, \"test\", (15, 30), cv2.FONT_HERSHEY_COMPLEX, 0.5, (10, 10, 200), 1)\n",
    "        cv2.imshow(\"DetectionResults\", frame)\n",
    "\n",
    "\n",
    "        cv2.waitKey(1)\n",
    "        \n",
    "        # Tab key\n",
    "        if key == 27:\n",
    "            break\n",
    "        # ESC key\n",
    "        if key == 9:\n",
    "            break   \n",
    "        \n",
    "cv2.destroyAllWindows()\n",
    "del net"
   ]
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
