GluonCV - Квантованная оценка позы / модели обнаружения объектов

Я использую сценарий (реализованный с помощью MXNet GluonCV) для оценки позы в видеопотоке с камеры, и я хотел попробовать улучшить скорость вывода. В настоящее время я использую:

  • модель SSD, предварительно обученная на COCO (ssd_512_mobilenet1.0_coco), для обнаружения людей в кадре;
  • simple_pose_resnet18_v1b для оценки позы, исключающей обнаружение из предыдущей модели;

И это код, который я использую.

      import mxnet as mx
from gluoncv.data import mscoco
from gluoncv.model_zoo import get_model
from gluoncv.data.transforms.pose import detector_to_simple_pose, heatmap_to_coord
from gluoncv.utils.viz import cv_plot_image, cv_plot_keypoints
from mxnet.contrib.quantization import *

def main():

    ctx = mx.cpu()
    net_obj_det = gcv.model_zoo.get_model('ssd_512_mobilenet1.0_coco', pretrained=True, ctx=ctx)
    net_pose_est = get_model('simple_pose_resnet18_v1b', pretrained='ccd24037', ctx=ctx)
    net_obj_det.hybridize()
    net_pose_est.hybridize()

    cap = cv2.VideoCapture(0)


    while(True):

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

        key = cv2.waitKey(1)
        if (key == ord('q')) or (ret == False):
            cv2.destroyAllWindows()
            cap.release()
            break

        ## Image pre-processing
        frame_nd_orig = mx.nd.array(cv2.cvtColor(frame_np_orig, cv2.COLOR_BGR2RGB)).astype('uint8')
        frame_nd_new, frame_np_new = gcv.data.transforms.presets.ssd.transform_test(frame_nd_orig, short=512, max_size=700)
        frame_nd_new = frame_nd_new.as_in_context(ctx)

        ## run frame through network, detect keypoints for persons
        frame_nd_new = frame_nd_new.as_in_context(ctx)
        class_IDs, scores, bounding_boxes = net_obj_det(frame_nd_new)

        ## select only one class (person)
        selected_indices_person = np.where( ((class_IDs[0].asnumpy() == 0) & (scores[0].asnumpy() >= 0.35)))[0]
        selected_bboxes_person = bounding_boxes[0].asnumpy()[selected_indices_person]
        selected_class_IDs = class_IDs[0].asnumpy()[selected_indices_person]        # unused so far
        selected_scores = scores[0].asnumpy()[selected_indices_person]              # unused so far

        ## detection of body keypoints
        pose_input, upscale_bbox = detector_to_simple_pose(frame_np_new, class_IDs, scores, bounding_boxes, output_shape=(256, 192), ctx=ctx)
        if len(upscale_bbox) > 0:
            predicted_heatmap = net_pose_est(pose_input)
            pred_coords, confidence = heatmap_to_coord(predicted_heatmap, upscale_bbox)
            img = cv_plot_keypoints(frame_np_new, pred_coords, confidence, class_IDs, bounding_boxes, scores, box_thresh=0.5, keypoint_thresh=0.2)


        cv_plot_image(img)

    cap.release()


if __name__ == "__main__":
    main()

Я хотел попробовать квантование этих моделей и имел в виду эту страницу.. Итак, я пробовал что-то вроде этого:

      net_obj_det =   net_obj_det = quantize_net(net_obj_det, quantized_dtype='auto',
                               exclude_layers=None, exclude_layers_match=None,
                               calib_data=None, data_shapes=None,
                               calib_mode='naive', num_calib_examples=None,
                               ctx=mx.cpu(), logger=logging)

Но он возвращает ошибку:

       File "/home/lews/anaconda3/envs/gluon/lib/python3.8/site-packages/mxnet/contrib/quantization.py", line 820, in quantize_net
    data_shapes = dshapes
UnboundLocalError: local variable 'dshapes' referenced before assignment

На странице, на которую я указал, указано, что 'data_shapes' должен быть «Списком DataDesc, необходимого, если не указан calib_data». Как я должен это использовать?

0 ответов

Другие вопросы по тегам