aos/vision/blob/threshold.h

#ifndef AOS_VISION_BLOB_THRESHOLD_H_
#define AOS_VISION_BLOB_THRESHOLD_H_

#include "aos/vision/blob/range_image.h"
#include "aos/vision/image/image_types.h"

namespace aos {
namespace vision {
namespace threshold_internal {

// Performs thresholding in a given region using a function which determines
// whether a given point is in or out of the region.
//
// fn must return a bool when called with two integers (x, y).
template <typename PointTestFn>
RangeImage ThresholdPointsWithFunction(ImageFormat fmt, PointTestFn &&fn) {
  static_assert(
      std::is_convertible<PointTestFn, std::function<bool(int, int)>>::value,
      "Invalid threshold function");
  std::vector<std::vector<ImageRange>> result;
  result.reserve(fmt.h);
  // Iterate through each row.
  for (int y = 0; y < fmt.h; ++y) {
    // Whether we're currently in a range.
    bool in_range = false;
    int current_range_start = -1;
    std::vector<ImageRange> current_row_ranges;
    // Iterate through each pixel.
    for (int x = 0; x < fmt.w; ++x) {
      if (fn(x, y) != in_range) {
        if (in_range) {
          current_row_ranges.emplace_back(ImageRange(current_range_start, x));
        } else {
          current_range_start = x;
        }
        in_range = !in_range;
      }
    }
    if (in_range) {
      current_row_ranges.emplace_back(ImageRange(current_range_start, fmt.w));
    }
    result.push_back(current_row_ranges);
  }
  return RangeImage(0, std::move(result));
}

}  // namespace threshold_internal

// Thresholds an image using a function which determines whether a given pixel
// value is in or out of the region.
//
// fn must return a bool when called with a PixelRef.
template <typename ThresholdFn>
RangeImage ThresholdImageWithFunction(const ImagePtr &img, ThresholdFn &&fn) {
  static_assert(
      std::is_convertible<ThresholdFn, std::function<bool(PixelRef)>>::value,
      "Invalid threshold function");
  return threshold_internal::ThresholdPointsWithFunction(
      img.fmt(), [&](int x, int y) { return fn(img.get_px(x, y)); });
}

// Thresholds an image in YUYV format, selecting pixels with a Y (luma) greater
// than value.
//
// This is implemented via a simple function that pulls out the Y values and
// compares them each. It mostly exists for tests to compare against
// FastYuyvYThreshold, because it's obviously correct.
inline RangeImage SlowYuyvYThreshold(ImageFormat fmt, const char *data,
                                     uint8_t value) {
  return threshold_internal::ThresholdPointsWithFunction(
      fmt, [&](int x, int y) {
        uint8_t v = data[x * 2 + y * fmt.w * 2];
        return v > value;
      });
}

// Thresholds an image in YUYV format, selecting pixels with a Y (luma) greater
// than value. The width must be a multiple of 4.
//
// This is implemented via some tricky bit shuffling that goes fast.
RangeImage FastYuyvYThreshold(ImageFormat fmt, const char *data, uint8_t value);

}  // namespace vision
}  // namespace aos

#endif  //  AOS_VISION_BLOB_THRESHOLD_H_