{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "yolov3.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "p0y3wIkfSuIT",
        "colab_type": "code",
        "outputId": "eeedd664-406a-43ff-aa5e-bd48963494c4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 53
        }
      },
      "source": [
        "%tensorflow_version 1.x\n",
        "## Check your google colab/drive settings!!! (libraries, argument paths, ...)\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "\n",
        "## variables for notebook\n",
        "training = True\n",
        "\n",
        "##### changes\n",
        "### changed some variable names because of argument conflicts\n",
        "### last two parts are train, test mode code. you can switch the mode with above variable, 'training'\n",
        "### there are some difficulties for separating train/eval code (making into functions), because of variable dependencies"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "TensorFlow 1.x selected.\n",
            "Drive already mounted at /content/gdrive; to attempt to forcibly remount, call drive.mount(\"/content/gdrive\", force_remount=True).\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Yh3RWBkgAjZx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "## TFRecord utils here\n",
        "import tensorflow as tf\n",
        "from itertools import tee\n",
        "\n",
        "class TFRecordIterator:\n",
        "    def __init__(self, path, compression=None):\n",
        "        self._core = tf.python_io.tf_record_iterator(path, tf.python_io.TFRecordOptions(compression))\n",
        "        self._iterator = iter(self._core)\n",
        "        self._iterator, self._iterator_temp = tee(self._iterator)\n",
        "        self._total_cnt = sum(1 for _ in self._iterator_temp)\n",
        "\n",
        "    def _read_value(self, feature):\n",
        "        if len(feature.int64_list.value) > 0:\n",
        "            return feature.int64_list.value\n",
        "\n",
        "        if len(feature.bytes_list.value) > 0:\n",
        "            return feature.bytes_list.value\n",
        "\n",
        "        if len(feature.float_list.value) > 0:\n",
        "            return feature.float_list.value\n",
        "\n",
        "        return None\n",
        "\n",
        "    def _read_features(self, features):\n",
        "        d = dict()\n",
        "        for data in features:\n",
        "            d[data] = self._read_value(features[data])\n",
        "        return d\n",
        "\n",
        "    def __enter__(self):\n",
        "        return self\n",
        "\n",
        "    def __exit__(self, exception_type, exception_value, traceback):\n",
        "        pass\n",
        "\n",
        "    def __iter__(self):\n",
        "        return self\n",
        "\n",
        "    def __next__(self):\n",
        "        record = next(self._iterator)\n",
        "        example = tf.train.Example()\n",
        "        example.ParseFromString(record)\n",
        "        return self._read_features(example.features.feature)\n",
        "\n",
        "    def count(self):\n",
        "        return self._total_cnt\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oCVOPE2XC3qE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "## plot utils\n",
        "from __future__ import division, print_function\n",
        "\n",
        "import cv2\n",
        "import random\n",
        "\n",
        "def get_color_table(class_num, seed=2):\n",
        "    random.seed(seed)\n",
        "    color_table = {}\n",
        "    for i in range(class_num):\n",
        "        color_table[i] = [random.randint(0, 255) for _ in range(3)]\n",
        "    return color_table\n",
        "\n",
        "\n",
        "def plot_one_box(img, coord, label=None, color=None, line_thickness=None):\n",
        "    tl = line_thickness or int(round(0.002 * max(img.shape[0:2])))  # line thickness\n",
        "    color = color or [random.randint(0, 255) for _ in range(3)]\n",
        "    c1, c2 = (int(coord[0]), int(coord[1])), (int(coord[2]), int(coord[3]))\n",
        "    cv2.rectangle(img, c1, c2, color, thickness=tl)\n",
        "    if label:\n",
        "        tf = max(tl - 1, 1)  # font thickness\n",
        "        t_size = cv2.getTextSize(label, 0, fontScale=float(tl) / 3, thickness=tf)[0]\n",
        "        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3\n",
        "        cv2.rectangle(img, c1, c2, color, -1)  # filled\n",
        "        cv2.putText(img, label, (c1[0], c1[1] - 2), 0, float(tl) / 3, [0, 0, 0], thickness=tf, lineType=cv2.LINE_AA)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SY10K9LoDJOZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "## nms utils\n",
        "import numpy as np\n",
        "\n",
        "def gpu_nms(boxes, scores, num_classes, max_boxes=50, score_thresh=0.5, nms_thresh=0.5):\n",
        "    boxes_list, label_list, score_list = [], [], []\n",
        "    max_boxes = tf.constant(max_boxes, dtype='int32')\n",
        "\n",
        "    boxes = tf.reshape(boxes, [-1, 4]) # '-1' means we don't konw the exact number of boxes\n",
        "    score = tf.reshape(scores, [-1, num_classes])\n",
        "\n",
        "    # Step 1: Create a filtering mask based on \"box_class_scores\" by using \"threshold\".\n",
        "    mask = tf.greater_equal(score, tf.constant(score_thresh))\n",
        "    # Step 2: Do non_max_suppression for each class\n",
        "    for i in range(num_classes):\n",
        "        # Step 3: Apply the mask to scores, boxes and pick them out\n",
        "        filter_boxes = tf.boolean_mask(boxes, mask[:,i])\n",
        "        filter_score = tf.boolean_mask(score[:,i], mask[:,i])\n",
        "        nms_indices = tf.image.non_max_suppression(boxes=filter_boxes,\n",
        "                                                   scores=filter_score,\n",
        "                                                   max_output_size=max_boxes,\n",
        "                                                   iou_threshold=nms_thresh, name='nms_indices')\n",
        "        label_list.append(tf.ones_like(tf.gather(filter_score, nms_indices), 'int32')*i)\n",
        "        boxes_list.append(tf.gather(filter_boxes, nms_indices))\n",
        "        score_list.append(tf.gather(filter_score, nms_indices))\n",
        "\n",
        "    boxes = tf.concat(boxes_list, axis=0)\n",
        "    score = tf.concat(score_list, axis=0)\n",
        "    label = tf.concat(label_list, axis=0)\n",
        "\n",
        "    return boxes, score, label\n",
        "\n",
        "\n",
        "def py_nms(boxes, scores, max_boxes=50, iou_thresh=0.5):\n",
        "    assert boxes.shape[1] == 4 and len(scores.shape) == 1\n",
        "\n",
        "    x1 = boxes[:, 0]\n",
        "    y1 = boxes[:, 1]\n",
        "    x2 = boxes[:, 2]\n",
        "    y2 = boxes[:, 3]\n",
        "\n",
        "    areas = (x2 - x1) * (y2 - y1)\n",
        "    order = scores.argsort()[::-1]\n",
        "\n",
        "    keep = []\n",
        "    while order.size > 0:\n",
        "        i = order[0]\n",
        "        keep.append(i)\n",
        "        xx1 = np.maximum(x1[i], x1[order[1:]])\n",
        "        yy1 = np.maximum(y1[i], y1[order[1:]])\n",
        "        xx2 = np.minimum(x2[i], x2[order[1:]])\n",
        "        yy2 = np.minimum(y2[i], y2[order[1:]])\n",
        "\n",
        "        w = np.maximum(0.0, xx2 - xx1 + 1)\n",
        "        h = np.maximum(0.0, yy2 - yy1 + 1)\n",
        "        inter = w * h\n",
        "        ovr = inter / (areas[i] + areas[order[1:]] - inter)\n",
        "\n",
        "        inds = np.where(ovr <= iou_thresh)[0]\n",
        "        order = order[inds + 1]\n",
        "\n",
        "    return keep[:max_boxes]\n",
        "\n",
        "\n",
        "def cpu_nms(boxes, scores, num_classes, max_boxes=50, score_thresh=0.5, iou_thresh=0.5):\n",
        "    boxes = boxes.reshape(-1, 4)\n",
        "    scores = scores.reshape(-1, num_classes)\n",
        "    picked_boxes, picked_score, picked_label = [], [], []\n",
        "\n",
        "    for i in range(num_classes):\n",
        "        indices = np.where(scores[:,i] >= score_thresh)\n",
        "        filter_boxes = boxes[indices]\n",
        "        filter_scores = scores[:,i][indices]\n",
        "        if len(filter_boxes) == 0: \n",
        "            continue\n",
        "\n",
        "        indices = py_nms(filter_boxes, filter_scores,\n",
        "                         max_boxes=max_boxes, iou_thresh=iou_thresh)\n",
        "        picked_boxes.append(filter_boxes[indices])\n",
        "        picked_score.append(filter_scores[indices])\n",
        "        picked_label.append(np.ones(len(indices), dtype='int32')*i)\n",
        "    if len(picked_boxes) == 0: \n",
        "        return None, None, None\n",
        "\n",
        "    boxes = np.concatenate(picked_boxes, axis=0)\n",
        "    score = np.concatenate(picked_score, axis=0)\n",
        "    label = np.concatenate(picked_label, axis=0)\n",
        "\n",
        "    return boxes, score, label"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Dg-ZKHmRDlPp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "## misc utils\n",
        "class AverageMeter(object):\n",
        "    def __init__(self):\n",
        "        self.reset()\n",
        "\n",
        "    def reset(self):\n",
        "        self.val = 0\n",
        "        self.average = 0\n",
        "        self.sum = 0\n",
        "        self.count = 0\n",
        "\n",
        "    def update(self, val, n=1):\n",
        "        self.val = val\n",
        "        self.sum += val * n\n",
        "        self.count += n\n",
        "        self.average = self.sum / float(self.count)\n",
        "\n",
        "\n",
        "def parse_anchors(anchor_path):\n",
        "    anchors = np.reshape(np.asarray(open(anchor_path, 'r').read().split(','), np.float32), [-1, 2])\n",
        "    return anchors\n",
        "\n",
        "\n",
        "def read_class_names(class_name_path):\n",
        "    names = {}\n",
        "    with open(class_name_path, 'r') as data:\n",
        "        for ID, name in enumerate(data):\n",
        "            names[ID] = name.strip('\\n')\n",
        "    return names\n",
        "\n",
        "\n",
        "def shuffle_and_overwrite(file_name):\n",
        "    content = open(file_name, 'r').readlines()\n",
        "    random.shuffle(content)\n",
        "    with open(file_name, 'w') as f:\n",
        "        for line in content:\n",
        "            f.write(line)\n",
        "\n",
        "\n",
        "def update_dict(ori_dict, new_dict):\n",
        "    if not ori_dict:\n",
        "        return new_dict\n",
        "    for key in ori_dict:\n",
        "        ori_dict[key] += new_dict[key]\n",
        "    return ori_dict\n",
        "\n",
        "\n",
        "def list_add(ori_list, new_list):\n",
        "    for i in range(len(ori_list)):\n",
        "        ori_list[i] += new_list[i]\n",
        "    return ori_list\n",
        "\n",
        "\n",
        "def load_weights(var_list, weights_file):\n",
        "    with open(weights_file, \"rb\") as fp:\n",
        "        np.fromfile(fp, dtype=np.int32, count=5)\n",
        "        weights = np.fromfile(fp, dtype=np.float32)\n",
        "\n",
        "    ptr = 0\n",
        "    i = 0\n",
        "    assign_ops = []\n",
        "    while i < len(var_list) - 1:\n",
        "        var1 = var_list[i]\n",
        "        var2 = var_list[i + 1]\n",
        "        if 'Conv' in var1.name.split('/')[-2]:\n",
        "            if 'BatchNorm' in var2.name.split('/')[-2]:\n",
        "                gamma, beta, mean, var = var_list[i + 1:i + 5]\n",
        "                batch_norm_vars = [beta, gamma, mean, var]\n",
        "                for var in batch_norm_vars:\n",
        "                    shape = var.shape.as_list()\n",
        "                    num_params = np.prod(shape)\n",
        "                    var_weights = weights[ptr:ptr + num_params].reshape(shape)\n",
        "                    ptr += num_params\n",
        "                    assign_ops.append(tf.assign(var, var_weights, validate_shape=True))\n",
        "                i += 4\n",
        "            elif 'Conv' in var2.name.split('/')[-2]:\n",
        "                # load biases\n",
        "                bias = var2\n",
        "                bias_shape = bias.shape.as_list()\n",
        "                bias_params = np.prod(bias_shape)\n",
        "                bias_weights = weights[ptr:ptr +\n",
        "                                       bias_params].reshape(bias_shape)\n",
        "                ptr += bias_params\n",
        "                assign_ops.append(tf.assign(bias, bias_weights, validate_shape=True))\n",
        "                i += 1\n",
        "\n",
        "            shape = var1.shape.as_list()\n",
        "            num_params = np.prod(shape)\n",
        "\n",
        "            var_weights = weights[ptr:ptr + num_params].reshape(\n",
        "                (shape[3], shape[2], shape[0], shape[1]))\n",
        "\n",
        "            var_weights = np.transpose(var_weights, (2, 3, 1, 0))\n",
        "            ptr += num_params\n",
        "            assign_ops.append(\n",
        "                tf.assign(var1, var_weights, validate_shape=True))\n",
        "            i += 1\n",
        "\n",
        "    return assign_ops\n",
        "\n",
        "\n",
        "def config_learning_rate(global_step):\n",
        "    ## fixes for removing arg paramter\n",
        "    global lr_type, learning_rate_init, lr_decay_freq, lr_decay_factor, lr_lower_bound, total_epoches, use_warm_up, warm_up_epoch, train_batch_num, lr_lower_bound, pw_boundaries, pw_values\n",
        "\n",
        "    if lr_type == 'exponential':\n",
        "        lr_tmp = tf.train.exponential_decay(learning_rate_init, global_step, lr_decay_freq,\n",
        "                                            lr_decay_factor, staircase=True, name='exponential_learning_rate')\n",
        "        return tf.maximum(lr_tmp, lr_lower_bound)\n",
        "    elif lr_type == 'cosine_decay':\n",
        "        train_steps = (total_epoches - float(use_warm_up) * warm_up_epoch) * train_batch_num\n",
        "        return lr_lower_bound + 0.5 * (learning_rate_init - lr_lower_bound) * \\\n",
        "            (1 + tf.cos(global_step / train_steps * np.pi))\n",
        "    elif lr_type == 'cosine_decay_restart':\n",
        "        return tf.train.cosine_decay_restarts(learning_rate_init, global_step, \n",
        "                                              lr_decay_freq, t_mul=2.0, m_mul=1.0, \n",
        "                                              name='cosine_decay_learning_rate_restart')\n",
        "    elif lr_type == 'fixed':\n",
        "        return tf.convert_to_tensor(learning_rate_init, name='fixed_learning_rate')\n",
        "    elif lr_type == 'piecewise':\n",
        "        return tf.train.piecewise_constant(global_step, boundaries=pw_boundaries, values=pw_values,\n",
        "                                           name='piecewise_learning_rate')\n",
        "    else:\n",
        "        raise ValueError('Unsupported learning rate type!')\n",
        "\n",
        "\n",
        "def config_optimizer(optimizer_name, learning_rate, decay=0.9, momentum=0.9):\n",
        "    if optimizer_name == 'momentum':\n",
        "        return tf.train.MomentumOptimizer(learning_rate, momentum=momentum)\n",
        "    elif optimizer_name == 'rmsprop':\n",
        "        return tf.train.RMSPropOptimizer(learning_rate, decay=decay, momentum=momentum)\n",
        "    elif optimizer_name == 'adam':\n",
        "        return tf.train.AdamOptimizer(learning_rate)\n",
        "    elif optimizer_name == 'sgd':\n",
        "        return tf.train.GradientDescentOptimizer(learning_rate)\n",
        "    else:\n",
        "        raise ValueError('Unsupported optimizer type!')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YIlZhFLYD0d8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "## data utils\n",
        "\n",
        "import sys\n",
        "\n",
        "PY_VERSION = sys.version_info[0]\n",
        "iter_cnt = 0\n",
        "\n",
        "def _parse_tfrecord(data):\n",
        "    example = tf.train.Example()\n",
        "    example.ParseFromString(data)\n",
        "    features = example.features.feature\n",
        "    return features\n",
        "\n",
        "def parse_tfrecord(data):\n",
        "    # tfrecord parser for TFRecordDataset (raw data)\n",
        "    features = _parse_tfrecord(data)\n",
        "    index = features['index'].int64_list.value[0]\n",
        "    encoded_image = np.frombuffer(features['image'].bytes_list.value[0], dtype = np.uint8)\n",
        "    width = features['width'].int64_list.value[0]\n",
        "    height = features['height'].int64_list.value[0]\n",
        "    boxes = features['boxes'].int64_list.value\n",
        "\n",
        "    assert len(boxes) % 5 == 0, 'Annotation error occured in box array.'\n",
        "    box_cnt = len(boxes) // 5\n",
        "\n",
        "    aligned_boxes = []\n",
        "    labels = []\n",
        "\n",
        "    for i in range(box_cnt):\n",
        "        label, x_min, y_min, x_max, y_max = int(boxes[i * 5]), float(boxes[i * 5 + 1]), float(boxes[i * 5 + 2]), float(boxes[i * 5 + 3]), float(boxes[i * 5 + 4]) ## do we need to change int to float? is there float rectangle sample?\n",
        "        aligned_boxes.append([x_min, y_min, x_max, y_max])\n",
        "        labels.append(label)\n",
        "\n",
        "    aligned_boxes = np.asarray(aligned_boxes, np.float32)\n",
        "    labels = np.asarray(labels, np.int64)\n",
        "\n",
        "    return index, encoded_image, aligned_boxes, labels, width, height\n",
        "\n",
        "def parse_record(features):\n",
        "    # tfrecord parser for TFRecordIterator (primitive data)\n",
        "\n",
        "    index = int(features['index'][0])\n",
        "    encoded_image = np.frombuffer(features['image'][0], dtype = np.uint8)\n",
        "    width = int(features['width'][0])\n",
        "    height = int(features['height'][0])\n",
        "    boxes = features['boxes']\n",
        "\n",
        "    assert len(boxes) % 5 == 0, 'Annotation error occured in box array.'\n",
        "    box_cnt = len(boxes) // 5\n",
        "\n",
        "    aligned_boxes = []\n",
        "    labels = []\n",
        "\n",
        "    for i in range(box_cnt):\n",
        "        label, x_min, y_min, x_max, y_max = int(boxes[i * 5]), float(boxes[i * 5 + 1]), float(boxes[i * 5 + 2]), float(boxes[i * 5 + 3]), float(boxes[i * 5 + 4])\n",
        "        aligned_boxes.append([x_min, y_min, x_max, y_max])\n",
        "        labels.append(label)\n",
        "\n",
        "    aligned_boxes = np.asarray(aligned_boxes, np.float32)\n",
        "    labels = np.asarray(labels, np.int64)\n",
        "\n",
        "    return index, encoded_image, aligned_boxes, labels, width, height\n",
        "\n",
        "def bbox_crop(bbox, crop_box=None, allow_outside_center=True):\n",
        "    bbox = bbox.copy()\n",
        "    if crop_box is None:\n",
        "        return bbox\n",
        "    if not len(crop_box) == 4:\n",
        "        raise ValueError(\n",
        "            \"Invalid crop_box parameter, requires length 4, given {}\".format(str(crop_box)))\n",
        "    if sum([int(c is None) for c in crop_box]) == 4:\n",
        "        return bbox\n",
        "\n",
        "    l, t, w, h = crop_box\n",
        "\n",
        "    left = l if l else 0\n",
        "    top = t if t else 0\n",
        "    right = left + (w if w else np.inf)\n",
        "    bottom = top + (h if h else np.inf)\n",
        "    crop_bbox = np.array((left, top, right, bottom))\n",
        "\n",
        "    if allow_outside_center:\n",
        "        mask = np.ones(bbox.shape[0], dtype=bool)\n",
        "    else:\n",
        "        centers = (bbox[:, :2] + bbox[:, 2:4]) / 2\n",
        "        mask = np.logical_and(crop_bbox[:2] <= centers, centers < crop_bbox[2:]).all(axis=1)\n",
        "\n",
        "    # transform borders\n",
        "    bbox[:, :2] = np.maximum(bbox[:, :2], crop_bbox[:2])\n",
        "    bbox[:, 2:4] = np.minimum(bbox[:, 2:4], crop_bbox[2:4])\n",
        "    bbox[:, :2] -= crop_bbox[:2]\n",
        "    bbox[:, 2:4] -= crop_bbox[:2]\n",
        "\n",
        "    mask = np.logical_and(mask, (bbox[:, :2] < bbox[:, 2:4]).all(axis=1))\n",
        "    bbox = bbox[mask]\n",
        "    return bbox\n",
        "\n",
        "def bbox_iou(bbox_a, bbox_b, offset=0):\n",
        "    if bbox_a.shape[1] < 4 or bbox_b.shape[1] < 4:\n",
        "        raise IndexError(\"Bounding boxes axis 1 must have at least length 4\")\n",
        "\n",
        "    tl = np.maximum(bbox_a[:, None, :2], bbox_b[:, :2])\n",
        "    br = np.minimum(bbox_a[:, None, 2:4], bbox_b[:, 2:4])\n",
        "\n",
        "    area_i = np.prod(br - tl + offset, axis=2) * (tl < br).all(axis=2)\n",
        "    area_a = np.prod(bbox_a[:, 2:4] - bbox_a[:, :2] + offset, axis=1)\n",
        "    area_b = np.prod(bbox_b[:, 2:4] - bbox_b[:, :2] + offset, axis=1)\n",
        "    return area_i / (area_a[:, None] + area_b - area_i)\n",
        "\n",
        "\n",
        "def random_crop_with_constraints(bbox, size, min_scale=0.3, max_scale=1,\n",
        "                                 max_aspect_ratio=2, constraints=None,\n",
        "                                 max_trial=50):\n",
        "    # default params in paper\n",
        "    if constraints is None:\n",
        "        constraints = (\n",
        "            (0.1, None),\n",
        "            (0.3, None),\n",
        "            (0.5, None),\n",
        "            (0.7, None),\n",
        "            (0.9, None),\n",
        "            (None, 1),\n",
        "        )\n",
        "\n",
        "    w, h = size\n",
        "\n",
        "    candidates = [(0, 0, w, h)]\n",
        "    for min_iou, max_iou in constraints:\n",
        "        min_iou = -np.inf if min_iou is None else min_iou\n",
        "        max_iou = np.inf if max_iou is None else max_iou\n",
        "\n",
        "        for _ in range(max_trial):\n",
        "            scale = random.uniform(min_scale, max_scale)\n",
        "            aspect_ratio = random.uniform(\n",
        "                max(1 / max_aspect_ratio, scale * scale),\n",
        "                min(max_aspect_ratio, 1 / (scale * scale)))\n",
        "            crop_h = int(h * scale / np.sqrt(aspect_ratio))\n",
        "            crop_w = int(w * scale * np.sqrt(aspect_ratio))\n",
        "\n",
        "            crop_t = random.randrange(h - crop_h)\n",
        "            crop_l = random.randrange(w - crop_w)\n",
        "            crop_bb = np.array((crop_l, crop_t, crop_l + crop_w, crop_t + crop_h))\n",
        "\n",
        "            if len(bbox) == 0:\n",
        "                top, bottom = crop_t, crop_t + crop_h\n",
        "                left, right = crop_l, crop_l + crop_w\n",
        "                return bbox, (left, top, right-left, bottom-top)\n",
        "\n",
        "            iou = bbox_iou(bbox, crop_bb[np.newaxis])\n",
        "            if min_iou <= iou.min() and iou.max() <= max_iou:\n",
        "                top, bottom = crop_t, crop_t + crop_h\n",
        "                left, right = crop_l, crop_l + crop_w\n",
        "                candidates.append((left, top, right-left, bottom-top))\n",
        "                break\n",
        "\n",
        "    # random select one\n",
        "    while candidates:\n",
        "        crop = candidates.pop(np.random.randint(0, len(candidates)))\n",
        "        new_bbox = bbox_crop(bbox, crop, allow_outside_center=False)\n",
        "        if new_bbox.size < 1:\n",
        "            continue\n",
        "        new_crop = (crop[0], crop[1], crop[2], crop[3])\n",
        "        return new_bbox, new_crop\n",
        "    return bbox, (0, 0, w, h)\n",
        "\n",
        "\n",
        "def random_color_distort(img, brightness_delta=32, hue_vari=18, sat_vari=0.5, val_vari=0.5):\n",
        "    def random_hue(img_hsv, hue_vari, p=0.5):\n",
        "        if np.random.uniform(0, 1) > p:\n",
        "            hue_delta = np.random.randint(-hue_vari, hue_vari)\n",
        "            img_hsv[:, :, 0] = (img_hsv[:, :, 0] + hue_delta) % 180\n",
        "        return img_hsv\n",
        "\n",
        "    def random_saturation(img_hsv, sat_vari, p=0.5):\n",
        "        if np.random.uniform(0, 1) > p:\n",
        "            sat_mult = 1 + np.random.uniform(-sat_vari, sat_vari)\n",
        "            img_hsv[:, :, 1] *= sat_mult\n",
        "        return img_hsv\n",
        "\n",
        "    def random_value(img_hsv, val_vari, p=0.5):\n",
        "        if np.random.uniform(0, 1) > p:\n",
        "            val_mult = 1 + np.random.uniform(-val_vari, val_vari)\n",
        "            img_hsv[:, :, 2] *= val_mult\n",
        "        return img_hsv\n",
        "\n",
        "    def random_brightness(img, brightness_delta, p=0.5):\n",
        "        if np.random.uniform(0, 1) > p:\n",
        "            img = img.astype(np.float32)\n",
        "            brightness_delta = int(np.random.uniform(-brightness_delta, brightness_delta))\n",
        "            img = img + brightness_delta\n",
        "        return np.clip(img, 0, 255)\n",
        "\n",
        "    # brightness\n",
        "    img = random_brightness(img, brightness_delta)\n",
        "    img = img.astype(np.uint8)\n",
        "\n",
        "    # color jitter\n",
        "    img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float32)\n",
        "\n",
        "    if np.random.randint(0, 2):\n",
        "        img_hsv = random_value(img_hsv, val_vari)\n",
        "        img_hsv = random_saturation(img_hsv, sat_vari)\n",
        "        img_hsv = random_hue(img_hsv, hue_vari)\n",
        "    else:\n",
        "        img_hsv = random_saturation(img_hsv, sat_vari)\n",
        "        img_hsv = random_hue(img_hsv, hue_vari)\n",
        "        img_hsv = random_value(img_hsv, val_vari)\n",
        "\n",
        "    img_hsv = np.clip(img_hsv, 0, 255)\n",
        "    img = cv2.cvtColor(img_hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)\n",
        "\n",
        "    return img\n",
        "\n",
        "\n",
        "def letterbox_resize(img, new_width, new_height, interp=0):\n",
        "    ori_height, ori_width = img.shape[:2]\n",
        "\n",
        "    resize_ratio = min(new_width / ori_width, new_height / ori_height)\n",
        "\n",
        "    resize_w = int(resize_ratio * ori_width)\n",
        "    resize_h = int(resize_ratio * ori_height)\n",
        "\n",
        "    img = cv2.resize(img, (resize_w, resize_h), interpolation=interp)\n",
        "    image_padded = np.full((new_height, new_width, 3), 128, np.uint8)\n",
        "\n",
        "    dw = int((new_width - resize_w) / 2)\n",
        "    dh = int((new_height - resize_h) / 2)\n",
        "\n",
        "    image_padded[dh: resize_h + dh, dw: resize_w + dw, :] = img\n",
        "\n",
        "    return image_padded, resize_ratio, dw, dh\n",
        "\n",
        "\n",
        "def resize_with_bbox(img, bbox, new_width, new_height, interp=0, letterbox=False):\n",
        "    if letterbox:\n",
        "        image_padded, resize_ratio, dw, dh = letterbox_resize(img, new_width, new_height, interp)\n",
        "\n",
        "        # xmin, xmax\n",
        "        bbox[:, [0, 2]] = bbox[:, [0, 2]] * resize_ratio + dw\n",
        "        # ymin, ymax\n",
        "        bbox[:, [1, 3]] = bbox[:, [1, 3]] * resize_ratio + dh\n",
        "\n",
        "        return image_padded, bbox\n",
        "    else:\n",
        "        ori_height, ori_width = img.shape[:2]\n",
        "\n",
        "        img = cv2.resize(img, (new_width, new_height), interpolation=interp)\n",
        "\n",
        "        # xmin, xmax\n",
        "        bbox[:, [0, 2]] = bbox[:, [0, 2]] / ori_width * new_width\n",
        "        # ymin, ymax\n",
        "        bbox[:, [1, 3]] = bbox[:, [1, 3]] / ori_height * new_height\n",
        "\n",
        "        return img, bbox\n",
        "\n",
        "\n",
        "def random_flip(img, bbox, px=0, py=0):\n",
        "    height, width = img.shape[:2]\n",
        "    if np.random.uniform(0, 1) < px:\n",
        "        img = cv2.flip(img, 1)\n",
        "        xmax = width - bbox[:, 0]\n",
        "        xmin = width - bbox[:, 2]\n",
        "        bbox[:, 0] = xmin\n",
        "        bbox[:, 2] = xmax\n",
        "\n",
        "    if np.random.uniform(0, 1) < py:\n",
        "        img = cv2.flip(img, 0)\n",
        "        ymax = height - bbox[:, 1]\n",
        "        ymin = height - bbox[:, 3]\n",
        "        bbox[:, 1] = ymin\n",
        "        bbox[:, 3] = ymax\n",
        "    return img, bbox\n",
        "\n",
        "\n",
        "def random_expand(img, bbox, max_ratio=4, fill=0, keep_ratio=True):\n",
        "    h, w, c = img.shape\n",
        "    ratio_x = random.uniform(1, max_ratio)\n",
        "    if keep_ratio:\n",
        "        ratio_y = ratio_x\n",
        "    else:\n",
        "        ratio_y = random.uniform(1, max_ratio)\n",
        "\n",
        "    oh, ow = int(h * ratio_y), int(w * ratio_x)\n",
        "    off_y = random.randint(0, oh - h)\n",
        "    off_x = random.randint(0, ow - w)\n",
        "\n",
        "    dst = np.full(shape=(oh, ow, c), fill_value=fill, dtype=img.dtype)\n",
        "\n",
        "    dst[off_y:off_y + h, off_x:off_x + w, :] = img\n",
        "\n",
        "    # correct bbox\n",
        "    bbox[:, :2] += (off_x, off_y)\n",
        "    bbox[:, 2:4] += (off_x, off_y)\n",
        "\n",
        "    return dst, bbox\n",
        "\n",
        "def process_box(boxes, labels, img_size, class_num, anchors):\n",
        "    anchors_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]\n",
        "\n",
        "    # convert boxes form:\n",
        "    # shape: [N, 2]\n",
        "    # (x_center, y_center)\n",
        "    box_centers = (boxes[:, 0:2] + boxes[:, 2:4]) / 2\n",
        "    # (width, height)\n",
        "    box_sizes = boxes[:, 2:4] - boxes[:, 0:2]\n",
        "\n",
        "    # [13, 13, 3, 5+num_class+1] `5` means coords and labels. `1` means mix up weight. \n",
        "    y_true_13 = np.zeros((img_size[1] // 32, img_size[0] // 32, 3, 6 + class_num), np.float32)\n",
        "    y_true_26 = np.zeros((img_size[1] // 16, img_size[0] // 16, 3, 6 + class_num), np.float32)\n",
        "    y_true_52 = np.zeros((img_size[1] // 8, img_size[0] // 8, 3, 6 + class_num), np.float32)\n",
        "\n",
        "    # mix up weight default to 1.\n",
        "    y_true_13[..., -1] = 1.\n",
        "    y_true_26[..., -1] = 1.\n",
        "    y_true_52[..., -1] = 1.\n",
        "\n",
        "    y_true = [y_true_13, y_true_26, y_true_52]\n",
        "\n",
        "    # [N, 1, 2]\n",
        "    box_sizes = np.expand_dims(box_sizes, 1)\n",
        "    # broadcast tricks\n",
        "    # [N, 1, 2] & [9, 2] ==> [N, 9, 2]\n",
        "    mins = np.maximum(- box_sizes / 2, - anchors / 2)\n",
        "    maxs = np.minimum(box_sizes / 2, anchors / 2)\n",
        "    # [N, 9, 2]\n",
        "    whs = maxs - mins\n",
        "\n",
        "    # [N, 9]\n",
        "    iou = (whs[:, :, 0] * whs[:, :, 1]) / (\n",
        "                box_sizes[:, :, 0] * box_sizes[:, :, 1] + anchors[:, 0] * anchors[:, 1] - whs[:, :, 0] * whs[:, :,\n",
        "                                                                                                         1] + 1e-10)\n",
        "    # [N]\n",
        "    best_match_idx = np.argmax(iou, axis=1)\n",
        "\n",
        "    ratio_dict = {1.: 8., 2.: 16., 3.: 32.}\n",
        "    for i, idx in enumerate(best_match_idx):\n",
        "        # idx: 0,1,2 ==> 2; 3,4,5 ==> 1; 6,7,8 ==> 0\n",
        "        feature_map_group = 2 - idx // 3\n",
        "        # scale ratio: 0,1,2 ==> 8; 3,4,5 ==> 16; 6,7,8 ==> 32\n",
        "        ratio = ratio_dict[np.ceil((idx + 1) / 3.)]\n",
        "        x = int(np.floor(box_centers[i, 0] / ratio))\n",
        "        y = int(np.floor(box_centers[i, 1] / ratio))\n",
        "        k = anchors_mask[feature_map_group].index(idx)\n",
        "        c = labels[i]\n",
        "        # print(feature_map_group, '|', y,x,k,c)\n",
        "\n",
        "        y_true[feature_map_group][y, x, k, :2] = box_centers[i]\n",
        "        y_true[feature_map_group][y, x, k, 2:4] = box_sizes[i]\n",
        "        y_true[feature_map_group][y, x, k, 4] = 1.\n",
        "        y_true[feature_map_group][y, x, k, 5 + c] = 1.\n",
        "        y_true[feature_map_group][y, x, k, -1] = boxes[i, -1]\n",
        "\n",
        "    return y_true_13, y_true_26, y_true_52\n",
        "\n",
        "\n",
        "def parse_data(data, class_num, img_size, anchors, is_training, letterbox_resize):\n",
        "    \n",
        "    img_idx, encoded_img, boxes, labels, _, _ = parse_tfrecord(data)\n",
        "    img = cv2.imdecode(encoded_img, cv2.IMREAD_COLOR)\n",
        "    boxes = np.concatenate((boxes, np.full(shape=(boxes.shape[0], 1), fill_value=1., dtype=np.float32)), axis=-1)\n",
        "\n",
        "    ## I erased mix-up method here\n",
        "\n",
        "    if is_training:\n",
        "        # random color distortion\n",
        "        img = random_color_distort(img)\n",
        "\n",
        "        # random expansion with prob 0.5\n",
        "        if np.random.uniform(0, 1) > 0.5:\n",
        "            img, boxes = random_expand(img, boxes, 4)\n",
        "\n",
        "        # random cropping\n",
        "        h, w, _ = img.shape\n",
        "        boxes, crop = random_crop_with_constraints(boxes, (w, h))\n",
        "        x0, y0, w, h = crop\n",
        "        img = img[y0: y0+h, x0: x0+w]\n",
        "\n",
        "        # resize with random interpolation\n",
        "        h, w, _ = img.shape\n",
        "        interp = np.random.randint(0, 5)\n",
        "        img, boxes = resize_with_bbox(img, boxes, img_size[0], img_size[1], interp=interp, letterbox=letterbox_resize)\n",
        "\n",
        "        # random horizontal flip\n",
        "        h, w, _ = img.shape\n",
        "        img, boxes = random_flip(img, boxes, px=0.5)\n",
        "    else:\n",
        "        img, boxes = resize_with_bbox(img, boxes, img_size[0], img_size[1], interp=1, letterbox=letterbox_resize)\n",
        "\n",
        "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32)\n",
        "\n",
        "    # the input of yolo_v3 should be in range 0~1\n",
        "    img = img / 255.\n",
        "\n",
        "    y_true_13, y_true_26, y_true_52 = process_box(boxes, labels, img_size, class_num, anchors)\n",
        "\n",
        "    return img_idx, img, y_true_13, y_true_26, y_true_52\n",
        "\n",
        "\n",
        "def get_batch_data(records, class_num, img_size, anchors, is_training, multi_scale=False, mix_up=False, letterbox_resize=True, interval=10):\n",
        "    global iter_cnt\n",
        "\n",
        "    # multi_scale training\n",
        "    if multi_scale and is_training:\n",
        "        random.seed(iter_cnt // interval)\n",
        "        random_img_size = [[x * 32, x * 32] for x in range(10, 20)]\n",
        "        img_size = random.sample(random_img_size, 1)[0]\n",
        "    iter_cnt += 1\n",
        "\n",
        "    img_idx_batch, img_batch, y_true_13_batch, y_true_26_batch, y_true_52_batch = [], [], [], [], []\n",
        "\n",
        "    # deleted mix up strategy\n",
        "    \n",
        "    for data in records:\n",
        "        img_idx, img, y_true_13, y_true_26, y_true_52 = parse_data(data, class_num, img_size, anchors, is_training, letterbox_resize)\n",
        "\n",
        "        img_idx_batch.append(img_idx)\n",
        "        img_batch.append(img)\n",
        "        y_true_13_batch.append(y_true_13)\n",
        "        y_true_26_batch.append(y_true_26)\n",
        "        y_true_52_batch.append(y_true_52)\n",
        "\n",
        "    img_idx_batch, img_batch, y_true_13_batch, y_true_26_batch, y_true_52_batch = np.asarray(img_idx_batch, np.int64), np.asarray(img_batch), np.asarray(y_true_13_batch), np.asarray(y_true_26_batch), np.asarray(y_true_52_batch)\n",
        "\n",
        "    return img_idx_batch, img_batch, y_true_13_batch, y_true_26_batch, y_true_52_batch"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sd9Pk3XgDqxt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "## evaluation utils\n",
        "\n",
        "from collections import Counter\n",
        "\n",
        "def calc_iou(pred_boxes, true_boxes):\n",
        "    pred_boxes = np.expand_dims(pred_boxes, -2)\n",
        "    true_boxes = np.expand_dims(true_boxes, 0)\n",
        "\n",
        "    intersect_mins = np.maximum(pred_boxes[..., :2], true_boxes[..., :2])\n",
        "    intersect_maxs = np.minimum(pred_boxes[..., 2:], true_boxes[..., 2:])\n",
        "    intersect_wh = np.maximum(intersect_maxs - intersect_mins, 0.)\n",
        "\n",
        "    intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]\n",
        "    pred_box_wh = pred_boxes[..., 2:] - pred_boxes[..., :2]\n",
        "    pred_box_area = pred_box_wh[..., 0] * pred_box_wh[..., 1]\n",
        "    true_boxes_wh = true_boxes[..., 2:] - true_boxes[..., :2]\n",
        "    true_boxes_area = true_boxes_wh[..., 0] * true_boxes_wh[..., 1]\n",
        "\n",
        "    iou = intersect_area / (pred_box_area + true_boxes_area - intersect_area + 1e-10)\n",
        "\n",
        "    return iou\n",
        "\n",
        "\n",
        "def evaluate_on_cpu(y_pred, y_true, num_classes, calc_now=True, max_boxes=50, score_thresh=0.5, iou_thresh=0.5):\n",
        "    num_images = y_true[0].shape[0]\n",
        "    true_labels_dict = {i: 0 for i in range(num_classes)}\n",
        "    pred_labels_dict = {i: 0 for i in range(num_classes)}\n",
        "    true_positive_dict = {i: 0 for i in range(num_classes)}\n",
        "\n",
        "    for i in range(num_images):\n",
        "        true_labels_list, true_boxes_list = [], []\n",
        "        for j in range(3):\n",
        "            true_probs_temp = y_true[j][i][..., 5:-1]\n",
        "            true_boxes_temp = y_true[j][i][..., 0:4]\n",
        "\n",
        "            object_mask = true_probs_temp.sum(axis=-1) > 0\n",
        "\n",
        "            true_probs_temp = true_probs_temp[object_mask]\n",
        "            true_boxes_temp = true_boxes_temp[object_mask]\n",
        "\n",
        "            true_labels_list += np.argmax(true_probs_temp, axis=-1).tolist()\n",
        "            true_boxes_list += true_boxes_temp.tolist()\n",
        "\n",
        "        if len(true_labels_list) != 0:\n",
        "            for cls, count in Counter(true_labels_list).items():\n",
        "                true_labels_dict[cls] += count\n",
        "\n",
        "        true_boxes = np.array(true_boxes_list)\n",
        "        box_centers, box_sizes = true_boxes[:, 0:2], true_boxes[:, 2:4]\n",
        "        true_boxes[:, 0:2] = box_centers - box_sizes / 2.\n",
        "        true_boxes[:, 2:4] = true_boxes[:, 0:2] + box_sizes\n",
        "\n",
        "        pred_boxes = y_pred[0][i:i + 1]\n",
        "        pred_confs = y_pred[1][i:i + 1]\n",
        "        pred_probs = y_pred[2][i:i + 1]\n",
        "\n",
        "        pred_boxes, pred_confs, pred_labels = cpu_nms(pred_boxes, pred_confs * pred_probs, num_classes, max_boxes=max_boxes, score_thresh=score_thresh, iou_thresh=iou_thresh)\n",
        "\n",
        "        pred_labels_list = [] if pred_labels is None else pred_labels.tolist()\n",
        "        if pred_labels_list == []:\n",
        "            continue\n",
        "\n",
        "        # calc iou\n",
        "        iou_matrix = calc_iou(pred_boxes, true_boxes)\n",
        "        max_iou_idx = np.argmax(iou_matrix, axis=-1)\n",
        "\n",
        "        correct_idx = []\n",
        "        correct_conf = []\n",
        "\n",
        "        for k in range(max_iou_idx.shape[0]):\n",
        "            pred_labels_dict[pred_labels_list[k]] += 1\n",
        "            match_idx = max_iou_idx[k]  # V level\n",
        "            if iou_matrix[k, match_idx] > iou_thresh and true_labels_list[match_idx] == pred_labels_list[k]:\n",
        "                if match_idx not in correct_idx:\n",
        "                    correct_idx.append(match_idx)\n",
        "                    correct_conf.append(pred_confs[k])\n",
        "                else:\n",
        "                    same_idx = correct_idx.index(match_idx)\n",
        "                    if pred_confs[k] > correct_conf[same_idx]:\n",
        "                        correct_idx.pop(same_idx)\n",
        "                        correct_conf.pop(same_idx)\n",
        "                        correct_idx.append(match_idx)\n",
        "                        correct_conf.append(pred_confs[k])\n",
        "\n",
        "        for t in correct_idx:\n",
        "            true_positive_dict[true_labels_list[t]] += 1\n",
        "\n",
        "    if calc_now:\n",
        "        # avoid divided by 0\n",
        "        recall = sum(true_positive_dict.values()) / (sum(true_labels_dict.values()) + 1e-6)\n",
        "        precision = sum(true_positive_dict.values()) / (sum(pred_labels_dict.values()) + 1e-6)\n",
        "\n",
        "        return recall, precision\n",
        "    else:\n",
        "        return true_positive_dict, true_labels_dict, pred_labels_dict\n",
        "\n",
        "\n",
        "def evaluate_on_gpu(sess, gpu_nms_op, pred_boxes_flag, pred_scores_flag, y_pred, y_true, num_classes, iou_thresh=0.5, calc_now=True):\n",
        "    num_images = y_true[0].shape[0]\n",
        "    true_labels_dict = {i: 0 for i in range(num_classes)}\n",
        "    pred_labels_dict = {i: 0 for i in range(num_classes)}\n",
        "    true_positive_dict = {i: 0 for i in range(num_classes)}\n",
        "\n",
        "    for i in range(num_images):\n",
        "        true_labels_list, true_boxes_list = [], []\n",
        "        for j in range(3):\n",
        "            true_probs_temp = y_true[j][i][..., 5:-1]\n",
        "            true_boxes_temp = y_true[j][i][..., 0:4]\n",
        "\n",
        "            object_mask = true_probs_temp.sum(axis=-1) > 0\n",
        "\n",
        "            true_probs_temp = true_probs_temp[object_mask]\n",
        "            true_boxes_temp = true_boxes_temp[object_mask]\n",
        "\n",
        "            true_labels_list += np.argmax(true_probs_temp, axis=-1).tolist()\n",
        "            true_boxes_list += true_boxes_temp.tolist()\n",
        "\n",
        "        if len(true_labels_list) != 0:\n",
        "            for cls, count in Counter(true_labels_list).items():\n",
        "                true_labels_dict[cls] += count\n",
        "\n",
        "        true_boxes = np.array(true_boxes_list)\n",
        "        box_centers, box_sizes = true_boxes[:, 0:2], true_boxes[:, 2:4]\n",
        "        true_boxes[:, 0:2] = box_centers - box_sizes / 2.\n",
        "        true_boxes[:, 2:4] = true_boxes[:, 0:2] + box_sizes\n",
        "\n",
        "        pred_boxes = y_pred[0][i:i + 1]\n",
        "        pred_confs = y_pred[1][i:i + 1]\n",
        "        pred_probs = y_pred[2][i:i + 1]\n",
        "\n",
        "        pred_boxes, pred_confs, pred_labels = sess.run(gpu_nms_op, feed_dict={pred_boxes_flag: pred_boxes, pred_scores_flag: pred_confs * pred_probs})\n",
        "\n",
        "        pred_labels_list = [] if pred_labels is None else pred_labels.tolist()\n",
        "        if pred_labels_list == []:\n",
        "            continue\n",
        "\n",
        "        # calc iou\n",
        "        iou_matrix = calc_iou(pred_boxes, true_boxes)\n",
        "        max_iou_idx = np.argmax(iou_matrix, axis=-1)\n",
        "\n",
        "        correct_idx = []\n",
        "        correct_conf = []\n",
        "        for k in range(max_iou_idx.shape[0]):\n",
        "            pred_labels_dict[pred_labels_list[k]] += 1\n",
        "            match_idx = max_iou_idx[k]  # V level\n",
        "            if iou_matrix[k, match_idx] > iou_thresh and true_labels_list[match_idx] == pred_labels_list[k]:\n",
        "                if match_idx not in correct_idx:\n",
        "                    correct_idx.append(match_idx)\n",
        "                    correct_conf.append(pred_confs[k])\n",
        "                else:\n",
        "                    same_idx = correct_idx.index(match_idx)\n",
        "                    if pred_confs[k] > correct_conf[same_idx]:\n",
        "                        correct_idx.pop(same_idx)\n",
        "                        correct_conf.pop(same_idx)\n",
        "                        correct_idx.append(match_idx)\n",
        "                        correct_conf.append(pred_confs[k])\n",
        "\n",
        "        for t in correct_idx:\n",
        "            true_positive_dict[true_labels_list[t]] += 1\n",
        "\n",
        "    if calc_now:\n",
        "        # avoid divided by 0\n",
        "        recall = sum(true_positive_dict.values()) / (sum(true_labels_dict.values()) + 1e-6)\n",
        "        precision = sum(true_positive_dict.values()) / (sum(pred_labels_dict.values()) + 1e-6)\n",
        "\n",
        "        return recall, precision\n",
        "    else:\n",
        "        return true_positive_dict, true_labels_dict, pred_labels_dict\n",
        "\n",
        "\n",
        "def get_preds_gpu(sess, gpu_nms_op, pred_boxes_flag, pred_scores_flag, image_ids, y_pred):\n",
        "    image_id = image_ids[0]\n",
        "\n",
        "    pred_boxes = y_pred[0][0:1]\n",
        "    pred_confs = y_pred[1][0:1]\n",
        "    pred_probs = y_pred[2][0:1]\n",
        "\n",
        "    boxes, scores, labels = sess.run(gpu_nms_op, feed_dict={pred_boxes_flag: pred_boxes, pred_scores_flag: pred_confs * pred_probs})\n",
        "\n",
        "    pred_content = []\n",
        "    for i in range(len(labels)):\n",
        "        x_min, y_min, x_max, y_max = boxes[i]\n",
        "        score = scores[i]\n",
        "        label = labels[i]\n",
        "        pred_content.append([image_id, x_min, y_min, x_max, y_max, score, label])\n",
        "\n",
        "    return pred_content\n",
        "\n",
        "gt_dict = {}  # key: img_id, value: gt object list\n",
        "def parse_gt_rec(gt_filename, compression_type, target_img_size, letterbox_resize=True):\n",
        "    global gt_dict\n",
        "\n",
        "    if not gt_dict:\n",
        "        new_width, new_height = target_img_size\n",
        "\n",
        "        with TFRecordIterator(gt_filename, compression_type) as reader:\n",
        "            for data in reader:\n",
        "                img_id, image, boxes, labels, ori_width, ori_height = parse_record(data)\n",
        "\n",
        "                objects = []\n",
        "                for i in range(len(labels)):\n",
        "                    x_min, y_min, x_max, y_max = boxes[i]\n",
        "                    label = labels[i]\n",
        "\n",
        "                    if letterbox_resize:\n",
        "                        resize_ratio = min(new_width / ori_width, new_height / ori_height)\n",
        "\n",
        "                        resize_w = int(resize_ratio * ori_width)\n",
        "                        resize_h = int(resize_ratio * ori_height)\n",
        "\n",
        "                        dw = int((new_width - resize_w) / 2)\n",
        "                        dh = int((new_height - resize_h) / 2)\n",
        "\n",
        "                        objects.append([x_min * resize_ratio + dw,\n",
        "                                        y_min * resize_ratio + dh,\n",
        "                                        x_max * resize_ratio + dw,\n",
        "                                        y_max * resize_ratio + dh,\n",
        "                                        label])\n",
        "                    else:\n",
        "                        objects.append([x_min * new_width / ori_width,\n",
        "                                        y_min * new_height / ori_height,\n",
        "                                        x_max * new_width / ori_width,\n",
        "                                        y_max * new_height / ori_height,\n",
        "                                        label])\n",
        "                gt_dict[img_id] = objects\n",
        "    return gt_dict\n",
        "\n",
        "\n",
        "# The following two functions are modified from FAIR's Detectron repo to calculate mAP:\n",
        "# https://github.com/facebookresearch/Detectron/blob/master/detectron/datasets/voc_eval.py\n",
        "def voc_ap(rec, prec, use_07_metric=False):\n",
        "    if use_07_metric:\n",
        "        ap = 0.\n",
        "        for t in np.arange(0., 1.1, 0.1):\n",
        "            if np.sum(rec >= t) == 0:\n",
        "                p = 0\n",
        "            else:\n",
        "                p = np.max(prec[rec >= t])\n",
        "            ap = ap + p / 11.\n",
        "    else:\n",
        "        mrec = np.concatenate(([0.], rec, [1.]))\n",
        "        mpre = np.concatenate(([0.], prec, [0.]))\n",
        "\n",
        "        for i in range(mpre.size - 1, 0, -1):\n",
        "            mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])\n",
        "\n",
        "        i = np.where(mrec[1:] != mrec[:-1])[0]\n",
        "\n",
        "        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])\n",
        "    return ap\n",
        "\n",
        "\n",
        "def voc_eval(gt_dict, val_preds, classidx, iou_thres=0.5, use_07_metric=False):\n",
        "    # 1.obtain gt: extract all gt objects for this class\n",
        "    class_recs = {}\n",
        "    npos = 0\n",
        "    for img_id in gt_dict:\n",
        "        R = [obj for obj in gt_dict[img_id] if obj[-1] == classidx]\n",
        "        bbox = np.array([x[:4] for x in R])\n",
        "        det = [False] * len(R)\n",
        "        npos += len(R)\n",
        "        class_recs[img_id] = {'bbox': bbox, 'det': det}\n",
        "\n",
        "    # 2. obtain pred results\n",
        "    pred = [x for x in val_preds if x[-1] == classidx]\n",
        "    img_ids = [x[0] for x in pred]\n",
        "    confidence = np.array([x[-2] for x in pred])\n",
        "    BB = np.array([[x[1], x[2], x[3], x[4]] for x in pred])\n",
        "\n",
        "    # 3. sort by confidence\n",
        "    sorted_ind = np.argsort(-confidence)\n",
        "    try:\n",
        "        BB = BB[sorted_ind, :]\n",
        "    except:\n",
        "        print('no box, ignore')\n",
        "        return 1e-6, 1e-6, 0, 0, 0\n",
        "    img_ids = [img_ids[x] for x in sorted_ind]\n",
        "\n",
        "    # 4. mark TPs and FPs\n",
        "    nd = len(img_ids)\n",
        "    tp = np.zeros(nd)\n",
        "    fp = np.zeros(nd)\n",
        "\n",
        "    for d in range(nd):\n",
        "        R = class_recs[img_ids[d]]\n",
        "        bb = BB[d, :]\n",
        "        ovmax = -np.Inf\n",
        "        BBGT = R['bbox']\n",
        "\n",
        "        if BBGT.size > 0:\n",
        "            ixmin = np.maximum(BBGT[:, 0], bb[0])\n",
        "            iymin = np.maximum(BBGT[:, 1], bb[1])\n",
        "            ixmax = np.minimum(BBGT[:, 2], bb[2])\n",
        "            iymax = np.minimum(BBGT[:, 3], bb[3])\n",
        "            iw = np.maximum(ixmax - ixmin + 1., 0.)\n",
        "            ih = np.maximum(iymax - iymin + 1., 0.)\n",
        "            inters = iw * ih\n",
        "\n",
        "            uni = ((bb[2] - bb[0] + 1.) * (bb[3] - bb[1] + 1.) + (BBGT[:, 2] - BBGT[:, 0] + 1.) * (\n",
        "                        BBGT[:, 3] - BBGT[:, 1] + 1.) - inters)\n",
        "\n",
        "            overlaps = inters / uni\n",
        "            ovmax = np.max(overlaps)\n",
        "            jmax = np.argmax(overlaps)\n",
        "\n",
        "        if ovmax > iou_thres:\n",
        "            # gt not matched yet\n",
        "            if not R['det'][jmax]:\n",
        "                tp[d] = 1.\n",
        "                R['det'][jmax] = 1\n",
        "            else:\n",
        "                fp[d] = 1.\n",
        "        else:\n",
        "            fp[d] = 1.\n",
        "\n",
        "    fp = np.cumsum(fp)\n",
        "    tp = np.cumsum(tp)\n",
        "    rec = tp / float(npos)\n",
        "    # avoid divide by zero in case the first detection matches a difficult\n",
        "    prec = tp / np.maximum(tp + fp, np.finfo(np.float64).eps)\n",
        "    ap = voc_ap(rec, prec, use_07_metric)\n",
        "\n",
        "    # return rec, prec, ap\n",
        "    return npos, nd, tp[-1] / float(npos), tp[-1] / float(nd), ap"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X4uQxNl0FRli",
        "colab_type": "code",
        "outputId": "c2b22c73-6195-4b80-d1b4-5ada76ef3da8",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 161
        }
      },
      "source": [
        "## model\n",
        "\n",
        "slim = tf.contrib.slim\n",
        "\n",
        "def conv2d(inputs, filters, kernel_size, strides=1):\n",
        "    def _fixed_padding(inputs, kernel_size):\n",
        "        pad_total = kernel_size - 1\n",
        "        pad_beg = pad_total // 2\n",
        "        pad_end = pad_total - pad_beg\n",
        "\n",
        "        padded_inputs = tf.pad(inputs, [[0, 0], [pad_beg, pad_end],\n",
        "                                        [pad_beg, pad_end], [0, 0]], mode='CONSTANT')\n",
        "        return padded_inputs\n",
        "    if strides > 1: \n",
        "        inputs = _fixed_padding(inputs, kernel_size)\n",
        "    inputs = slim.conv2d(inputs, filters, kernel_size, stride=strides,\n",
        "                         padding=('SAME' if strides == 1 else 'VALID'))\n",
        "    return inputs\n",
        "\n",
        "def darknet53_body(inputs):\n",
        "    def res_block(inputs, filters):\n",
        "        shortcut = inputs\n",
        "        net = conv2d(inputs, filters * 1, 1)\n",
        "        net = conv2d(net, filters * 2, 3)\n",
        "\n",
        "        net = net + shortcut\n",
        "\n",
        "        return net\n",
        "    \n",
        "    # first two conv2d layers\n",
        "    net = conv2d(inputs, 32,  3, strides=1)\n",
        "    net = conv2d(net, 64,  3, strides=2)\n",
        "\n",
        "    # res_block * 1\n",
        "    net = res_block(net, 32)\n",
        "\n",
        "    net = conv2d(net, 128, 3, strides=2)\n",
        "\n",
        "    # res_block * 2\n",
        "    for i in range(2):\n",
        "        net = res_block(net, 64)\n",
        "\n",
        "    net = conv2d(net, 256, 3, strides=2)\n",
        "\n",
        "    # res_block * 8\n",
        "    for i in range(8):\n",
        "        net = res_block(net, 128)\n",
        "\n",
        "    route_1 = net\n",
        "    net = conv2d(net, 512, 3, strides=2)\n",
        "\n",
        "    # res_block * 8\n",
        "    for i in range(8):\n",
        "        net = res_block(net, 256)\n",
        "\n",
        "    route_2 = net\n",
        "    net = conv2d(net, 1024, 3, strides=2)\n",
        "\n",
        "    # res_block * 4\n",
        "    for i in range(4):\n",
        "        net = res_block(net, 512)\n",
        "    route_3 = net\n",
        "\n",
        "    return route_1, route_2, route_3\n",
        "\n",
        "\n",
        "def yolo_block(inputs, filters):\n",
        "    net = conv2d(inputs, filters * 1, 1)\n",
        "    net = conv2d(net, filters * 2, 3)\n",
        "    net = conv2d(net, filters * 1, 1)\n",
        "    net = conv2d(net, filters * 2, 3)\n",
        "    net = conv2d(net, filters * 1, 1)\n",
        "    route = net\n",
        "    net = conv2d(net, filters * 2, 3)\n",
        "    return route, net\n",
        "\n",
        "\n",
        "def upsample_layer(inputs, out_shape):\n",
        "    new_height, new_width = out_shape[1], out_shape[2]\n",
        "    # NOTE: here height is the first\n",
        "    inputs = tf.image.resize_nearest_neighbor(inputs, (new_height, new_width), name='upsampled')\n",
        "    return inputs\n",
        "\n",
        "class yolov3(object):\n",
        "\n",
        "    def __init__(self, class_num, anchors, use_label_smooth=False, use_focal_loss=False, batch_norm_decay=0.999, weight_decay=5e-4, use_static_shape=True):\n",
        "        self.class_num = class_num\n",
        "        self.anchors = anchors\n",
        "        self.batch_norm_decay = batch_norm_decay\n",
        "        self.use_label_smooth = use_label_smooth\n",
        "        self.use_focal_loss = use_focal_loss\n",
        "        self.weight_decay = weight_decay\n",
        "        self.use_static_shape = use_static_shape\n",
        "\n",
        "    def forward(self, inputs, is_training=False, reuse=False):\n",
        "        # the input size: [height, weight] format\n",
        "        self.img_size = tf.shape(inputs)[1:3]\n",
        "        print(\"Img size:\", self.img_size)\n",
        "\t\t\n",
        "        batch_norm_params = {\n",
        "            'decay': self.batch_norm_decay,\n",
        "            'epsilon': 1e-05,\n",
        "            'scale': True,\n",
        "            'is_training': is_training,\n",
        "            'fused': None,\n",
        "        }\n",
        "\n",
        "        with slim.arg_scope([slim.conv2d, slim.batch_norm], reuse=reuse):\n",
        "            with slim.arg_scope([slim.conv2d], \n",
        "                                normalizer_fn=slim.batch_norm,\n",
        "                                normalizer_params=batch_norm_params,\n",
        "                                biases_initializer=None,\n",
        "                                activation_fn=lambda x: tf.nn.leaky_relu(x, alpha=0.1),\n",
        "                                weights_regularizer=slim.l2_regularizer(self.weight_decay)):\n",
        "\n",
        "                with tf.variable_scope('darknet53_body'):\n",
        "                    route_1, route_2, route_3 = darknet53_body(inputs)\n",
        "\n",
        "                with tf.variable_scope('yolov3_head'):\n",
        "                    inter1, net = yolo_block(route_3, 512)\n",
        "                    feature_map_1 = slim.conv2d(net, 3 * (5 + self.class_num), 1,\n",
        "                                                stride=1, normalizer_fn=None,\n",
        "                                                activation_fn=None, biases_initializer=tf.zeros_initializer())\n",
        "                    feature_map_1 = tf.identity(feature_map_1, name='feature_map_1')\n",
        "\n",
        "                    inter1 = conv2d(inter1, 256, 1)\n",
        "                    inter1 = upsample_layer(inter1, route_2.get_shape().as_list() if self.use_static_shape else tf.shape(route_2))\n",
        "                    concat1 = tf.concat([inter1, route_2], axis=3)\n",
        "\n",
        "                    inter2, net = yolo_block(concat1, 256)\n",
        "                    feature_map_2 = slim.conv2d(net, 3 * (5 + self.class_num), 1,\n",
        "                                                stride=1, normalizer_fn=None,\n",
        "                                                activation_fn=None, biases_initializer=tf.zeros_initializer())\n",
        "                    feature_map_2 = tf.identity(feature_map_2, name='feature_map_2')\n",
        "\n",
        "                    inter2 = conv2d(inter2, 128, 1)\n",
        "                    inter2 = upsample_layer(inter2, route_1.get_shape().as_list() if self.use_static_shape else tf.shape(route_1))\n",
        "                    concat2 = tf.concat([inter2, route_1], axis=3)\n",
        "\n",
        "                    _, feature_map_3 = yolo_block(concat2, 128)\n",
        "                    feature_map_3 = slim.conv2d(feature_map_3, 3 * (5 + self.class_num), 1,\n",
        "                                                stride=1, normalizer_fn=None,\n",
        "                                                activation_fn=None, biases_initializer=tf.zeros_initializer())\n",
        "                    feature_map_3 = tf.identity(feature_map_3, name='feature_map_3')\n",
        "\n",
        "            return feature_map_1, feature_map_2, feature_map_3\n",
        "\n",
        "    def reorg_layer(self, feature_map, anchors):\t\n",
        "        # size : [h, w] format\n",
        "        grid_size = feature_map.get_shape().as_list()[1:3] if self.use_static_shape else tf.shape(feature_map)[1:3]  # [13, 13]\n",
        "        ratio = tf.cast(self.img_size / grid_size, tf.float32)\n",
        "\t\t\n",
        "        # anchor : [w, h] format\n",
        "        rescaled_anchors = [(anchor[0] / ratio[1], anchor[1] / ratio[0]) for anchor in anchors]\n",
        "\n",
        "        feature_map = tf.reshape(feature_map, [-1, grid_size[0], grid_size[1], 3, 5 + self.class_num])\n",
        "\t\t\n",
        "        box_centers, box_sizes, conf_logits, prob_logits = tf.split(feature_map, [2, 2, 1, self.class_num], axis=-1)\n",
        "        box_centers = tf.nn.sigmoid(box_centers)\n",
        "\n",
        "        grid_x = tf.range(grid_size[1], dtype=tf.int32)\n",
        "        grid_y = tf.range(grid_size[0], dtype=tf.int32)\n",
        "        grid_x, grid_y = tf.meshgrid(grid_x, grid_y)\n",
        "        x_offset = tf.reshape(grid_x, (-1, 1))\n",
        "        y_offset = tf.reshape(grid_y, (-1, 1))\n",
        "        x_y_offset = tf.concat([x_offset, y_offset], axis=-1)\n",
        "\t\t\n",
        "        x_y_offset = tf.cast(tf.reshape(x_y_offset, [grid_size[0], grid_size[1], 1, 2]), tf.float32)\n",
        "\n",
        "        box_centers = box_centers + x_y_offset\n",
        "        box_centers = box_centers * ratio[::-1]\n",
        "\n",
        "        box_sizes = tf.exp(box_sizes) * rescaled_anchors\n",
        "        box_sizes = box_sizes * ratio[::-1]\n",
        "\n",
        "        boxes = tf.concat([box_centers, box_sizes], axis=-1)\n",
        "\n",
        "        return x_y_offset, boxes, conf_logits, prob_logits\n",
        "    \n",
        "    def predict(self, feature_maps):\n",
        "        feature_map_1, feature_map_2, feature_map_3 = feature_maps\n",
        "\n",
        "        feature_map_anchors = [(feature_map_1, self.anchors[6:9]),\n",
        "                               (feature_map_2, self.anchors[3:6]),\n",
        "                               (feature_map_3, self.anchors[0:3])]\n",
        "        reorg_results = [self.reorg_layer(feature_map, anchors) for (feature_map, anchors) in feature_map_anchors]\n",
        "\n",
        "        def _reshape_logit(result):\n",
        "            x_y_offset, boxes, conf_logits, prob_logits = result\n",
        "            grid_size = x_y_offset.get_shape().as_list()[:2] if self.use_static_shape else tf.shape(x_y_offset)[:2]\n",
        "            boxes = tf.reshape(boxes, [-1, grid_size[0] * grid_size[1] * 3, 4])\n",
        "            conf_logits = tf.reshape(conf_logits, [-1, grid_size[0] * grid_size[1] * 3, 1])\n",
        "            prob_logits = tf.reshape(prob_logits, [-1, grid_size[0] * grid_size[1] * 3, self.class_num])\n",
        "            return boxes, conf_logits, prob_logits\n",
        "\n",
        "        boxes_list, confs_list, probs_list = [], [], []\n",
        "\t\t\n",
        "        for result in reorg_results:\n",
        "            boxes, conf_logits, prob_logits = _reshape_logit(result)\n",
        "            confs = tf.sigmoid(conf_logits)\n",
        "            probs = tf.sigmoid(prob_logits)\n",
        "            boxes_list.append(boxes)\n",
        "            confs_list.append(confs)\n",
        "            probs_list.append(probs)\n",
        "        \n",
        "        boxes = tf.concat(boxes_list, axis=1)\n",
        "        confs = tf.concat(confs_list, axis=1)\n",
        "        probs = tf.concat(probs_list, axis=1)\n",
        "\n",
        "        center_x, center_y, width, height = tf.split(boxes, [1, 1, 1, 1], axis=-1)\n",
        "        x_min = center_x - width / 2\n",
        "        y_min = center_y - height / 2\n",
        "        x_max = center_x + width / 2\n",
        "        y_max = center_y + height / 2\n",
        "\n",
        "        boxes = tf.concat([x_min, y_min, x_max, y_max], axis=-1)\n",
        "\n",
        "        return boxes, confs, probs\n",
        "    \n",
        "    def loss_layer(self, feature_map_i, y_true, anchors):\n",
        "        grid_size = tf.shape(feature_map_i)[1:3]\n",
        "        ratio = tf.cast(self.img_size / grid_size, tf.float32)\n",
        "        # N: batch_size\n",
        "        N = tf.cast(tf.shape(feature_map_i)[0], tf.float32)\n",
        "\n",
        "        x_y_offset, pred_boxes, pred_conf_logits, pred_prob_logits = self.reorg_layer(feature_map_i, anchors)\n",
        "\n",
        "\t\t### mask\n",
        "        object_mask = y_true[..., 4:5]\n",
        "        ignore_mask = tf.TensorArray(tf.float32, size=0, dynamic_size=True)\n",
        "\t\t\n",
        "        def loop_cond(idx, ignore_mask):\n",
        "            return tf.less(idx, tf.cast(N, tf.int32))\n",
        "\t\t\t\n",
        "        def loop_body(idx, ignore_mask):\n",
        "            valid_true_boxes = tf.boolean_mask(y_true[idx, ..., 0:4], tf.cast(object_mask[idx, ..., 0], 'bool'))\n",
        "\t\t\t\n",
        "            iou = self.box_iou(pred_boxes[idx], valid_true_boxes)\t\t\t\n",
        "            best_iou = tf.reduce_max(iou, axis=-1)\n",
        "\t\t\t\n",
        "            ignore_mask_tmp = tf.cast(best_iou < 0.5, tf.float32)\n",
        "\t\t\t\n",
        "            ignore_mask = ignore_mask.write(idx, ignore_mask_tmp)\n",
        "            return idx + 1, ignore_mask\n",
        "\t\t\t\n",
        "        _, ignore_mask = tf.while_loop(cond=loop_cond, body=loop_body, loop_vars=[0, ignore_mask])\n",
        "        ignore_mask = ignore_mask.stack()\n",
        "        ignore_mask = tf.expand_dims(ignore_mask, -1)\n",
        "\n",
        "        pred_box_xy = pred_boxes[..., 0:2]\n",
        "        pred_box_wh = pred_boxes[..., 2:4]\n",
        "\n",
        "        true_xy = y_true[..., 0:2] / ratio[::-1] - x_y_offset\n",
        "        pred_xy = pred_box_xy / ratio[::-1] - x_y_offset\n",
        "\n",
        "        true_tw_th = y_true[..., 2:4] / anchors\n",
        "        pred_tw_th = pred_box_wh / anchors\n",
        "\t\t\n",
        "        true_tw_th = tf.where(condition=tf.equal(true_tw_th, 0),\n",
        "                              x=tf.ones_like(true_tw_th), y=true_tw_th)\n",
        "        pred_tw_th = tf.where(condition=tf.equal(pred_tw_th, 0),\n",
        "                              x=tf.ones_like(pred_tw_th), y=pred_tw_th)\n",
        "        true_tw_th = tf.log(tf.clip_by_value(true_tw_th, 1e-9, 1e9))\n",
        "        pred_tw_th = tf.log(tf.clip_by_value(pred_tw_th, 1e-9, 1e9))\n",
        "\n",
        "        box_loss_scale = 2. - (y_true[..., 2:3] / tf.cast(self.img_size[1], tf.float32)) * (y_true[..., 3:4] / tf.cast(self.img_size[0], tf.float32))\n",
        "\n",
        "        ### loss\n",
        "\t\t\n",
        "        mix_w = y_true[..., -1:]\n",
        "\t\t\n",
        "        xy_loss = tf.reduce_sum(tf.square(true_xy - pred_xy) * object_mask * box_loss_scale * mix_w) / N\n",
        "        wh_loss = tf.reduce_sum(tf.square(true_tw_th - pred_tw_th) * object_mask * box_loss_scale * mix_w) / N\n",
        "\n",
        "        conf_pos_mask = object_mask\n",
        "        conf_neg_mask = (1 - object_mask) * ignore_mask\n",
        "        conf_loss_pos = conf_pos_mask * tf.nn.sigmoid_cross_entropy_with_logits(labels=object_mask, logits=pred_conf_logits)\n",
        "        conf_loss_neg = conf_neg_mask * tf.nn.sigmoid_cross_entropy_with_logits(labels=object_mask, logits=pred_conf_logits)\n",
        "\t\t\n",
        "        conf_loss = conf_loss_pos + conf_loss_neg\n",
        "\n",
        "        if self.use_focal_loss:\n",
        "            alpha = 1.0\n",
        "            gamma = 2.0\n",
        "            focal_mask = alpha * tf.pow(tf.abs(object_mask - tf.sigmoid(pred_conf_logits)), gamma)\n",
        "            conf_loss *= focal_mask\n",
        "        conf_loss = tf.reduce_sum(conf_loss * mix_w) / N\n",
        "\n",
        "        if self.use_label_smooth:\n",
        "            delta = 0.01\n",
        "            label_target = (1 - delta) * y_true[..., 5:-1] + delta * 1. / self.class_num\n",
        "        else:\n",
        "            label_target = y_true[..., 5:-1]\n",
        "\t\t\t\n",
        "        class_loss = object_mask * tf.nn.sigmoid_cross_entropy_with_logits(labels=label_target, logits=pred_prob_logits) * mix_w\n",
        "        class_loss = tf.reduce_sum(class_loss) / N\n",
        "\n",
        "        return xy_loss, wh_loss, conf_loss, class_loss\n",
        "    \n",
        "\n",
        "    def box_iou(self, pred_boxes, valid_true_boxes):\n",
        "        pred_box_xy = pred_boxes[..., 0:2]\n",
        "        pred_box_wh = pred_boxes[..., 2:4]\n",
        "\n",
        "        pred_box_xy = tf.expand_dims(pred_box_xy, -2)\n",
        "        pred_box_wh = tf.expand_dims(pred_box_wh, -2)\n",
        "\n",
        "        true_box_xy = valid_true_boxes[:, 0:2]\n",
        "        true_box_wh = valid_true_boxes[:, 2:4]\n",
        "\n",
        "        intersect_mins = tf.maximum(pred_box_xy - pred_box_wh / 2.,\n",
        "                                    true_box_xy - true_box_wh / 2.)\n",
        "        intersect_maxs = tf.minimum(pred_box_xy + pred_box_wh / 2.,\n",
        "                                    true_box_xy + true_box_wh / 2.)\n",
        "        intersect_wh = tf.maximum(intersect_maxs - intersect_mins, 0.)\n",
        "\n",
        "        intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]\n",
        "        pred_box_area = pred_box_wh[..., 0] * pred_box_wh[..., 1]\n",
        "        true_box_area = true_box_wh[..., 0] * true_box_wh[..., 1]\n",
        "        true_box_area = tf.expand_dims(true_box_area, axis=0)\n",
        "\n",
        "        iou = intersect_area / (pred_box_area + true_box_area - intersect_area + 1e-10)\n",
        "\n",
        "        return iou\n",
        "\n",
        "    \n",
        "    def compute_loss(self, y_pred, y_true):\n",
        "        loss_xy, loss_wh, loss_conf, loss_class = 0., 0., 0., 0.\n",
        "        anchor_group = [self.anchors[6:9], self.anchors[3:6], self.anchors[0:3]]\n",
        "\n",
        "        for i in range(len(y_pred)):\n",
        "            result = self.loss_layer(y_pred[i], y_true[i], anchor_group[i])\n",
        "            loss_xy += result[0]\n",
        "            loss_wh += result[1]\n",
        "            loss_conf += result[2]\n",
        "            loss_class += result[3]\n",
        "        total_loss = loss_xy + loss_wh + loss_conf + loss_class\n",
        "        return [total_loss, loss_xy, loss_wh, loss_conf, loss_class]"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:\n",
            "The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
            "For more information, please see:\n",
            "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
            "  * https://github.com/tensorflow/addons\n",
            "  * https://github.com/tensorflow/io (for I/O related ops)\n",
            "If you depend on functionality not listed there, please file an issue.\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nlddq-K7AJin",
        "colab_type": "code",
        "outputId": "c5baed55-0d4e-4c65-fa7d-340b27baf8f9",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 89
        }
      },
      "source": [
        "## arguments\n",
        "\n",
        "import math\n",
        "\n",
        "\n",
        "### Some paths\n",
        "\n",
        "data_path = '/content/gdrive/My Drive/yolo/data/'\n",
        "train_file = data_path + 'train.tfrecord'  # The path of the training txt file.\n",
        "val_file = data_path + 'val.tfrecord'  # The path of the validation txt file.\n",
        "restore_path = data_path + 'darknet_weights/yolov3.ckpt'  # The path of the weights to restore.\n",
        "save_dir = '/content/gdrive/My Drive/yolo/checkpoint/'  # The directory of the weights to save.\n",
        "\n",
        "### we are not using tensorboard logs in this code\n",
        "\n",
        "log_dir = data_path + 'logs/'  # The directory to store the tensorboard log files.\n",
        "progress_log_path = data_path + 'progress.log'  # The path to record the training progress.\n",
        "\n",
        "anchor_path = data_path + 'yolo_anchors.txt'  # The path of the anchor txt file.\n",
        "class_name_path = data_path + 'classes.txt'  # The path of the class names.\n",
        "\n",
        "### Training releated numbers\n",
        "batch_size = 4\n",
        "img_size = [416, 416]  # Images will be resized to `img_size` and fed to the network, size format: [width, height]\n",
        "letterbox_resizing = True  # Whether to use the letterbox resize, i.e., keep the original aspect ratio in the resized image.\n",
        "total_epoches = 10\n",
        "train_evaluation_step = 10  # Evaluate on the training batch after some steps.\n",
        "val_evaluation_epoch = 2  # Evaluate on the whole validation dataset after some epochs. Set to None to evaluate every epoch.\n",
        "save_epoch = 5  # Save the model after some epochs.\n",
        "batch_norm_decay = 0.99  # decay in bn ops\n",
        "weight_decay = 5e-4  # l2 weight decay\n",
        "current_global_step = 0  # used when resuming training\n",
        "\n",
        "### tf.data parameters\n",
        "num_threads = 10  # Number of threads for image processing used in tf.data pipeline.\n",
        "prefetech_buffer = 5  # Prefetech_buffer used in tf.data pipeline.\n",
        "\n",
        "### Learning rate and optimizer\n",
        "optimizer_name = 'momentum'  # Chosen from [sgd, momentum, adam, rmsprop]\n",
        "save_optimizer = True  # Whether to save the optimizer parameters into the checkpoint file.\n",
        "learning_rate_init = 1e-4\n",
        "lr_type = 'piecewise'  # Chosen from [fixed, exponential, cosine_decay, cosine_decay_restart, piecewise]\n",
        "lr_decay_epoch = 5  # Epochs after which learning rate decays. Int or float. Used when chosen `exponential` and `cosine_decay_restart` lr_type.\n",
        "lr_decay_factor = 0.96  # The learning rate decay factor. Used when chosen `exponential` lr_type.\n",
        "lr_lower_bound = 1e-6  # The minimum learning rate.\n",
        "# only used in piecewise lr type\n",
        "pw_boundaries = [30, 50]  # epoch based boundaries\n",
        "pw_values = [learning_rate_init, 3e-5, 1e-5]\n",
        "\n",
        "### Load and finetune\n",
        "# Choose the parts you want to restore the weights. List form.\n",
        "# restore_include: None, restore_exclude: None  => restore the whole model\n",
        "# restore_include: None, restore_exclude: scope  => restore the whole model except `scope`\n",
        "# restore_include: scope1, restore_exclude: scope2  => if scope1 contains scope2, restore scope1 and not restore scope2 (scope1 - scope2)\n",
        "# choise 1: only restore the darknet body\n",
        "# restore_include = ['yolov3/darknet53_body']\n",
        "# restore_exclude = None\n",
        "# choise 2: restore all layers except the last 3 conv2d layers in 3 scale\n",
        "restore_include = None\n",
        "restore_exclude = ['yolov3/yolov3_head/Conv_14', 'yolov3/yolov3_head/Conv_6', 'yolov3/yolov3_head/Conv_22']\n",
        "# Choose the parts you want to finetune. List form.\n",
        "# Set to None to train the whole model.\n",
        "\n",
        "update_part = ['yolov3/yolov3_head']\n",
        "\n",
        "### other training strategies\n",
        "multi_scale_train = True  # Whether to apply multi-scale training strategy. Image size varies from [320, 320] to [640, 640] by default.\n",
        "use_label_smooth = True # Whether to use class label smoothing strategy.\n",
        "use_focal_loss = True  # Whether to apply focal loss on the conf loss.\n",
        "use_mix_up = True  # Whether to use mix up data augmentation strategy. \n",
        "use_warm_up = True  # whether to use warm up strategy to prevent from gradient exploding.\n",
        "warm_up_epoch = 2  # Warm up training epoches. Set to a larger value if gradient explodes.\n",
        "\n",
        "### some constants in validation\n",
        "# nms\n",
        "nms_threshold = 0.45  # iou threshold in nms operation\n",
        "score_threshold = 0.01  # threshold of the probability of the classes in nms operation, i.e. score = pred_confs * pred_probs. set lower for higher recall.\n",
        "nms_topk = 150  # keep at most nms_topk outputs after nms\n",
        "# mAP eval\n",
        "eval_threshold = 0.5  # the iou threshold applied in mAP evaluation\n",
        "use_voc_07_metric = False  # whether to use voc 2007 evaluation metric, i.e. the 11-point metric\n",
        "\n",
        "### parse some params\n",
        "anchors = parse_anchors(anchor_path)\n",
        "classes = read_class_names(class_name_path)\n",
        "class_num = len(classes)\n",
        "train_img_cnt = TFRecordIterator(train_file, 'GZIP').count()\n",
        "val_img_cnt = TFRecordIterator(val_file, 'GZIP').count()\n",
        "train_batch_num = int(math.ceil(float(train_img_cnt) / batch_size))\n",
        "\n",
        "lr_decay_freq = int(train_batch_num * lr_decay_epoch)\n",
        "pw_boundaries = [float(i) * train_batch_num + current_global_step for i in pw_boundaries]\n"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From <ipython-input-2-ea7f0591b13c>:7: tf_record_iterator (from tensorflow.python.lib.io.tf_record) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Use eager execution and: \n",
            "`tf.data.TFRecordDataset(path)`\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NagT2oNZFf0q",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "## train\n",
        "\n",
        "import os\n",
        "from tqdm import trange\n",
        "\n",
        "if training:\n",
        "  is_training = tf.placeholder(tf.bool, name=\"phase_train\")\n",
        "  handle_flag = tf.placeholder(tf.string, [], name='iterator_handle_flag')\n",
        "\n",
        "  pred_boxes_flag = tf.placeholder(tf.float32, [1, None, None])\n",
        "  pred_scores_flag = tf.placeholder(tf.float32, [1, None, None])\n",
        "  gpu_nms_op = gpu_nms(pred_boxes_flag, pred_scores_flag, class_num, nms_topk, score_threshold, nms_threshold)\n",
        "\n",
        "  ### tf.data pipeline\n",
        "  train_dataset = tf.data.TFRecordDataset(filenames=train_file, compression_type='GZIP')\n",
        "  train_dataset = train_dataset.shuffle(train_img_cnt)\n",
        "  train_dataset = train_dataset.batch(batch_size)\n",
        "  train_dataset = train_dataset.map(\n",
        "      lambda x: tf.py_func(get_batch_data,\n",
        "                          inp=[x, class_num, img_size, anchors, True, multi_scale_train, use_mix_up, letterbox_resizing],\n",
        "                          Tout=[tf.int64, tf.float32, tf.float32, tf.float32, tf.float32]),\n",
        "      num_parallel_calls=num_threads\n",
        "  )\n",
        "  train_dataset = train_dataset.prefetch(prefetech_buffer)\n",
        "\n",
        "  val_dataset = tf.data.TFRecordDataset(filenames=val_file, compression_type='GZIP')\n",
        "  val_dataset = val_dataset.batch(1)\n",
        "  val_dataset = val_dataset.map(\n",
        "      lambda x: tf.py_func(get_batch_data,\n",
        "                          inp=[x, class_num, img_size, anchors, False, False, False, letterbox_resizing],\n",
        "                          Tout=[tf.int64, tf.float32, tf.float32, tf.float32, tf.float32]),\n",
        "      num_parallel_calls=num_threads\n",
        "  )\n",
        "  val_dataset.prefetch(prefetech_buffer)\n",
        "\n",
        "  iterator = tf.data.Iterator.from_structure(train_dataset.output_types, train_dataset.output_shapes)\n",
        "  train_init_op = iterator.make_initializer(train_dataset)\n",
        "  val_init_op = iterator.make_initializer(val_dataset)\n",
        "\n",
        "  image_ids, image, y_true_13, y_true_26, y_true_52 = iterator.get_next()\n",
        "  y_true = [y_true_13, y_true_26, y_true_52]\n",
        "\n",
        "  image_ids.set_shape([None])\n",
        "  image.set_shape([None, None, None, 3])\n",
        "  for y in y_true:\n",
        "      y.set_shape([None, None, None, None, None])\n",
        "\n",
        "\n",
        "  ### Model definition\n",
        "  yolo_model = yolov3(class_num, anchors, use_label_smooth, use_focal_loss, batch_norm_decay, weight_decay, use_static_shape=False)\n",
        "\n",
        "  with tf.variable_scope('yolov3'):\n",
        "      pred_feature_maps = yolo_model.forward(image, is_training=is_training)\n",
        "\n",
        "  loss = yolo_model.compute_loss(pred_feature_maps, y_true)\n",
        "  y_pred = yolo_model.predict(pred_feature_maps)\n",
        "\n",
        "  l2_loss = tf.losses.get_regularization_loss()\n",
        "\n",
        "  saver_to_restore = tf.train.Saver(var_list=tf.contrib.framework.get_variables_to_restore(include=restore_include, exclude=restore_exclude))\n",
        "  update_vars = tf.contrib.framework.get_variables_to_restore(include=update_part)\n",
        "\n",
        "\n",
        "  global_step = tf.Variable(float(current_global_step), trainable=False, collections=[tf.GraphKeys.LOCAL_VARIABLES])\n",
        "  if use_warm_up:\n",
        "      learning_rate = tf.cond(tf.less(global_step, train_batch_num * warm_up_epoch), \n",
        "                              lambda: learning_rate_init * global_step / (train_batch_num * warm_up_epoch),\n",
        "                              lambda: config_learning_rate(global_step - train_batch_num * warm_up_epoch))\n",
        "  else:\n",
        "      learning_rate = config_learning_rate(global_step)\n",
        "\n",
        "  optimizer = config_optimizer(optimizer_name, learning_rate)\n",
        "\n",
        "  update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)\n",
        "\n",
        "  with tf.control_dependencies(update_ops):\n",
        "      gvs = optimizer.compute_gradients(loss[0] + l2_loss, var_list=update_vars)\n",
        "      clip_grad_var = [gv if gv[0] is None else [\n",
        "            tf.clip_by_norm(gv[0], 100.), gv[1]] for gv in gvs]\n",
        "      train_op = optimizer.apply_gradients(clip_grad_var, global_step=global_step)\n",
        "\n",
        "  if save_optimizer:\n",
        "      print('Saving optimizer parameters: ON')\n",
        "      saver_to_save = tf.train.Saver()\n",
        "      saver_best = tf.train.Saver()\n",
        "  else:\n",
        "      print('Saving optimizer parameters: OFF')\n",
        "\n",
        "\n",
        "  with tf.Session() as sess:\n",
        "      sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])\n",
        "\n",
        "      if os.path.exists(restore_path):\n",
        "          saver_to_restore.restore(sess, restore_path)\n",
        "\n",
        "      print('\\nStart training...: Total epoches =', total_epoches, '\\n')\n",
        "\n",
        "      best_mAP = -np.Inf\n",
        "\n",
        "      for epoch in range(total_epoches):\n",
        "          sess.run(train_init_op)\n",
        "          loss_total, loss_xy, loss_wh, loss_conf, loss_class = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter()\n",
        "\n",
        "          ### train part\n",
        "          for i in trange(train_batch_num):\n",
        "              _, __y_pred, __y_true, __loss, __global_step, __lr = sess.run(\n",
        "                  [train_op, y_pred, y_true, loss, global_step, learning_rate],\n",
        "                  feed_dict={is_training: True})\n",
        "\n",
        "              loss_total.update(__loss[0], len(__y_pred[0]))\n",
        "              loss_xy.update(__loss[1], len(__y_pred[0]))\n",
        "              loss_wh.update(__loss[2], len(__y_pred[0]))\n",
        "              loss_conf.update(__loss[3], len(__y_pred[0]))\n",
        "              loss_class.update(__loss[4], len(__y_pred[0]))\n",
        "\n",
        "              if __global_step % train_evaluation_step == 0 and __global_step > 0:\n",
        "                  recall, precision = evaluate_on_gpu(sess, gpu_nms_op, pred_boxes_flag, pred_scores_flag, __y_pred, __y_true, class_num, nms_threshold)\n",
        "\n",
        "                  info = \"Epoch: {}, global_step: {} | loss: total: {:.2f}, xy: {:.2f}, wh: {:.2f}, conf: {:.2f}, class: {:.2f} | \".format(\n",
        "                          epoch, int(__global_step), loss_total.average, loss_xy.average, loss_wh.average, loss_conf.average, loss_class.average)\n",
        "                  info += 'Last batch: rec: {:.3f}, prec: {:.3f} | lr: {:.5g}'.format(recall, precision, __lr)\n",
        "                  print(info)\n",
        "          \n",
        "                  if np.isnan(loss_total.average):\n",
        "                      print('****' * 10)\n",
        "                      raise ArithmeticError('Gradient exploded!')\n",
        "\n",
        "          ## train end (saving parameters)\n",
        "          if save_optimizer and epoch % save_epoch == 0 and epoch > 0:\n",
        "              if loss_total.average <= 2.:\n",
        "                  saver_to_save.save(sess, save_dir + 'model-epoch_{}_step_{}_loss_{:.4f}_lr_{:.5g}'.format(epoch, int(__global_step), loss_total.average, __lr))\n",
        "\n",
        "          ### validation part\n",
        "          if epoch % val_evaluation_epoch == 0 and epoch >= warm_up_epoch:\n",
        "              sess.run(val_init_op)\n",
        "\n",
        "              val_loss_total, val_loss_xy, val_loss_wh, val_loss_conf, val_loss_class = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter()\n",
        "\n",
        "              val_preds = []\n",
        "\n",
        "              for j in trange(val_img_cnt):\n",
        "                  __image_ids, __y_pred, __loss = sess.run([image_ids, y_pred, loss],\n",
        "                                                          feed_dict={is_training: False})\n",
        "                  pred_content = get_preds_gpu(sess, gpu_nms_op, pred_boxes_flag, pred_scores_flag, __image_ids, __y_pred)\n",
        "                  val_preds.extend(pred_content)\n",
        "                  val_loss_total.update(__loss[0])\n",
        "                  val_loss_xy.update(__loss[1])\n",
        "                  val_loss_wh.update(__loss[2])\n",
        "                  val_loss_conf.update(__loss[3])\n",
        "                  val_loss_class.update(__loss[4])\n",
        "\n",
        "              # calc mAP\n",
        "              rec_total, prec_total, ap_total = AverageMeter(), AverageMeter(), AverageMeter()\n",
        "              gt_dict = parse_gt_rec(val_file, 'GZIP', img_size, letterbox_resize)\n",
        "\n",
        "              info = '======> Epoch: {}, global_step: {}, lr: {:.6g} <======\\n'.format(epoch, __global_step, __lr)\n",
        "\n",
        "              for ii in range(class_num):\n",
        "                  npos, nd, rec, prec, ap = voc_eval(gt_dict, val_preds, ii, iou_thres=eval_threshold, use_07_metric=use_voc_07_metric)\n",
        "                  info += 'EVAL: Class {}: Recall: {:.4f}, Precision: {:.4f}, AP: {:.4f}\\n'.format(ii, rec, prec, ap)\n",
        "                  rec_total.update(rec, npos)\n",
        "                  prec_total.update(prec, nd)\n",
        "                  ap_total.update(ap, 1)\n",
        "\n",
        "              mAP = ap_total.average\n",
        "              info += 'EVAL: Recall: {:.4f}, Precison: {:.4f}, mAP: {:.4f}\\n'.format(rec_total.average, prec_total.average, mAP)\n",
        "              info += 'EVAL: loss: total: {:.2f}, xy: {:.2f}, wh: {:.2f}, conf: {:.2f}, class: {:.2f}\\n'.format(\n",
        "                  val_loss_total.average, val_loss_xy.average, val_loss_wh.average, val_loss_conf.average, val_loss_class.average)\n",
        "              print(info)\n",
        "\n",
        "              if save_optimizer and mAP > best_mAP:\n",
        "                  best_mAP = mAP\n",
        "                  saver_best.save(sess, save_dir + 'best_model_Epoch_{}_step_{}_mAP_{:.4f}_loss_{:.4f}_lr_{:.7g}'.format(\n",
        "                                    epoch, int(__global_step), best_mAP, val_loss_total.average, __lr))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HmoSmKIuOpyC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "## evaluation (test)\n",
        "\n",
        "import argparse\n",
        "\n",
        "if not training:\n",
        "\n",
        "  ### ArgumentParser\n",
        "  parser = argparse.ArgumentParser(description=\"YOLO-V3 eval procedure.\")\n",
        "\n",
        "  # paths\n",
        "  parser.add_argument(\"--eval_file\", type=str, default=\"/content/gdrive/My Drive/yolo/data/test.tfrecord\",\n",
        "                      help=\"The path of the validation or test txt file.\")\n",
        "\n",
        "  parser.add_argument(\"--restore_path\", type=str, default=\"/content/gdrive/My Drive/yolo/data/darknet_weights/yolov3.ckpt\",\n",
        "                      help=\"The path of the weights to restore.\")\n",
        "\n",
        "  parser.add_argument(\"--anchor_path\", type=str, default=\"./content/gdrive/My Drive/yolo/data/yolo_anchors.txt\",\n",
        "                      help=\"The path of the anchor txt file.\")\n",
        "\n",
        "  parser.add_argument(\"--class_name_path\", type=str, default=\"/content/gdrive/My Drive/yolo/data/classes.txt\",\n",
        "                      help=\"The path of the class names.\")\n",
        "\n",
        "  # some numbers\n",
        "  parser.add_argument(\"--img_size\", nargs='*', type=int, default=[416, 416],\n",
        "                      help=\"Resize the input image to `img_size`, size format: [width, height]\")\n",
        "\n",
        "  parser.add_argument(\"--letterbox_resize\", type=lambda x: (str(x).lower() == 'true'), default=False,\n",
        "                      help=\"Whether to use the letterbox resize, i.e., keep the original image aspect ratio.\")\n",
        "\n",
        "  parser.add_argument(\"--num_threads\", type=int, default=10,\n",
        "                      help=\"Number of threads for image processing used in tf.data pipeline.\")\n",
        "\n",
        "  parser.add_argument(\"--prefetech_buffer\", type=int, default=5,\n",
        "                      help=\"Prefetech_buffer used in tf.data pipeline.\")\n",
        "\n",
        "  parser.add_argument(\"--nms_threshold\", type=float, default=0.45,\n",
        "                      help=\"IOU threshold in nms operation.\")\n",
        "\n",
        "  parser.add_argument(\"--score_threshold\", type=float, default=0.01,\n",
        "                      help=\"Threshold of the probability of the classes in nms operation.\")\n",
        "\n",
        "  parser.add_argument(\"--nms_topk\", type=int, default=400,\n",
        "                      help=\"Keep at most nms_topk outputs after nms.\")\n",
        "\n",
        "  parser.add_argument(\"--use_voc_07_metric\", type=lambda x: (str(x).lower() == 'true'), default=False,\n",
        "                      help=\"Whether to use the voc 2007 mAP metrics.\")\n",
        "\n",
        "  args = parser.parse_args()\n",
        "\n",
        "  # args params\n",
        "  args.anchors = parse_anchors(args.anchor_path)\n",
        "  args.classes = read_class_names(args.class_name_path)\n",
        "  args.class_num = len(args.classes)\n",
        "  args.img_cnt = len(open(args.eval_file, 'r').readlines())\n",
        "\n",
        "  # setting placeholders\n",
        "  is_training = tf.placeholder(dtype=tf.bool, name=\"phase_train\")\n",
        "  handle_flag = tf.placeholder(tf.string, [], name='iterator_handle_flag')\n",
        "  pred_boxes_flag = tf.placeholder(tf.float32, [1, None, None])\n",
        "  pred_scores_flag = tf.placeholder(tf.float32, [1, None, None])\n",
        "  gpu_nms_op = gpu_nms(pred_boxes_flag, pred_scores_flag, args.class_num, args.nms_topk, args.score_threshold, args.nms_threshold)\n",
        "\n",
        "  ### tf.data pipeline\n",
        "  val_dataset = tf.data.TFRecordDataset(filenames=args.eval_file, compression_type='GZIP')\n",
        "  val_dataset = val_dataset.batch(1)\n",
        "  val_dataset = val_dataset.map(\n",
        "      lambda x: tf.py_func(get_batch_data, [x, args.class_num, args.img_size, args.anchors, False, False, False, args.letterbox_resize], [tf.int64, tf.float32, tf.float32, tf.float32, tf.float32]),\n",
        "      num_parallel_calls=args.num_threads\n",
        "  )\n",
        "  val_dataset.prefetch(args.prefetech_buffer)\n",
        "  iterator = val_dataset.make_one_shot_iterator()\n",
        "\n",
        "  image_ids, image, y_true_13, y_true_26, y_true_52 = iterator.get_next()\n",
        "  image_ids.set_shape([None])\n",
        "  y_true = [y_true_13, y_true_26, y_true_52]\n",
        "  image.set_shape([None, args.img_size[1], args.img_size[0], 3])\n",
        "  for y in y_true:\n",
        "      y.set_shape([None, None, None, None, None])\n",
        "\n",
        "  ### Model definition\n",
        "  yolo_model = yolov3(args.class_num, args.anchors)\n",
        "  with tf.variable_scope('yolov3'):\n",
        "      pred_feature_maps = yolo_model.forward(image, is_training=is_training)\n",
        "  loss = yolo_model.compute_loss(pred_feature_maps, y_true)\n",
        "  y_pred = yolo_model.predict(pred_feature_maps)\n",
        "\n",
        "  saver_to_restore = tf.train.Saver()\n",
        "\n",
        "\n",
        "  with tf.Session() as sess:\n",
        "      sess.run([tf.global_variables_initializer()])\n",
        "      if os.path.exists(args.restore_path):\n",
        "          saver_to_restore.restore(sess, args.restore_path)\n",
        "      else:\n",
        "          raise ValueError('there is no model to evaluate. You should move/create the checkpoint file to restore path')\n",
        "\n",
        "      print('\\nStart evaluation...\\n')\n",
        "\n",
        "      val_loss_total, val_loss_xy, val_loss_wh, val_loss_conf, val_loss_class = \\\n",
        "          AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter()\n",
        "      val_preds = []\n",
        "\n",
        "      for j in trange(args.img_cnt):\n",
        "          __image_ids, __y_pred, __loss = sess.run([image_ids, y_pred, loss], feed_dict={is_training: False})\n",
        "          pred_content = get_preds_gpu(sess, gpu_nms_op, pred_boxes_flag, pred_scores_flag, __image_ids, __y_pred)\n",
        "\n",
        "          val_preds.extend(pred_content)\n",
        "          val_loss_total.update(__loss[0])\n",
        "          val_loss_xy.update(__loss[1])\n",
        "          val_loss_wh.update(__loss[2])\n",
        "          val_loss_conf.update(__loss[3])\n",
        "          val_loss_class.update(__loss[4])\n",
        "\n",
        "      rec_total, prec_total, ap_total = AverageMeter(), AverageMeter(), AverageMeter()\n",
        "      gt_dict = parse_gt_rec(args.eval_file, 'GZIP', args.img_size, args.letterbox_resize)\n",
        "      print('mAP eval:')\n",
        "      for ii in range(args.class_num):\n",
        "          npos, nd, rec, prec, ap = voc_eval(gt_dict, val_preds, ii, iou_thres=0.5, use_07_metric=args.use_voc_07_metric)\n",
        "          rec_total.update(rec, npos)\n",
        "          prec_total.update(prec, nd)\n",
        "          ap_total.update(ap, 1)\n",
        "          print('Class {}: Recall: {:.4f}, Precision: {:.4f}, AP: {:.4f}'.format(ii, rec, prec, ap))\n",
        "\n",
        "      mAP = ap_total.average\n",
        "      print('final mAP: {:.4f}'.format(mAP))\n",
        "      print(\"recall: {:.3f}, precision: {:.3f}\".format(rec_total.average, prec_total.average))\n",
        "      print(\"total_loss: {:.3f}, loss_xy: {:.3f}, loss_wh: {:.3f}, loss_conf: {:.3f}, loss_class: {:.3f}\".format(\n",
        "          val_loss_total.average, val_loss_xy.average, val_loss_wh.average, val_loss_conf.average, val_loss_class.average\n",
        "      ))"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}