本文作者Zhreshold，原文載于其知乎主頁(yè)，雷鋒網(wǎng)獲其授權(quán)發(fā)布。

本文先為大家介紹目前流行的目標(biāo)檢測(cè)算法SSD (Single-Shot MultiBox Object Detection)和實(shí)驗(yàn)過(guò)程中的數(shù)據(jù)集。訓(xùn)練、測(cè)試過(guò)程及結(jié)果參見(jiàn)《從零開(kāi)始碼一個(gè)皮卡丘檢測(cè)器-CNN目標(biāo)檢測(cè)入門教程(下)》

目標(biāo)檢測(cè)通俗的來(lái)說(shuō)是為了找到圖像或者視頻里的所有目標(biāo)物體。在下面這張圖中，兩狗一貓的位置，包括它們所屬的類（狗／貓），需要被正確的檢測(cè)到。

所以和圖像分類不同的地方在于，目標(biāo)檢測(cè)需要找到盡量多的目標(biāo)物體，而且要準(zhǔn)確的定位物體的位置，一般用矩形框來(lái)表示。

在接下來(lái)的章節(jié)里，我們先介紹一個(gè)流行的目標(biāo)檢測(cè)算法，SSD (Single-Shot MultiBox Object Detection).

友情提示：本章節(jié)特別長(zhǎng)，千萬(wàn)不要在蹲坑的時(shí)候點(diǎn)開(kāi)。
本文中涉及MXNet 0.11最新的發(fā)布的gluon接口，參考MXNet 0.11發(fā)布，加入動(dòng)態(tài)圖接口Gluon，還有兩位CMU教授的親筆教程（https://zhuanlan.zhihu.com/p/28648399）

SSD:  Single Shot MultiBox Detector

顧名思義，算法的核心是用卷積神經(jīng)網(wǎng)絡(luò)一次前向推導(dǎo)求出大量多尺度（幾百到幾千）的方框來(lái)表示目標(biāo)檢測(cè)的結(jié)果。網(wǎng)絡(luò)的結(jié)構(gòu)用下圖表示。

跟所有的圖像相關(guān)的網(wǎng)絡(luò)一樣，我們需要一個(gè)主干網(wǎng)絡(luò)來(lái)提取特征，同時(shí)也是作為第一個(gè)預(yù)測(cè)特征層。網(wǎng)絡(luò)在當(dāng)前層產(chǎn)生大量的預(yù)設(shè)框，和與之對(duì)應(yīng)的每個(gè)方框的分類概率（背景，貓，狗等等）以及真正的物體和預(yù)設(shè)框的偏移量。在完成當(dāng)前層的預(yù)測(cè)后，我們會(huì)下采樣當(dāng)前特征層，作為新的預(yù)測(cè)層，重新產(chǎn)生新的預(yù)設(shè)框，分類概率，偏移量。這個(gè)過(guò)程往往會(huì)重復(fù)好幾次，直到預(yù)測(cè)特征層到達(dá)全局尺度（ )。

接下來(lái)我們用例子解釋每個(gè)細(xì)節(jié)實(shí)現(xiàn)。

預(yù)設(shè)框 Default anchor boxes

預(yù)設(shè)框的形狀和大小可以由參數(shù)控制，我們往往設(shè)置一堆預(yù)設(shè)框，以期望任意圖像上的物體都能有一個(gè)預(yù)設(shè)框能大致重合，由于每個(gè)預(yù)設(shè)框需要對(duì)應(yīng)的預(yù)測(cè)網(wǎng)絡(luò)預(yù)測(cè)值，所以希望對(duì)于每個(gè)物體都有100%重合的預(yù)設(shè)框是不現(xiàn)實(shí)的，可能會(huì)需要幾十萬(wàn)甚至幾百萬(wàn)的預(yù)設(shè)框，但是采樣的預(yù)設(shè)框越多，重合概率越好，用幾千到上萬(wàn)個(gè)預(yù)設(shè)框基本能實(shí)現(xiàn)略大于70%的最好重合率，同時(shí)保證了檢測(cè)的速度。

為了保證重合覆蓋率，對(duì)于每個(gè)特征層上的像素點(diǎn)，我們用不同的大小和長(zhǎng)寬比來(lái)采樣預(yù)設(shè)框。 假設(shè)在某個(gè)特定的特征層（ ），每個(gè)預(yù)設(shè)框的中心點(diǎn)就是特征像素點(diǎn)的中心，然后我們用如下的公式采樣預(yù)設(shè)框：

• 預(yù)設(shè)框的中心為特征像素點(diǎn)的中心

• 對(duì)于長(zhǎng)寬比 ,  尺寸 , 生成的預(yù)設(shè)框大小 。

• 對(duì)于長(zhǎng)寬比 同時(shí) , 生成的預(yù)設(shè)框大小 , 其中 是第一個(gè)預(yù)設(shè)的尺寸。

這里例子里，我們用事先實(shí)現(xiàn)的層MultiBoxPrior產(chǎn)生預(yù)設(shè)框，輸入 個(gè)預(yù)設(shè)尺寸，和 個(gè)預(yù)設(shè)的長(zhǎng)寬比，輸出為 個(gè)方框而不是 個(gè)。 當(dāng)然我們完全可以使用其他的算法產(chǎn)生不同的預(yù)設(shè)框，但是實(shí)踐中我們發(fā)現(xiàn)上述的方法覆蓋率和相應(yīng)需要的預(yù)設(shè)框數(shù)量比較合適。

import mxnet as mx

from mxnet import nd

from mxnet.contrib.ndarray import MultiBoxPrior

n=40

#  輸入形狀: batch x channel x height x weight

x=nd.random_uniform(shape=(1,3,n,n))

y=MultiBoxPrior(x,sizes=[.5,.25,.1],ratios=[1,2,.5])

# 取位于 (20,20) 像素點(diǎn)的第一個(gè)預(yù)設(shè)框

# 格式為 (x_min, y_min, x_max, y_max)

boxes=y.reshape((n,n,-1,4))

print('The first anchor box at row 21, column 21:',boxes[20,20,0,:])

The first anchor box at row 21, column 21:
[ 0.26249999  0.26249999  0.76249999  0.76249999]
<NDArray 4 @cpu(0)>

看著數(shù)字不夠直觀的話，我們把框框畫出來(lái)。取最中心像素的所有預(yù)設(shè)框，畫在圖上的話，我們看到已經(jīng)可以覆蓋幾種尺寸和位置的物體了。把所有位置的組合起來(lái)，就是相當(dāng)可觀的預(yù)設(shè)框集合了。

import matplotlib.pyplot as plt

def box_to_rect(box,color,linewidth=3):

    """convert an anchor box to a matplotlib rectangle"""

    box=box.asnumpy()

    returnplt.Rectangle(

        (box[0],box[1]),(box[2]-box[0]),(box[3]-box[1]),

        fill=False,edgecolor=color,linewidth=linewidth)

colors=['blue','green','red','black','magenta']

plt.imshow(nd.ones((n,n,3)).asnumpy())

anchors=boxes[20,20,:,:]

for i in range(anchors.shape[0]):

    plt.gca().add_patch(box_to_rect(anchors[i,:]*n,colors[i]))

plt.show()

分類預(yù)測(cè) Predict classes

這個(gè)部分的目標(biāo)很簡(jiǎn)單，就是預(yù)測(cè)每個(gè)預(yù)設(shè)框?qū)?yīng)的分類。我們用 ， Padding (填充)   的卷積來(lái)預(yù)測(cè)分類概率，這樣可以做到卷積后空間上的形狀不會(huì)變化，而且由于卷積的特點(diǎn)，每個(gè)卷積核會(huì)掃過(guò)每個(gè)預(yù)測(cè)特征層的所有像素點(diǎn)，得到 個(gè)預(yù)測(cè)值， 對(duì)應(yīng)了空間上所有的像素點(diǎn)，每個(gè)通道對(duì)應(yīng)特定的預(yù)設(shè)框 。假設(shè)有10個(gè)正類，每個(gè)像素5個(gè)預(yù)設(shè)框，那么我們就需要 個(gè)通道。 具體的來(lái)說(shuō)，對(duì)于每個(gè)像素第 個(gè)預(yù)設(shè)框：

• 通道 的值對(duì)應(yīng)背景（非物體）的得分

• 通道 對(duì)應(yīng)了第 類的得分

from mxnet.gluon import nn

def class_predictor(num_anchors,num_classes):

    """return a layer to predict classes"""

    returnnn.Conv2D(num_anchors*(num_classes+1),3,padding=1)

cls_pred=class_predictor(5,10)

cls_pred.initialize()

x=nd.zeros((2,3,20,20))

print('Class prediction',cls_pred(x).shape)

Class prediction (2, 55, 20, 20)

預(yù)測(cè)預(yù)設(shè)框偏移 Predict anchor boxes

為了找到物體準(zhǔn)確的位置，光靠預(yù)設(shè)框本身是不行的，我們還需要預(yù)測(cè)偏移量以便把真正的物體框出來(lái)。

假設(shè)    是某個(gè)預(yù)設(shè)框， 是目標(biāo)物體的真實(shí)矩形框，我們需要預(yù)測(cè)的偏移為 ，全都是長(zhǎng)度為4的向量， 我們求得偏移

• 

• 

• 

• 

所有的偏移量都除以預(yù)設(shè)框的長(zhǎng)或?qū)捠菫榱烁玫氖諗俊?/p>

類似分類概率，我們同樣用 填充 的卷積來(lái)預(yù)測(cè)偏移。這次不同的是，對(duì)于每個(gè)預(yù)設(shè)框，我們只需要 個(gè)通道來(lái)預(yù)測(cè)偏移量， 一共需要 個(gè)通道，第    個(gè)預(yù)設(shè)框?qū)?yīng)的偏移量存在通道 到通道    之間。

def box_predictor(num_anchors):

    """return a layer to predict delta locations"""

    returnnn.Conv2D(num_anchors*4,3,padding=1)

box_pred=box_predictor(10)

box_pred.initialize()

x=nd.zeros((2,3,20,20))

print('Box prediction',box_pred(x).shape)

Box prediction (2, 40, 20, 20)

下采樣特征層 Down-sample features

每次我們下采樣特征層到一半的長(zhǎng)寬，用Pooling(池化)操作就可以輕松的做到，當(dāng)然也可以用stride(步長(zhǎng))為2的卷積直接得到。在下采樣之前，我們會(huì)希望增加幾層卷積層作為緩沖，防止特征值對(duì)應(yīng)多尺度帶來(lái)的混亂，同時(shí)又能增加網(wǎng)絡(luò)的深度，得到更好的抽象。

def down_sample(num_filters):

    """stack two Conv-BatchNorm-Relu blocks and then a pooling layer

    to halve the feature size"""

    out=nn.HybridSequential()

    for_inrange(2):

        out.add(nn.Conv2D(num_filters,3,strides=1,padding=1))

        out.add(nn.BatchNorm(in_channels=num_filters))

        out.add(nn.Activation('relu'))

    out.add(nn.MaxPool2D(2))

    return out

blk=down_sample(10)

blk.initialize()

x=nd.zeros((2,3,20,20))

print('Before',x.shape,'after',blk(x).shape)

Before (2, 3, 20, 20) after (2, 10, 10, 10)

整合多個(gè)特征層預(yù)測(cè)值 Manage predictions from multiple layers

SSD算法的一個(gè)關(guān)鍵點(diǎn)在于它用到了多尺度的特征層來(lái)預(yù)測(cè)不同大小的物體。相對(duì)來(lái)說(shuō)，淺層的特征層的空間尺度更大，越到網(wǎng)絡(luò)的深層，空間尺度越小，最后我們往往下采樣直到 ，用來(lái)預(yù)測(cè)全圖大小的物體。所以每個(gè)特征層產(chǎn)生的預(yù)設(shè)框，分類概率，框偏移量需要被整合起來(lái)統(tǒng)一在全圖與真實(shí)的物體比較。 為了做到一一對(duì)應(yīng)，我們統(tǒng)一把所有的預(yù)設(shè)框， 分類概率，框偏移量 平鋪再連接。得到的是按順序排列但是攤平的所有預(yù)測(cè)值和預(yù)設(shè)框。

# 隨便創(chuàng)建一個(gè)大小為 20x20的預(yù)測(cè)層

feat1=nd.zeros((2,8,20,20))

print('Feature map 1',feat1.shape)

cls_pred1=class_predictor(5,10)

cls_pred1.initialize()

y1=cls_pred1(feat1)

print('Class prediction for feature map 1',y1.shape)

# 下采樣

ds=down_sample(16)

ds.initialize()

feat2=ds(feat1)

print('Feature map 2',feat2.shape)

cls_pred2=class_predictor(3,10)

cls_pred2.initialize()

y2=cls_pred2(feat2)

print('Class prediction for feature map 2',y2.shape)

Feature map 1 (2, 8, 20, 20)
Class prediction for feature map 1 (2, 55, 20, 20)
Feature map 2 (2, 16, 10, 10)
Class prediction for feature map 2 (2, 33, 10, 10)

def flatten_prediction(pred):

    return nd.flatten(nd.transpose(pred,axes=(0,2,3,1)))

def concat_predictions(preds):

    return nd.concat(*preds,dim=1)

flat_y1=flatten_prediction(y1)

print('Flatten class prediction 1',flat_y1.shape)

flat_y2=flatten_prediction(y2)

print('Flatten class prediction 2',flat_y2.shape)

print('Concat class predictions',concat_predictions([flat_y1,flat_y2]).shape)

Flatten class prediction 1 (2, 22000)
Flatten class prediction 2 (2, 3300)
Concat class predictions (2, 25300)

我們總是確保在 上連接，以免打亂一一對(duì)應(yīng)的關(guān)系。

主干網(wǎng)絡(luò) Body network

主干網(wǎng)絡(luò)用來(lái)從原始圖像輸入提取特征。 一般來(lái)說(shuō)我們會(huì)用預(yù)先訓(xùn)練好的用于分類的高性能網(wǎng)絡(luò)（VGG, ResNet等）來(lái)提取特征。

在這里我們就簡(jiǎn)單地堆疊幾層卷積和下采樣層作為主干網(wǎng)絡(luò)的演示。

from mxnet import gluon

def body():

    """return the body network"""

    out=nn.HybridSequential()

    for nfilters in [16,32,64]:

        out.add(down_sample(nfilters))

    return out

bnet=body()

bnet.initialize()

x=nd.zeros((2,3,256,256))

print('Body network',[y.shape for y in bnet(x)])

Body network [(64, 32, 32), (64, 32, 32)]

設(shè)計(jì)一個(gè)簡(jiǎn)單的SSD示意網(wǎng)絡(luò) Create a toy SSD model

我們這里介紹一個(gè)示意用的簡(jiǎn)單SSD網(wǎng)絡(luò)，出于速度的考量，輸入圖像尺寸定為 。

def toy_ssd_model(num_anchors,num_classes):

    """return SSD modules"""

    downsamples=nn.Sequential()

    class_preds=nn.Sequential()

    box_preds=nn.Sequential()

    downsamples.add(down_sample(128))

    downsamples.add(down_sample(128))

    downsamples.add(down_sample(128))

    for scale in range(5):

        class_preds.add(class_predictor(num_anchors,num_classes))

        box_preds.add(box_predictor(num_anchors))

    return body(),downsamples,class_preds,box_preds

print(toy_ssd_model(5,2))

(HybridSequential(
(0): HybridSequential(
(0): Conv2D(16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm(fix_gamma=False, axis=1, momentum=0.9, eps=1e-05, in_channels=16)
(2): Activation(relu)
(3): Conv2D(16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(4): BatchNorm(fix_gamma=False, axis=1, momentum=0.9, eps=1e-05, in_channels=16)
(5): Activation(relu)
(6): MaxPool2D(size=(2, 2), stride=(2, 2), padding=(0, 0), ceil_mode=False)
)
(1): HybridSequential(
(0): Conv2D(32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm(fix_gamma=False, axis=1, momentum=0.9, eps=1e-05, in_channels=32)
(2): Activation(relu)
(3): Conv2D(32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(4): BatchNorm(fix_gamma=False, axis=1, momentum=0.9, eps=1e-05, in_channels=32)
(5): Activation(relu)
(6): MaxPool2D(size=(2, 2), stride=(2, 2), padding=(0, 0), ceil_mode=False)
 )
(2): HybridSequential(
(0): Conv2D(64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm(fix_gamma=False, axis=1, momentum=0.9, eps=1e-05, in_channels=64)
(2): Activation(relu)
(3): Conv2D(64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(4): BatchNorm(fix_gamma=False, axis=1, momentum=0.9, eps=1e-05, in_channels=64)
(5): Activation(relu)
(6): MaxPool2D(size=(2, 2), stride=(2, 2), padding=(0, 0), ceil_mode=False)
)
), Sequential(
(0): HybridSequential(
(0): Conv2D(128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm(fix_gamma=False, axis=1, momentum=0.9, eps=1e-05, in_channels=128)
(2): Activation(relu)
(3): Conv2D(128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(4): BatchNorm(fix_gamma=False, axis=1, momentum=0.9, eps=1e-05, in_channels=128)
(5): Activation(relu)
(6): MaxPool2D(size=(2, 2), stride=(2, 2), padding=(0, 0), ceil_mode=False)
)
(1): HybridSequential(
(0): Conv2D(128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm(fix_gamma=False, axis=1, momentum=0.9, eps=1e-05, in_channels=128)
(2): Activation(relu)
(3): Conv2D(128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(4): BatchNorm(fix_gamma=False, axis=1, momentum=0.9, eps=1e-05, in_channels=128)
(5): Activation(relu)
(6): MaxPool2D(size=(2, 2), stride=(2, 2), padding=(0, 0), ceil_mode=False)
)
(2): HybridSequential(
(0): Conv2D(128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm(fix_gamma=False, axis=1, momentum=0.9, eps=1e-05, in_channels=128)
(2): Activation(relu)
(3): Conv2D(128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(4): BatchNorm(fix_gamma=False, axis=1, momentum=0.9, eps=1e-05, in_channels=128)
(5): Activation(relu)
(6): MaxPool2D(size=(2, 2), stride=(2, 2), padding=(0, 0), ceil_mode=False)
)
), Sequential(
(0): Conv2D(15, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): Conv2D(15, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(2): Conv2D(15, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): Conv2D(15, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(4): Conv2D(15, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
), Sequential(
(0): Conv2D(20, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): Conv2D(20, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(2): Conv2D(20, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): Conv2D(20, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(4): Conv2D(20, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
))

網(wǎng)絡(luò)前向推導(dǎo) Forward

既然我們已經(jīng)設(shè)計(jì)完網(wǎng)絡(luò)結(jié)構(gòu)了，接下來(lái)可以定義網(wǎng)絡(luò)前向推導(dǎo)的步驟。

首先得到主干網(wǎng)絡(luò)的輸出，然后對(duì)于每一個(gè)特征預(yù)測(cè)層，推導(dǎo)當(dāng)前層的預(yù)設(shè)框，分類概率和偏移量。最后我們把這些輸入攤平，連接，作為網(wǎng)絡(luò)的輸出。

打包收工 Put all things together

from mxnet import gluon

class ToySSD(gluon.Block):
    def __init__(self, num_classes, **kwargs):
        super(ToySSD, self).__init__(**kwargs)
        # 5個(gè)預(yù)測(cè)層，每層負(fù)責(zé)的預(yù)設(shè)框尺寸不同，由小到大，符合網(wǎng)絡(luò)的形狀
        self.anchor_sizes = [[.2, .272], [.37, .447], [.54, .619], [.71, .79], [.88, .961]]
        # 每層的預(yù)設(shè)框都用 1，2，0.5作為長(zhǎng)寬比候選
        self.anchor_ratios = [[1, 2, .5]] * 5
        self.num_classes = num_classes

        with self.name_scope():
            self.body, self.downsamples, self.class_preds, self.box_preds = toy_ssd_model(4, num_classes)

    def forward(self, x):
        default_anchors, predicted_classes, predicted_boxes = toy_ssd_forward(x, self.body, self.downsamples,
            self.class_preds, self.box_preds, self.anchor_sizes, self.anchor_ratios)
         # 把從每個(gè)預(yù)測(cè)層輸入的結(jié)果攤平并連接，以確保一一對(duì)應(yīng)
         anchors = concat_predictions(default_anchors)
         box_preds = concat_predictions(predicted_boxes)
         class_preds = concat_predictions(predicted_classes)
         # 改變下形狀，為了更方便地計(jì)算softmax
         class_preds = nd.reshape(class_preds, shape=(0, -1, self.num_classes + 1))

         return anchors, class_preds, box_preds

網(wǎng)絡(luò)輸出示意 Outputs of ToySSD

# 新建一個(gè)2個(gè)正類的SSD網(wǎng)絡(luò)

net = ToySSD(2)

net.initialize()

x = nd.zeros((1, 3, 256, 256))

default_anchors, class_predictions, box_predictions = net(x)

print('Outputs:', 'anchors', default_anchors.shape, 'class prediction', class_predictions.shape, 'box prediction', box_predictions.shape)

Outputs: anchors (1, 5444, 4) class prediction (1, 5444, 3) box prediction (1, 21776)

數(shù)據(jù)集 Dataset

聊了半天怎么構(gòu)建一個(gè)虛無(wú)的網(wǎng)絡(luò)，接下來(lái)看看真正有意思的東西。

我們用3D建模批量生成了一個(gè)皮卡丘的數(shù)據(jù)集，產(chǎn)生了1000張圖片作為這個(gè)展示用的訓(xùn)練集。這個(gè)數(shù)據(jù)集里面，皮神會(huì)以各種角度，各種姿勢(shì)出現(xiàn)在各種背景圖中，就像Pokemon Go里增強(qiáng)現(xiàn)實(shí)那樣炫酷。

因?yàn)槭巧傻臄?shù)據(jù)集，我們自然可以得到每只皮神的真實(shí)坐標(biāo)和大小，用來(lái)作為訓(xùn)練的真實(shí)標(biāo)記。

下載數(shù)據(jù)集 Download dataset

下載提前準(zhǔn)備好的數(shù)據(jù)集并驗(yàn)證

from mxnet.test_utils import download

import os.path as osp

def verified(file_path, sha1hash):
    import hashlib
    sha1 = hashlib.sha1()
    with open(file_path, 'rb') as f:
        while True:
            data = f.read(1048576)
            if not data:
                break
            sha1.update(data)
    matched = sha1.hexdigest() == sha1hash
    if not matched:
        print('Found hash mismatch in file {}, possibly due to incomplete download.'.format(file_path))
    return matched

url_format = 'https://apache-mxnet.s3-accelerate.amazonaws.com/gluon/datasets/pikachu/{}'

hashes = {'train.rec': 'e6bcb6ffba1ac04ff8a9b1115e650af56ee969c8',
         'train.idx': 'dcf7318b2602c06428b9988470c731621716c393',
         'val.rec': 'd6c33f799b4d058e82f2cb5bd9a976f69d72d520'}

for k, v in hashes.items():
    fname = 'pikachu_' + k
    target = osp.join('data', fname)
    url = url_format.format(k)
    if not osp.exists(target) or not verified(target, v):
        print('Downloading', target, url)
        download(url, fname=fname, dirname='data', overwrite=True)

加載數(shù)據(jù) Load dataset

加載數(shù)據(jù)可以用mxnet.image.ImageDetIter，同時(shí)還提供了大量數(shù)據(jù)增強(qiáng)的選項(xiàng)，比如翻轉(zhuǎn)，隨機(jī)截取等等。

DataBatch: data shapes: [(32, 3, 256, 256)] label shapes: [(32, 1, 5)]

示意圖 Illustration

加載的訓(xùn)練數(shù)據(jù)還可以顯示出來(lái)看看到底是怎么樣的。

import numpy as np

img = batch.data[0][0].asnumpy()  # 取第一批數(shù)據(jù)中的第一張，轉(zhuǎn)成numpy

img = img.transpose((1, 2, 0))  # 交換下通道的順序

img += np.array([123, 117, 104])

img = img.astype(np.uint8)  # 圖片應(yīng)該用0-255的范圍

# 在圖上畫出真實(shí)標(biāo)簽的方框

for label in batch.label[0][0].asnumpy():
    if label[0] < 0:
        break
    print(label)
    xmin, ymin, xmax, ymax = [int(x * data_shape) for x in label[1:5]]
    rect = plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, edgecolor=(1, 0, 0), linewidth=3)
    plt.gca().add_patch(rect)

plt.imshow(img)

plt.show()

[ 0. 0.75724518  0.34316057  0.93332517  0.70017999]

雷峰網(wǎng)版權(quán)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見(jiàn)轉(zhuǎn)載須知。