GluonCV, обнаружение объектов с предварительно обученной моделью из модели ZOO, обнаруживает только один класс - PullRequest
Новая
       99

GluonCV, обнаружение объектов с предварительно обученной моделью из модели ZOO, обнаруживает только один класс

0 голосов
/ 25 февраля 2020

Мне было интересно, как я могу использовать одну из предварительно обученных моделей из модели ZOO GluonCV для обнаружения только одного определенного класса c объектов.

Например, если я использую 'ssd_512_resnet50_v1_coco', я хотите обнаружить только один из 80 классов COCO на изображениях (скажем, «человек»). Я пытался нарезать массив в результате логического вывода: 

class_IDs, scores, bounding_boxes = net(x)

, но проблема в том, что логические маски не поддерживаются, поэтому, если я пытаюсь выбрать: 

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

it не работает (внутреннее условие возвращает массив 0 и 1 вместо 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)

Но это это как-то связано ... могу ли я что-нибудь еще сделать?

...