Tensorflow2.0 實現 YOLOv3(四):utils.py

文章說明

本系列文章旨在對 Github 上 malin9402 提供的代碼進行說明,在這篇文章中,我們會對 YOLOv3 項目中的 utils.py 文件進行說明。

如果只是想運行 Github 上的代碼,可以參考對 YOLOv3 代碼的說明一文。

導入需要的庫

import cv2
import random
import colorsys
import numpy as np
from core.config import cfg

image_preprocess

輸入:

  • image:要處理的原始圖片;
  • target_size:處理後的圖片尺寸;
  • gt_boxes:是否加入預測框。

步驟:

  • 得到目標圖像的高(ih)和寬(iw)以及原始圖像的高(h)和寬(w);
  • 得到縮放係數(scale),即 iw/w 和 ih/h 的最小值;
  • 將原始圖像 resize 成(scale x h, scale x w);
  • 將缺少的像素用 128 填充;
  • 歸一化;
  • 如果需要加入預測框則畫在圖上。
def image_preprocess(image, target_size, gt_boxes=None):

    ih, iw    = target_size  # ih=416, iw=416
    h,  w, _  = image.shape  # h=900, w=1352

    scale = min(iw/w, ih/h)  # scale=0.3077
    nw, nh  = int(scale * w), int(scale * h)  # nw=416, nh=276
    image_resized = cv2.resize(image, (nw, nh)  # image_resized.shape=(276, 416)

    image_paded = np.full(shape=[ih, iw, 3], fill_value=128.0)  # image_paded.shape=(416, 416)

    dw, dh = (iw - nw) // 2, (ih-nh) // 2  # dw=0, dh=70
    image_paded[dh:nh+dh, dw:nw+dw, :] = image_resized
    image_paded = image_paded / 255.

    if gt_boxes is None:
        return image_paded

    else:
        gt_boxes[:, [0, 2]] = gt_boxes[:, [0, 2]] * scale + dw
        gt_boxes[:, [1, 3]] = gt_boxes[:, [1, 3]] * scale + dh
        return image_paded, gt_boxes

在上面的代碼中,cv2.resize 的高放在後面,寬放在前面,所以是 (nw, nh) 而不是 (nh, nw)。

load_weights

load_weights 函數的作用是將預訓練好的權值文件 yolov3.weights 加載到模型中。

首先,我們需要知道模型中一共有多少個卷積層:

count = 0
for layer in model.layers:
    if layer.name[0] == 'c':
        count += 1

得到 count=75,說明 model 中一共有 75 個卷積層,另外,通過打印所有層的名稱,我們得知 conv2d_58,conv2d_66 和 conv2d_74 這三層卷積層後面是不跟 BN 層的(但是這三層的卷積核上有閾值)。

然後,分別設置卷積層和 BN 層的權值,這個 model 就可以直接使用了。

def load_weights(model, weights_file):
    wf = open(weights_file, 'rb')
    major, minor, revision, seen, _ = np.fromfile(wf, dtype=np.int32, count=5)

    j = 0
    for i in range(75):  # 模型中一共有75個卷積層
        conv_layer_name = 'conv2d_%d' %i if i > 0 else 'conv2d'
        bn_layer_name = 'batch_normalization_%d' %j if j > 0 else 'batch_normalization'

        conv_layer = model.get_layer(conv_layer_name)
        filters = conv_layer.filters
        k_size = conv_layer.kernel_size[0]
        in_dim = conv_layer.input_shape[-1]

        if i not in [58, 66, 74]:
            # darknet weights: [beta, gamma, mean, variance]
            bn_weights = np.fromfile(wf, dtype=np.float32, count=4 * filters)
            # tf weights: [gamma, beta, mean, variance]
            bn_weights = bn_weights.reshape((4, filters))[[1, 0, 2, 3]] # 第一行和第二行對調
            bn_layer = model.get_layer(bn_layer_name)
            j += 1
        else:
            conv_bias = np.fromfile(wf, dtype=np.float32, count=filters)

        # darknet shape (out_dim, in_dim, height, width)
        conv_shape = (filters, in_dim, k_size, k_size)
        conv_weights = np.fromfile(wf, dtype=np.float32, count=np.product(conv_shape))
        # tf shape (height, width, in_dim, out_dim)
        conv_weights = conv_weights.reshape(conv_shape).transpose([2, 3, 1, 0])

        if i not in [58, 66, 74]:
            conv_layer.set_weights([conv_weights])
            bn_layer.set_weights(bn_weights)
        else:
            conv_layer.set_weights([conv_weights, conv_bias])

    assert len(wf.read()) == 0, 'failed to read all data'  # len(wf.read()) == 0 時不會報錯
    wf.close()

在以上代碼中有一些函數需要被說明:

  • model.get_layer(conv_layer_name):得到 model 中的名稱爲 conv_layer_name 的層。
  • conv_layer.filters:conv_layer 層的卷積核個數。
  • conv_layer.kernel_size:conv_layer 層的卷積核大小。
  • conv_layer.input_shape:conv_layer 層的輸入的形狀。
  • np.fromfile(wf, dtype, count):讀取 wf 文件中 count 個值,類似於一個迭代器,每次運行得到的值會往下順延。(wf 文件中的前5個數字似乎不是權值,給模型賦權值時沒有用到)
  • conv_layer.set_weights([conv_weights]):將 conv_weights 權值賦給模型的 conv_layer 層。

postprocess_boxes

postprocess_boxes 函數用來確定預測框的信息,即預測框位置、置信概率(判斷其中有檢測物體的概率)和類別。

輸入:

  • pred_bbox:圖片經過 YOLOv3 網絡後得到的所有檢測框信息的整合,因爲經過 YOLOv3 網絡後得到三個特徵圖中對應的檢測框,其形狀分別爲 (1, 52, 52, 3, 15),(1, 26, 26, 3, 15),(1, 13, 13, 3, 15),所以 pred_bbox 的形狀爲 (13x13x3+26x26x3+52x52x3, 15)。
  • org_img_shape:原始圖片的尺寸大小。
  • input_size:輸入網絡時的圖片大小。
  • score_threshold:分數閾值,我們只留下分數大於分數閾值的檢測框。

流程:

  • 因爲在通過 YOLOv3 網絡後我們得到的檢測框信息是寬度、高度和中心點座標,所以首先要將它們轉化成檢測框左上角和右下角的座標,即 (x, y, w, h) \rightarrow (xmin, ymin, xmax, ymax)。
  • 現在我們得到的預測框信息是在 feature map 上的,所以要將它們轉化爲在原圖上的信息,即 (xmin, ymin, xmax, ymax) \rightarrow (xmin_org, ymin_org, xmax_org, ymax_org)。
  • 如果某檢測框超出原始圖片邊界,即整個檢測框在圖像外面,則將這個檢測框的四個角的座標都設爲 0,那麼它的面積也是 0。
  • 計算所有檢測框的面積,如果面積爲 0,表示此檢測框超出原始圖片邊界,則棄之。
  • 計算每個檢測框的分數(類別概率*置信概率),如果某檢測框分數小於分數閾值,則棄之。

輸出:所有符合要求的檢測框在原始圖像上的信息,格式爲 [預測框的數量,預測框位置+分數+類別],形狀爲 [-1, 6]。

def postprocess_boxes(pred_bbox, org_img_shape, input_size, score_threshold):
    valid_scale = [0, np.inf]
    pred_bbox = np.array(pred_bbox)
    pred_xywh = pred_bbox[:, 0:4]
    pred_conf = pred_bbox[:, 4]
    pred_prob = pred_bbox[:, 5:]

    # # 流程1:轉換檢測框在 feature map 上的信息:(x, y, w, h) --> (xmin, ymin, xmax, ymax)
    pred_coor = np.concatenate([pred_xywh[:, :2] - pred_xywh[:, 2:] * 0.5,
                                pred_xywh[:, :2] + pred_xywh[:, 2:] * 0.5], axis=-1)

    # # 流程2:求解預測框在原圖上的位置:(xmin, ymin, xmax, ymax) -> (xmin_org, ymin_org, xmax_org, ymax_org)
    org_h, org_w = org_img_shape
    resize_ratio = min(input_size / org_w, input_size / org_h)  # 尺寸轉換率
    # dw 和 dh 中有一個是 0 (這取決於原始圖片是寬長還是高長)
    # 假設 dh 不是 0,那麼它代表着乘以尺寸轉換率後的高與 input_size 之間距離的一半
    dw = (input_size - resize_ratio * org_w) / 2
    dh = (input_size - resize_ratio * org_h) / 2
    # 減去 dw 和 dh 是爲了去除填充的部分
    pred_coor[:, 0::2] = 1.0 * (pred_coor[:, 0::2] - dw) / resize_ratio # 轉換xmin和xmax
    pred_coor[:, 1::2] = 1.0 * (pred_coor[:, 1::2] - dh) / resize_ratio # 轉換ymin和ymax

    # # 流程3:超出圖片邊界的預測框的座標變爲0
    # 下面將每個檢測框座標與邊界座標進行比較
    # 對xmin和ymin來說,如果它們中有小於0的,就用0替換它;
    # 對xmax和ymax來說,如果它們中有大於原圖右下角座標的,就用右下角座標替換它。
    pred_coor = np.concatenate([np.maximum(pred_coor[:, :2], [0, 0]),
                                np.minimum(pred_coor[:, 2:], [org_w - 1, org_h - 1])],
                                axis=-1)
    # 假設有檢測框的座標爲 [-3, -2, -1, -2],那麼經過上一步後會變成 [0, 0, -1, -2],所以 invalid_mask 將被標記爲 True
    invalid_mask = np.logical_or((pred_coor[:, 0] > pred_coor[:, 2]), (pred_coor[:, 1] > pred_coor[:, 3]))
    pred_coor[invalid_mask] = 0
   
    # # 流程4:計算每個預測框的面積,找出面積爲0的,也就是找出超出邊界的檢測框的索引
    bboxes_scale = np.sqrt(np.multiply.reduce(pred_coor[:, 2:4] - pred_coor[:, 0:2], axis=-1))  # 計算面積
    scale_mask = np.logical_and((valid_scale[0] < bboxes_scale), (bboxes_scale < valid_scale[1]))  # 符合條件的索引的列表

    # # 流程5:計算每個預測框的分數,確定高於分數閾值的預測框索引
    # 分數 = 類別概率*置信概率  置信表示該預測框內有目標的概率,類別概率指的是概率最大下所對應的類別下的概率
    classes = np.argmax(pred_prob, axis=-1)
    scores = pred_conf * pred_prob[np.arange(len(pred_coor)), classes]
    score_mask = scores > score_threshold
    mask = np.logical_and(scale_mask, score_mask)

    coors, scores, classes = pred_coor[mask], scores[mask], classes[mask]

    # np.concatenate 將 coors, scores, classes 三列拼接在一起
    return np.concatenate([coors, scores[:, np.newaxis], classes[:, np.newaxis]], axis=-1)

nms

nms 函數用來清除冗餘的預測框,也就是極大值抑制處理。

輸入:

  • bboxes:原圖上的檢測框信息;
  • iou_threshold:IOU 閾值,如果某檢測框與同一類別下得分最高檢測框的 IOU 值大於此閾值,我們即認爲此檢測框是冗餘的,應棄之;
  • method:極大值抑制的兩種方法,‘nms’ 和 ‘soft-nms’;
  • sigma:用於 ‘soft-nms’ 方法的參數。

流程(以 ‘nms’ 方法爲例):

  • 確定每個檢測框檢測到的類別;

  • 對每個類別來說:

    • 找出相同類別的預測框;

    • 根據 IOU 閾值清除指定類別下冗餘的預測框,其流程爲:

      • 取出此類別下所有預測框中得分最高的一個,並將這個預測框跟其他的同類別的預測框進行 IOU 計算;
      • 將 IOU 值大於閾值的預測框視爲與剛取出的得分最高的預測框表示了同一個檢測物,故去掉;
      • 重複以上操作,直到在此類別下取不到其他的預測框爲止。

如果使用 ‘soft-nms’ 方法,其原理可以參考論文代碼

輸出:去除了冗餘檢測框的檢測框集合。

def nms(bboxes, iou_threshold, sigma=0.3, method='nms'):
    # 確定預測框的類別
    classes_in_img = list(set(bboxes[:, 5])) # 如果有5類,那麼classes_in_img應該是[0, 1, 2, 3, 4]

    best_bboxes = []  # 存放最終需要的檢測框
    for cls in classes_in_img: # 對於每個類別來說:
        # 1. 找出相同類別的預測框
        cls_mask = (bboxes[:, 5] == cls)
        cls_bboxes = bboxes[cls_mask]
        # 2. 清除指定類別下冗餘的預測框
        while len(cls_bboxes) > 0:
            max_ind = np.argmax(cls_bboxes[:, 4])  # 指定類別下預測框的最大分數的索引
            best_bbox = cls_bboxes[max_ind]  # 指定類別下擁有最大分數的預測框
            best_bboxes.append(best_bbox)   # 存儲該預測框
            cls_bboxes = np.concatenate([cls_bboxes[: max_ind], cls_bboxes[max_ind + 1:]])  # 將其他的預測框重新組合
            # 計算該預測框與其他預測框的iou
            iou = bboxes_iou(best_bbox[np.newaxis, :4], cls_bboxes[:, :4])
            # 確定符合條件的預測框
            # 針對每個預測框建立數值標籤weight,符合條件的數值爲1,不符合條件爲0
            weight = np.ones((len(iou),), dtype=np.float32)
            assert method in ['nms', 'soft-nms']
            # 'nms'表示Iou大於閾值的檢測框被刪除,'soft-nms'表示Iou經過計算小於0的檢測框被刪除
            if method == 'nms':
                iou_mask = iou > iou_threshold
                weight[iou_mask] = 0.0
            if method == 'soft-nms':
                weight = np.exp(-(1.0 * iou ** 2 / sigma))
            cls_bboxes[:, 4] = cls_bboxes[:, 4] * weight
            score_mask = cls_bboxes[:, 4] > 0.
            cls_bboxes = cls_bboxes[score_mask]

    return best_bboxes

bboxes_iou

bboxes_iou 函數被用來計算兩個檢測框之間的 IOU 值。

輸入:兩個(或一個對多個)檢測框的座標信息(左上角+右下角)。

IOU 值其實就是兩個框的交集面積比上它們的並集面積,這個值越大,代表這兩個框的位置越接近。用下面這個圖片表示:
在這裏插入圖片描述

def bboxes_iou(boxes1, boxes2):

    boxes1 = np.array(boxes1)  # 第一個檢測框的座標數據
    boxes2 = np.array(boxes2)  # 第二個檢測框的座標數據

    boxes1_area = (boxes1[..., 2] - boxes1[..., 0]) * (boxes1[..., 3] - boxes1[..., 1])  # 第一個檢測框的面積
    boxes2_area = (boxes2[..., 2] - boxes2[..., 0]) * (boxes2[..., 3] - boxes2[..., 1])  # 第二個檢測框的面積

    left_up = np.maximum(boxes1[..., :2], boxes2[..., :2])  # 對上圖來說,left_up=[xmin2, ymin2]
    right_down = np.minimum(boxes1[..., 2:], boxes2[..., 2:])  # 對上圖來說,right_down=[xmax1, ymax1]

    inter_section = np.maximum(right_down - left_up, 0.0)  # 交集區域
    inter_area = inter_section[..., 0] * inter_section[..., 1]  # 交集面積
    union_area = boxes1_area + boxes2_area - inter_area  # 並集面積
    ious = np.maximum(1.0 * inter_area / union_area, np.finfo(np.float32).eps)

    return ious

draw_bbox

draw_bbox 函數被用來在圖片上繪製檢測框。

輸入:

  • image:被繪製檢測框的圖片;
  • bboxes:檢測框,其格式爲 [x_min, y_min, x_max, y_max, probability, cls_id];
  • classes:是一個字典,包含着所有種類的索引,由 read_class_names 函數給出;
  • show_label:是否將種類寫在檢測框上。
def draw_bbox(image, bboxes, classes=read_class_names(cfg.YOLO.CLASSES), show_label=True):

    num_classes = len(classes)
    image_h, image_w, _ = image.shape
    # 爲每個類別設置顏色(當然自己設置顏色也行)
    hsv_tuples = [(1.0 * x / num_classes, 1., 1.) for x in range(num_classes)]
    colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
    colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors))

    random.seed(0)
    random.shuffle(colors)
    random.seed(None)

    for i, bbox in enumerate(bboxes):
        coor = np.array(bbox[:4], dtype=np.int32)
        fontScale = 0.5
        score = bbox[4]
        class_ind = int(bbox[5])
        # 檢測框顏色
        bbox_color = colors[class_ind]
        # 檢測框厚度
        bbox_thick = int(0.6 * (image_h + image_w) / 600)
        # 檢測框的左上和右下座標
        c1, c2 = (coor[0], coor[1]), (coor[2], coor[3])
        # 繪製檢測框
        cv2.rectangle(image, c1, c2, bbox_color, bbox_thick)

        # 如果要在框上增加類別信息
        if show_label:
            bbox_mess = '%s: %.2f' % (classes[class_ind], score)
            # 計算文本字符串的高度與寬度
            t_size = cv2.getTextSize(bbox_mess, 0, fontScale, thickness=bbox_thick//2)[0]
            # 繪製標註信息的框
            cv2.rectangle(image, c1, (c1[0] + t_size[0], c1[1] - t_size[1] - 3), bbox_color, -1)
            # 在圖片上增加標註信息,cv2.putText 的格式爲 [圖像,文字內容, 座標 ,字體,大小,顏色,字體厚度]
            cv2.putText(image, bbox_mess, (c1[0], c1[1]-2), cv2.FONT_HERSHEY_SIMPLEX,
                        fontScale, (0, 0, 0), bbox_thick//2, lineType=cv2.LINE_AA)
    return image

read_class_names

read_class_names 函數將 yymnist.names 文件中的類別索引讀取出來並放在一個字典中。

def read_class_names(class_file_name):
    names = {}
    with open(class_file_name, 'r') as data:
        for ID, name in enumerate(data):
            names[ID] = name.strip('\n')
    return names

將 “./data/classes/yymnist.names” 輸入將返回:

{0: '0',
 1: '1',
 2: '2',
 3: '3',
 4: '4',
 5: '5',
 6: '6',
 7: '7',
 8: '8',
 9: '9'}

get_anchors

get_anchors 函數的作用是將 basline_anchors.txt 中的先驗框讀取出來。

def get_anchors(anchors_path):
    '''loads the anchors from a file'''
    with open(anchors_path) as f:
        anchors = f.readline()
    anchors = np.array(anchors.split(','), dtype=np.float32)
    return anchors.reshape(3, 3, 2)

將 “./data/anchors/basline_anchors.txt” 輸入後得到:

array([[[ 1.25   ,  1.625  ],
        [ 2.     ,  3.75   ],
        [ 4.125  ,  2.875  ]],

       [[ 1.875  ,  3.8125 ],
        [ 3.875  ,  2.8125 ],
        [ 3.6875 ,  7.4375 ]],

       [[ 3.625  ,  2.8125 ],
        [ 4.875  ,  6.1875 ],
        [11.65625, 10.1875 ]]], dtype=float32)

完整代碼

import cv2
import random
import colorsys
import numpy as np
from core.config import cfg

def load_weights(model, weights_file):
    """
    I agree that this code is very ugly, but I don’t know any better way of doing it.
    """
    wf = open(weights_file, 'rb')
    major, minor, revision, seen, _ = np.fromfile(wf, dtype=np.int32, count=5)

    j = 0
    for i in range(75):
        conv_layer_name = 'conv2d_%d' %i if i > 0 else 'conv2d'
        bn_layer_name = 'batch_normalization_%d' %j if j > 0 else 'batch_normalization'

        conv_layer = model.get_layer(conv_layer_name)
        filters = conv_layer.filters
        k_size = conv_layer.kernel_size[0]
        in_dim = conv_layer.input_shape[-1]

        if i not in [58, 66, 74]:
            # darknet weights: [beta, gamma, mean, variance]
            bn_weights = np.fromfile(wf, dtype=np.float32, count=4 * filters)
            # tf weights: [gamma, beta, mean, variance]
            bn_weights = bn_weights.reshape((4, filters))[[1, 0, 2, 3]]
            bn_layer = model.get_layer(bn_layer_name)
            j += 1
        else:
            conv_bias = np.fromfile(wf, dtype=np.float32, count=filters)

        # darknet shape (out_dim, in_dim, height, width)
        conv_shape = (filters, in_dim, k_size, k_size)
        conv_weights = np.fromfile(wf, dtype=np.float32, count=np.product(conv_shape))
        # tf shape (height, width, in_dim, out_dim)
        conv_weights = conv_weights.reshape(conv_shape).transpose([2, 3, 1, 0])

        if i not in [58, 66, 74]:
            conv_layer.set_weights([conv_weights])
            bn_layer.set_weights(bn_weights)
        else:
            conv_layer.set_weights([conv_weights, conv_bias])

    assert len(wf.read()) == 0, 'failed to read all data'
    wf.close()


def read_class_names(class_file_name):
    '''loads class name from a file'''
    names = {}
    with open(class_file_name, 'r') as data:
        for ID, name in enumerate(data):
            names[ID] = name.strip('\n')
    return names


def get_anchors(anchors_path):
    '''loads the anchors from a file'''
    with open(anchors_path) as f:
        anchors = f.readline()
    anchors = np.array(anchors.split(','), dtype=np.float32)
    return anchors.reshape(3, 3, 2)


def image_preporcess(image, target_size, gt_boxes=None):

    ih, iw    = target_size
    h,  w, _  = image.shape

    scale = min(iw/w, ih/h)
    nw, nh  = int(scale * w), int(scale * h)
    image_resized = cv2.resize(image, (nw, nh))

    image_paded = np.full(shape=[ih, iw, 3], fill_value=128.0)
    dw, dh = (iw - nw) // 2, (ih-nh) // 2
    image_paded[dh:nh+dh, dw:nw+dw, :] = image_resized
    image_paded = image_paded / 255.

    if gt_boxes is None:
        return image_paded

    else:
        gt_boxes[:, [0, 2]] = gt_boxes[:, [0, 2]] * scale + dw
        gt_boxes[:, [1, 3]] = gt_boxes[:, [1, 3]] * scale + dh
        return image_paded, gt_boxes


def draw_bbox(image, bboxes, classes=read_class_names(cfg.YOLO.CLASSES), show_label=True):
    """
    bboxes: [x_min, y_min, x_max, y_max, probability, cls_id] format coordinates.
    """

    num_classes = len(classes)
    image_h, image_w, _ = image.shape
    hsv_tuples = [(1.0 * x / num_classes, 1., 1.) for x in range(num_classes)]
    colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
    colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors))

    random.seed(0)
    random.shuffle(colors)
    random.seed(None)

    for i, bbox in enumerate(bboxes):
        coor = np.array(bbox[:4], dtype=np.int32)
        fontScale = 0.5
        score = bbox[4]
        class_ind = int(bbox[5])
        bbox_color = colors[class_ind]
        bbox_thick = int(0.6 * (image_h + image_w) / 600)
        c1, c2 = (coor[0], coor[1]), (coor[2], coor[3])
        cv2.rectangle(image, c1, c2, bbox_color, bbox_thick)

        if show_label:
            bbox_mess = '%s: %.2f' % (classes[class_ind], score)
            t_size = cv2.getTextSize(bbox_mess, 0, fontScale, thickness=bbox_thick//2)[0]
            cv2.rectangle(image, c1, (c1[0] + t_size[0], c1[1] - t_size[1] - 3), bbox_color, -1)  # filled

            cv2.putText(image, bbox_mess, (c1[0], c1[1]-2), cv2.FONT_HERSHEY_SIMPLEX,
                        fontScale, (0, 0, 0), bbox_thick//2, lineType=cv2.LINE_AA)

    return image



def bboxes_iou(boxes1, boxes2):

    boxes1 = np.array(boxes1)
    boxes2 = np.array(boxes2)

    boxes1_area = (boxes1[..., 2] - boxes1[..., 0]) * (boxes1[..., 3] - boxes1[..., 1])
    boxes2_area = (boxes2[..., 2] - boxes2[..., 0]) * (boxes2[..., 3] - boxes2[..., 1])

    left_up       = np.maximum(boxes1[..., :2], boxes2[..., :2])
    right_down    = np.minimum(boxes1[..., 2:], boxes2[..., 2:])

    inter_section = np.maximum(right_down - left_up, 0.0)
    inter_area    = inter_section[..., 0] * inter_section[..., 1]
    union_area    = boxes1_area + boxes2_area - inter_area
    ious          = np.maximum(1.0 * inter_area / union_area, np.finfo(np.float32).eps)

    return ious


def nms(bboxes, iou_threshold, sigma=0.3, method='nms'):
    """
    :param bboxes: (xmin, ymin, xmax, ymax, score, class)
    Note: soft-nms, https://arxiv.org/pdf/1704.04503.pdf
          https://github.com/bharatsingh430/soft-nms
    """
    classes_in_img = list(set(bboxes[:, 5]))
    best_bboxes = []

    for cls in classes_in_img:
        cls_mask = (bboxes[:, 5] == cls)
        cls_bboxes = bboxes[cls_mask]

        while len(cls_bboxes) > 0:
            max_ind = np.argmax(cls_bboxes[:, 4])
            best_bbox = cls_bboxes[max_ind]
            best_bboxes.append(best_bbox)
            cls_bboxes = np.concatenate([cls_bboxes[: max_ind], cls_bboxes[max_ind + 1:]])
            iou = bboxes_iou(best_bbox[np.newaxis, :4], cls_bboxes[:, :4])
            weight = np.ones((len(iou),), dtype=np.float32)

            assert method in ['nms', 'soft-nms']

            if method == 'nms':
                iou_mask = iou > iou_threshold
                weight[iou_mask] = 0.0

            if method == 'soft-nms':
                weight = np.exp(-(1.0 * iou ** 2 / sigma))

            cls_bboxes[:, 4] = cls_bboxes[:, 4] * weight
            score_mask = cls_bboxes[:, 4] > 0.
            cls_bboxes = cls_bboxes[score_mask]

    return best_bboxes


def postprocess_boxes(pred_bbox, org_img_shape, input_size, score_threshold):

    valid_scale=[0, np.inf]
    pred_bbox = np.array(pred_bbox)

    pred_xywh = pred_bbox[:, 0:4]
    pred_conf = pred_bbox[:, 4]
    pred_prob = pred_bbox[:, 5:]

    # # (1) (x, y, w, h) --> (xmin, ymin, xmax, ymax)
    pred_coor = np.concatenate([pred_xywh[:, :2] - pred_xywh[:, 2:] * 0.5,
                                pred_xywh[:, :2] + pred_xywh[:, 2:] * 0.5], axis=-1)
    # # (2) (xmin, ymin, xmax, ymax) -> (xmin_org, ymin_org, xmax_org, ymax_org)
    org_h, org_w = org_img_shape
    resize_ratio = min(input_size / org_w, input_size / org_h)

    dw = (input_size - resize_ratio * org_w) / 2
    dh = (input_size - resize_ratio * org_h) / 2

    pred_coor[:, 0::2] = 1.0 * (pred_coor[:, 0::2] - dw) / resize_ratio
    pred_coor[:, 1::2] = 1.0 * (pred_coor[:, 1::2] - dh) / resize_ratio

    # # (3) clip some boxes those are out of range
    pred_coor = np.concatenate([np.maximum(pred_coor[:, :2], [0, 0]),
                                np.minimum(pred_coor[:, 2:], [org_w - 1, org_h - 1])], axis=-1)
    invalid_mask = np.logical_or((pred_coor[:, 0] > pred_coor[:, 2]), (pred_coor[:, 1] > pred_coor[:, 3]))
    pred_coor[invalid_mask] = 0

    # # (4) discard some invalid boxes
    bboxes_scale = np.sqrt(np.multiply.reduce(pred_coor[:, 2:4] - pred_coor[:, 0:2], axis=-1))
    scale_mask = np.logical_and((valid_scale[0] < bboxes_scale), (bboxes_scale < valid_scale[1]))

    # # (5) discard some boxes with low scores
    classes = np.argmax(pred_prob, axis=-1)
    scores = pred_conf * pred_prob[np.arange(len(pred_coor)), classes]
    score_mask = scores > score_threshold
    mask = np.logical_and(scale_mask, score_mask)
    coors, scores, classes = pred_coor[mask], scores[mask], classes[mask]

    return np.concatenate([coors, scores[:, np.newaxis], classes[:, np.newaxis]], axis=-1)
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章