allwpilib/cscore/src/main/native/cpp/Image.hpp

// Copyright (c) FIRST and other WPILib contributors.
// Open Source Software; you can modify and/or share it under the terms of
// the WPILib BSD license file in the root directory of this project.

#pragma once

#include <string_view>
#include <vector>

#include <opencv2/core/core.hpp>

#ifdef __linux__
#include "default_init_allocator.hpp"
#endif
#include "wpi/util/PixelFormat.hpp"

namespace wpi::cs {

class Frame;

class Image {
  friend class Frame;

 public:
#ifndef __linux__
  explicit Image(size_t capacity) { m_data.reserve(capacity); }
#else
  explicit Image(size_t capacity)
      : m_data{capacity, default_init_allocator<uchar>{}} {
    m_data.resize(0);
  }
#endif

  Image(const Image&) = delete;
  Image& operator=(const Image&) = delete;

  // Getters
  operator std::string_view() const {  // NOLINT
    return str();
  }
  std::string_view str() const { return {data(), size()}; }
  size_t capacity() const { return m_data.capacity(); }
  const char* data() const {
    return reinterpret_cast<const char*>(m_data.data());
  }
  char* data() { return reinterpret_cast<char*>(m_data.data()); }
  size_t size() const { return m_data.size(); }

  const std::vector<uchar>& vec() const { return m_data; }
  std::vector<uchar>& vec() { return m_data; }

  void resize(size_t size) { m_data.resize(size); }
  void SetSize(size_t size) { m_data.resize(size); }

  cv::Mat AsMat() {
    int type;
    switch (pixelFormat) {
      case wpi::util::PixelFormat::YUYV:
      case wpi::util::PixelFormat::RGB565:
      case wpi::util::PixelFormat::Y16:
      case wpi::util::PixelFormat::UYVY:
        type = CV_8UC2;
        break;
      case wpi::util::PixelFormat::BGR:
        type = CV_8UC3;
        break;
      case wpi::util::PixelFormat::BGRA:
        type = CV_8UC4;
        break;
      case wpi::util::PixelFormat::GRAY:
      case wpi::util::PixelFormat::MJPEG:
      default:
        type = CV_8UC1;
        break;
    }
    return cv::Mat{height, width, type, m_data.data()};
  }

  int GetStride() const {
    switch (pixelFormat) {
      case wpi::util::PixelFormat::YUYV:
      case wpi::util::PixelFormat::RGB565:
      case wpi::util::PixelFormat::Y16:
      case wpi::util::PixelFormat::UYVY:
        return 2 * width;
      case wpi::util::PixelFormat::BGR:
        return 3 * width;
      case wpi::util::PixelFormat::BGRA:
        return 4 * width;
      case wpi::util::PixelFormat::GRAY:
        return width;
      case wpi::util::PixelFormat::MJPEG:
      default:
        return 0;
    }
  }

  cv::_InputArray AsInputArray() { return cv::_InputArray{m_data}; }

  bool Is(int width_, int height_) {
    return width == width_ && height == height_;
  }
  bool Is(int width_, int height_, wpi::util::PixelFormat pixelFormat_) {
    return width == width_ && height == height_ && pixelFormat == pixelFormat_;
  }
  bool Is(int width_, int height_, wpi::util::PixelFormat pixelFormat_,
          int jpegQuality_) {
    // Consider +/-5 on JPEG quality to be "close enough"
    return width == width_ && height == height_ &&
           pixelFormat == pixelFormat_ &&
           (pixelFormat != wpi::util::PixelFormat::MJPEG ||
            jpegQuality_ == -1 ||
            (jpegQuality != -1 && std::abs(jpegQuality - jpegQuality_) <= 5));
  }
  bool IsLarger(int width_, int height_) {
    return width >= width_ && height >= height_;
  }
  bool IsLarger(const Image& oth) {
    return width >= oth.width && height >= oth.height;
  }
  bool IsSmaller(int width_, int height_) { return !IsLarger(width_, height_); }
  bool IsSmaller(const Image& oth) { return !IsLarger(oth); }

 private:
  std::vector<uchar> m_data;

 public:
  wpi::util::PixelFormat pixelFormat{wpi::util::PixelFormat::UNKNOWN};
  int width{0};
  int height{0};
  int jpegQuality{-1};
};

}  // namespace wpi::cs