mean_image.h 3 KB
Newer Older
1
// Copyright (C) 2018-2019 Intel Corporation
openvino-pushbot's avatar
openvino-pushbot committed
2 3 4 5 6 7 8
// SPDX-License-Identifier: Apache-2.0
//

#pragma once

#include "inference_engine.hpp"
#include "mkldnn_dims.h"
Alexey Suhov's avatar
Alexey Suhov committed
9
#include "ie_parallel.hpp"
openvino-pushbot's avatar
openvino-pushbot committed
10 11 12 13 14 15 16 17 18 19 20
#include <vector>
#include <limits>

namespace MKLDNNPlugin {

class MeanImage {
public:
    MeanImage();

public:
    void Load(const MKLDNNDims& inputDims, InferenceEngine::InputInfo::Ptr inputInfo);
21
    void Subtract(const MKLDNNDims &inputDims, float *input, InferenceEngine::Layout layout);
openvino-pushbot's avatar
openvino-pushbot committed
22 23

    template<typename T, typename std::enable_if<std::is_integral<T>::value>::type* = nullptr>
24
    void Subtract(const MKLDNNDims &inputDims, T *input, InferenceEngine::Layout layout) {
openvino-pushbot's avatar
openvino-pushbot committed
25 26 27 28 29 30
        IE_ASSERT(input != nullptr);

        if (inputDims.ndims() != 4) {
            THROW_IE_EXCEPTION << "Expecting input as 4 dimension blob with format NxCxHxW.";
        }

31 32 33 34
        if (layout != InferenceEngine::NCHW && layout != InferenceEngine::NHWC) {
            THROW_IE_EXCEPTION << "Expecting input layout NCHW or NHWC.";
        }

openvino-pushbot's avatar
openvino-pushbot committed
35 36 37 38 39
        int MB = inputDims[0];
        int srcSize = inputDims.size() / MB;

        if (meanBuffer && meanBuffer->size()) {
            const float * meanBufferValues = meanBuffer->readOnly();
Alexey Suhov's avatar
Alexey Suhov committed
40 41 42 43 44 45 46 47

            InferenceEngine::parallel_for2d(MB, srcSize, [&](int mb, int i) {
                int buf = input[srcSize * mb + i];
                buf -= meanBufferValues[i];
                if (buf < std::numeric_limits<T>::min()) buf = std::numeric_limits<T>::min();
                if (buf > std::numeric_limits<T>::max()) buf = std::numeric_limits<T>::max();
                input[srcSize * mb + i] = buf;
            });
openvino-pushbot's avatar
openvino-pushbot committed
48 49 50 51
        } else if (!meanValues.empty()) {
            int C = inputDims[1];
            srcSize /= inputDims[1];

52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70
            if (layout == InferenceEngine::NCHW) {
                InferenceEngine::parallel_for3d(MB, C, srcSize, [&](int mb, int c, int i) {
                    int buf = input[srcSize * mb * C + c * srcSize + i];
                    buf -= meanValues[c];
                    if (buf < std::numeric_limits<T>::min()) buf = std::numeric_limits<T>::min();
                    if (buf > std::numeric_limits<T>::max()) buf = std::numeric_limits<T>::max();
                    input[srcSize * mb * C + c * srcSize + i] = buf;
                });
            } else if (layout == InferenceEngine::NHWC) {
                InferenceEngine::parallel_for2d(MB, srcSize, [&](int mb, int i) {
                    for (int c = 0; c < C; c++) {
                        int buf = input[mb * srcSize * C + i * C + c];
                        buf -= meanValues[c];
                        if (buf < std::numeric_limits<T>::min()) buf = std::numeric_limits<T>::min();
                        if (buf > std::numeric_limits<T>::max()) buf = std::numeric_limits<T>::max();
                        input[mb * srcSize * C + i * C + c] = buf;
                    }
                });
            }
openvino-pushbot's avatar
openvino-pushbot committed
71 72 73 74 75 76 77 78 79 80
        }
    }

private:
    std::vector<float> meanValues;

    InferenceEngine::TBlob<float>::Ptr meanBuffer;
};

}  // namespace MKLDNNPlugin