Детектор объектов GluonCV - вывод на пакеты изображений - PullRequest
1 голос
/ 22 апреля 2020

Я запускаю обнаружение объектов в потоке веб-камеры или видеофайлах, используя Python в Ubuntu18.04. До сих пор я выполнял вывод покадрово, используя этот код:

def main():

    ctx = mx.gpu(0)

    # Load pretrained model
    net = gcv.model_zoo.get_model('ssd_512_mobilenet1.0_coco', pretrained=True)
    net.hybridize()

    # Load the webcam handler
    cap = cv2.VideoCapture(0)


    count_frame = 0


    while(True):
        print(f"Frame: {count_frame}")

        # Load frame from the camera
        ret, frame = cap.read()
        # print(type(frame))

        if (cv2.waitKey(25) & 0xFF == ord('q')) or (ret == False):
            cv2.destroyAllWindows()
            cap.release()
            print("Done!!!")
            break

        # Image pre-processing
        frame = mx.nd.array(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).astype('uint8')
        frame_nd, frame_np = gcv.data.transforms.presets.ssd.transform_test(frame, short=512, max_size=700)

        # Run frame through network
        frame_nd = frame_nd.as_in_context(ctx)
        class_IDs, scores, bounding_boxes = net(frame_nd)


        # Display result with cv
        img = gcv.utils.viz.cv_plot_bbox(frame_np, bounding_boxes[0], scores[0], class_IDs[0], thresh=0.3, class_names=net.classes)
        gcv.utils.viz.cv_plot_image(img)

        count_frame += 1



    cv2.destroyAllWindows()
    cap.release()



if __name__ == "__main__":
    main()


Я хотел попробовать альтернативную версию, где я не выполняю покадровое обнаружение , но в пакетах кадров.
Я пробовал следующим образом:

  • создать пустой список в начале;
  • добавить каждый кадр (после предварительной обработки изображения) к list;
  • после N кадров (скажем, N = 50) преобразовать список в массив mx.nd.array;
  • передать указанный массив в модель;

Итак, с этим кодом:

def main():

    ctx = mx.gpu(0)

    # Load a pretrained model
    net = gcv.model_zoo.get_model('ssd_512_mobilenet1.0_coco', pretrained=True)
    net.hybridize()

    # Load the webcam handler
    cap = cv2.VideoCapture(0)

    count_frame = 0

    batch = []

    while(True):
        print(f"Frame: {count_frame}")

        # Load frame from the camera
        ret, frame = cap.read()

        if (cv2.waitKey(25) & 0xFF == ord('q')) or (ret == False):
            cv2.destroyAllWindows()
            cap.release()
            print("Done!!!")
            break

        # Image pre-processing
        frame = mx.nd.array(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).astype('uint8')
        frame_nd, frame_np = gcv.data.transforms.presets.ssd.transform_test(frame, short=512, max_size=700)
        batch.append(frame_nd)


        if (count_frame>0) and (count_frame%50 == 0):
            print("\tStarting detection.")
            batch_nd = mx.nd.array(batch)
            batch_nd = batch_nd.as_in_context(ctx)
            class_IDs, scores, bounding_boxes = net(batch)
            print("\tDetection performed.")
        count_frame += 1



    cv2.destroyAllWindows()
    cap.release()




if __name__ == "__main__":
    main()

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

batch_nd = mx.nd.array(batch)

для справки, это это вывод:

Frame: 0
Frame: 1
Frame: 2
Frame: 3
Frame: 4
Frame: 5
Frame: 6
Frame: 7
Frame: 8
Frame: 9
Frame: 10
Frame: 11
Frame: 12
Frame: 13
Frame: 14
Frame: 15
Frame: 16
Frame: 17
Frame: 18
Frame: 19
Frame: 20
    Starting detection.

какая-либо подсказка, что я делаю не так? Есть ли лучший способ отправить партии кадров на модель?

Добро пожаловать на сайт PullRequest, где вы можете задавать вопросы и получать ответы от других членов сообщества.
...