GluonCV,使用模型ZOO中的预训练模型进行对象检测,仅检测一类

时间:2020-02-25 10:20:05

标签: object-detection mxnet

我想知道如何使用GluonCV模型ZOO中的一种预训练模型来仅检测一种特定的对象。

例如,如果我使用的是“ ssd_512_resnet50_v1_coco”,那么我只想检测图像中80个COCO类之一(例如“人”)。 我一直在尝试对由推理得出的数组进行切片:

class_IDs, scores, bounding_boxes = net(x)

但是问题是不支持布尔掩码,所以如果我尝试选择:

class_IDs[0][class_IDs[0]==0]

它不起作用(内部条件返回的是0和1s数组,而不是True和False)。

我设法通过使用以下代码来做到这一点:

    class_IDs, scores, bounding_boxes = net(x)

    selected_class_ID = []
    selected_scores = []
    selected_bbox = []

    for ID, score, box in zip(class_IDs[0].asnumpy(), scores[0].asnumpy(), bounding_boxes[0].asnumpy()):
        # using ID== 0 because is the classID for "person" in COCO
        if (ID==0) & (score>0.45):
            selected_class_ID.append(ID)
            selected_scores.append(score)
            selected_bbox.append(box)
    selected_class_ID = nd.array(selected_class_ID)
    selected_scores = nd.array(selected_scores)
    selected_bbox = nd.array(selected_bbox)

    ax = utils.viz.plot_bbox(img, selected_bbox, selected_scores,
                                     selected_class_ID, thresh=.4, class_names=net.classes)

但这有点牵连...我还能做些什么吗?

0 个答案:

没有答案