| 18 | namespace kernel { |
| 19 | template<typename T, bool IsColor> |
| 20 | void meanShift(Param<T> out, CParam<T> in, const float spatialSigma, |
| 21 | const float chromaticSigma, const unsigned numIterations) { |
| 22 | typedef typename std::conditional<std::is_same<T, double>::value, double, |
| 23 | float>::type AccType; |
| 24 | |
| 25 | const af::dim4 dims = in.dims(); |
| 26 | const af::dim4 istrides = in.strides(); |
| 27 | const af::dim4 ostrides = out.strides(); |
| 28 | const unsigned bCount = (IsColor ? 1 : dims[2]); |
| 29 | const unsigned channels = (IsColor ? dims[2] : 1); |
| 30 | const dim_t radius = std::max((int)(spatialSigma * 1.5f), 1); |
| 31 | const AccType cvar = chromaticSigma * chromaticSigma; |
| 32 | |
| 33 | std::array<AccType, 4> currentCenterColors{{0}}; |
| 34 | std::array<AccType, 4> currentMeanColors{{0}}; |
| 35 | std::array<AccType, 4> tempColors{{0}}; |
| 36 | for (dim_t b3 = 0; b3 < dims[3]; ++b3) { |
| 37 | for (unsigned b2 = 0; b2 < bCount; ++b2) { |
| 38 | T* outData = out.get() + b2 * ostrides[2] + b3 * ostrides[3]; |
| 39 | const T* inData = in.get() + b2 * istrides[2] + b3 * istrides[3]; |
| 40 | |
| 41 | for (dim_t j = 0; j < dims[1]; ++j) { |
| 42 | dim_t j_in_off = j * istrides[1]; |
| 43 | dim_t j_out_off = j * ostrides[1]; |
| 44 | |
| 45 | for (dim_t i = 0; i < dims[0]; ++i) { |
| 46 | dim_t i_in_off = i * istrides[0]; |
| 47 | dim_t i_out_off = i * ostrides[0]; |
| 48 | |
| 49 | for (unsigned ch = 0; ch < channels; ++ch) |
| 50 | currentCenterColors[ch] = static_cast<AccType>( |
| 51 | inData[j_in_off + i_in_off + ch * istrides[2]]); |
| 52 | |
| 53 | int meanPosJ = j; |
| 54 | int meanPosI = i; |
| 55 | |
| 56 | // scope of meanshift iterations begin |
| 57 | for (unsigned it = 0; it < numIterations; ++it) { |
| 58 | int oldMeanPosJ = meanPosJ; |
| 59 | int oldMeanPosI = meanPosI; |
| 60 | unsigned count = 0; |
| 61 | int shift_y = 0; |
| 62 | int shift_x = 0; |
| 63 | |
| 64 | currentMeanColors.fill(0); |
| 65 | // Windowing operation |
| 66 | for (dim_t wj = -radius; wj <= radius; ++wj) { |
| 67 | int hit_count = 0; |
| 68 | dim_t tj = meanPosJ + wj; |
| 69 | if (tj < 0 || tj > dims[1] - 1) continue; |
| 70 | |
| 71 | dim_t tjstride = tj * istrides[1]; |
| 72 | |
| 73 | for (dim_t wi = -radius; wi <= radius; ++wi) { |
| 74 | dim_t ti = meanPosI + wi; |
| 75 | if (ti < 0 || ti > dims[0] - 1) continue; |
| 76 | |
| 77 | dim_t tistride = ti * istrides[0]; |