infer_camera.py 文件源码

def inference():
    cap = cv2.VideoCapture(0)


    # load model
    model = Netmodel('eval-model', CLASSES)
    serializers.load_npz(MODEL_NAME, model)
    cuda.get_device(GPU_ID).use()
    model.to_gpu()

    LUT = fromHEX2RGB(stats_opts['colormap'] )
    fig3, axarr3 = plt.subplots(1, 1)

    batchRGB = np.zeros((1, 3, NEWSIZE[1], NEWSIZE[0]), dtype='float32')

    while(True):
            # Capture frame-by-frame
            ret, frame = cap.read()

        # process frame
        im = misc.imresize(frame, NEWSIZE, interp='bilinear')
        # convertion from HxWxCH to CHxWxH
        batchRGB[0,:,:,:] = im.astype(np.float32).transpose((2,1,0))
        batchRGBn = batchRGB  - 127.0

        # data ready
        batch = chainer.Variable(cuda.cupy.asarray(batchRGBn))

        # make predictions
        model((batch, []), test_mode=2)
        pred = model.probs.data.argmax(1)
        # move data back to CPU
        pred_ = cuda.to_cpu(pred)

        pred_ = LUT[pred_+1,:].squeeze()
        pred_ = pred_.transpose((1,0,2))
        pred2 = cv2.cvtColor(pred_, cv2.COLOR_BGR2RGB)

            # Display the resulting frame
        cv2.imshow('frame',frame)
        cv2.imshow('pred',pred2)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    # When everything done, release the capture
    cap.release()
    cv2.destroyAllWindows()