Simple dispersion spot finding implementation in C++

simple_dispersion.cc

#include <vector>

int dispersion_filter(std::vector<uint16_t> &image_in,
                      std::vector<uint16_t> &mask_out, 
                      int NY, int NX) {
  int32_t knl = 3;

  float sigma_b = 6.0f, sigma_s = 3.0f;

  std::vector<uint32_t> im(NY * NX);
  std::vector<uint32_t> m_sat(NY * NX);
  std::vector<uint32_t> i_sat(NY * NX);
  std::vector<uint32_t> i2_sat(NY * NX);

  // copy image data in, fill SATs
  for (uint32_t i = 0, k = 0; i < NY; i++) {
    uint32_t _m = 0, _i = 0, _i2 = 0;
    for (uint32_t j = 0; j < NX; j++, k++) {
      uint32_t m = image_in[k] > 0xfffd ? 0 : 1;
      uint32_t p = m * image_in[k];
      _m += m;
      _i += p;
      _i2 += p * p;
      im[k] = p;
      m_sat[k] = i > 0 ? _m + m_sat[k - NX] : _m;
      i_sat[k] = i > 0 ? _i + i_sat[k - NX] : _i;
      i2_sat[k] = i > 0 ? _i2 + i2_sat[k - NX] : _i2;
    }
  }

  // roll over now computing the mean, variance etc.
  for (uint32_t i = 0, k = 0; i < NY; i++) {
    for (uint32_t j = 0; j < NX; j++, k++) {
      int32_t j0 = j - knl - 1;
      int32_t j1 = j < (NX - knl) ? j + knl : NX - 1;
      int32_t i0 = i - knl - 1;
      int32_t i1 = i < (NY - knl) ? i + knl : NY - 1;

      int32_t a = i1 * NX + j1;
      int32_t b = i0 * NX + j1;
      int32_t c = i1 * NX + j0;
      int32_t d = i0 * NX + j0;

      uint32_t m_sum = m_sat[a], i_sum = i_sat[a], i2_sum = i2_sat[a];

      if (j0 >= 0 && i0 >= 0) {
        m_sum += m_sat[d] - m_sat[b] - m_sat[c];
        i_sum += i_sat[d] - i_sat[b] - i_sat[c];
        i2_sum += i2_sat[d] - i2_sat[b] - i2_sat[c];
      } else if (j0 >= 0) {
        m_sum -= m_sat[c];
        i_sum -= i_sat[c];
        i2_sum -= i2_sat[c];
      } else if (i0 >= 0) {
        m_sum -= m_sat[b];
        i_sum -= i_sat[b];
        i2_sum -= i2_sat[b];
      }

      uint16_t signal = 0;
      uint32_t p = im[k];

      if (p > 0 && m_sum >= 2) {
        float bg_lhs = (float)m_sum * i2_sum - (float)i_sum * i_sum -
                       (float)i_sum * (m_sum - 1);
        float bg_rhs = i_sum * sigma_b * sqrtf((float)2.0f * (m_sum - 1));
        uint16_t background = bg_lhs > bg_rhs;
        float fg_lhs = (float)m_sum * p - (float)i_sum;
        float fg_rhs = sigma_s * sqrtf((float)i_sum * m_sum);
        uint16_t foreground = fg_lhs > fg_rhs;
        signal = background && foreground;
      }

      // save the pixel value for later use in connected component labelling
      mask_out[k] = signal * p;
    }
  }

  return 0;
}