Caffe2 - C++ API
A deep learning, cross platform ML framework
norm_planar_yuv_op.cc
1 
17 #include <array>
18 #include "caffe2/core/operator.h"
19 #include "caffe2/utils/math.h"
20 
21 namespace caffe2 {
22 
23 namespace {
24 
25 class NormalizePlanarYUVOp : public Operator<CPUContext> {
26  public:
27  USE_OPERATOR_FUNCTIONS(CPUContext);
28  using Operator<CPUContext>::Operator;
29 
30  bool RunOnDevice() {
31  const auto& X = Input(0);
32  const auto& M = Input(1); // mean
33  const auto& S = Input(2); // standard deviation
34  auto* Z = Output(0);
35  Z->ResizeLike(X);
36 
37  CAFFE_ENFORCE(X.dims().size() == 4);
38 
39  const auto N = X.dim32(0);
40  auto C = X.dim(1);
41  const auto H = X.dim(2);
42  const auto W = X.dim(3);
43  CAFFE_ENFORCE(C == M.dim(1));
44  CAFFE_ENFORCE(C == S.dim(1));
45  const auto* Xdata = X.data<float>();
46  auto* Zdata = Z->mutable_data<float>();
47 
48  int offset = H * W;
49  for (auto n = 0; n < N; n++) { // realistically N will always be 1
50  int batch_offset = n * C * offset;
51  for (auto c = 0; c < C; c++) {
52  ConstEigenVectorMap<float> channel_s(
53  &Xdata[batch_offset + (c * offset)], offset);
54  EigenVectorMap<float> channel_d(
55  &Zdata[batch_offset + (c * offset)], offset);
56  channel_d = channel_s.array() - M.data<float>()[c];
57  channel_d = channel_d.array() / S.data<float>()[c];
58  }
59  }
60  return true;
61  }
62 };
63 
64 REGISTER_CPU_OPERATOR(NormalizePlanarYUV, NormalizePlanarYUVOp);
65 OPERATOR_SCHEMA(NormalizePlanarYUV).NumInputs(3).NumOutputs(1);
66 } // namespace
67 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.