[cscore] CvSink: Allow specifying output PixelFormat (wpilibsuite#5943)

srimanachanta · Nov 22, 2023 · 437cc91 · 437cc91
1 parent 25b7dca
commit 437cc91
Show file tree

Hide file tree

Showing 10 changed files with 70 additions and 35 deletions.
diff --git a/cscore/src/main/java/edu/wpi/first/cscore/CameraServerCvJNI.java b/cscore/src/main/java/edu/wpi/first/cscore/CameraServerCvJNI.java
@@ -64,7 +64,7 @@ public static native int createCvSource(
 
   public static native void putSourceFrame(int source, long imageNativeObj);
 
-  public static native int createCvSink(String name);
+  public static native int createCvSink(String name, int pixelFormat);
 
   // public static native int createCvSinkCallback(String name,
   //                            void (*processFrame)(long time));

diff --git a/cscore/src/main/java/edu/wpi/first/cscore/CvSink.java b/cscore/src/main/java/edu/wpi/first/cscore/CvSink.java
@@ -4,6 +4,7 @@
 
 package edu.wpi.first.cscore;
 
+import edu.wpi.first.cscore.VideoMode.PixelFormat;
 import org.opencv.core.Mat;
 
 /**
@@ -16,9 +17,20 @@ public class CvSink extends ImageSink {
    * get each new image.
    *
    * @param name Source name (arbitrary unique identifier)
+   * @param pixelFormat Source pixel format
+   */
+  public CvSink(String name, PixelFormat pixelFormat) {
+    super(CameraServerCvJNI.createCvSink(name, pixelFormat.getValue()));
+  }
+
+  /**
+   * Create a sink for accepting OpenCV images. WaitForFrame() must be called on the created sink to
+   * get each new image. Defaults to kBGR for pixelFormat
+   *
+   * @param name Source name (arbitrary unique identifier)
    */
   public CvSink(String name) {
-    super(CameraServerCvJNI.createCvSink(name));
+    this(name, PixelFormat.kBGR);
   }
 
   /// Create a sink for accepting OpenCV images in a separate thread.

diff --git a/cscore/src/main/native/cpp/CvSinkImpl.cpp b/cscore/src/main/native/cpp/CvSinkImpl.cpp
@@ -19,16 +19,18 @@
 using namespace cs;
 
 CvSinkImpl::CvSinkImpl(std::string_view name, wpi::Logger& logger,
-                       Notifier& notifier, Telemetry& telemetry)
-    : SinkImpl{name, logger, notifier, telemetry} {
+                       Notifier& notifier, Telemetry& telemetry,
+                       VideoMode::PixelFormat pixelFormat)
+    : SinkImpl{name, logger, notifier, telemetry}, m_pixelFormat{pixelFormat} {
   m_active = true;
   // m_thread = std::thread(&CvSinkImpl::ThreadMain, this);
 }
 
 CvSinkImpl::CvSinkImpl(std::string_view name, wpi::Logger& logger,
                        Notifier& notifier, Telemetry& telemetry,
+                       VideoMode::PixelFormat pixelFormat,
                        std::function<void(uint64_t time)> processFrame)
-    : SinkImpl{name, logger, notifier, telemetry} {}
+    : SinkImpl{name, logger, notifier, telemetry}, m_pixelFormat{pixelFormat} {}
 
 CvSinkImpl::~CvSinkImpl() {
   Stop();
@@ -65,7 +67,7 @@ uint64_t CvSinkImpl::GrabFrame(cv::Mat& image) {
     return 0;  // signal error
   }
 
-  if (!frame.GetCv(image)) {
+  if (!frame.GetCv(image, m_pixelFormat)) {
     // Shouldn't happen, but just in case...
     std::this_thread::sleep_for(std::chrono::milliseconds(20));
     return 0;
@@ -91,7 +93,7 @@ uint64_t CvSinkImpl::GrabFrame(cv::Mat& image, double timeout) {
     return 0;  // signal error
   }
 
-  if (!frame.GetCv(image)) {
+  if (!frame.GetCv(image, m_pixelFormat)) {
     // Shouldn't happen, but just in case...
     std::this_thread::sleep_for(std::chrono::milliseconds(20));
     return 0;
@@ -127,20 +129,23 @@ void CvSinkImpl::ThreadMain() {
 
 namespace cs {
 
-CS_Sink CreateCvSink(std::string_view name, CS_Status* status) {
+CS_Sink CreateCvSink(std::string_view name, VideoMode::PixelFormat pixelFormat,
+                     CS_Status* status) {
   auto& inst = Instance::GetInstance();
   return inst.CreateSink(
       CS_SINK_CV, std::make_shared<CvSinkImpl>(name, inst.logger, inst.notifier,
-                                               inst.telemetry));
+                                               inst.telemetry, pixelFormat));
 }
 
 CS_Sink CreateCvSinkCallback(std::string_view name,
+                             VideoMode::PixelFormat pixelFormat,
                              std::function<void(uint64_t time)> processFrame,
                              CS_Status* status) {
   auto& inst = Instance::GetInstance();
   return inst.CreateSink(
-      CS_SINK_CV, std::make_shared<CvSinkImpl>(name, inst.logger, inst.notifier,
-                                               inst.telemetry, processFrame));
+      CS_SINK_CV,
+      std::make_shared<CvSinkImpl>(name, inst.logger, inst.notifier,
+                                   inst.telemetry, pixelFormat, processFrame));
 }
 
 static constexpr unsigned SinkMask = CS_SINK_CV | CS_SINK_RAW;
@@ -206,15 +211,19 @@ void SetSinkEnabled(CS_Sink sink, bool enabled, CS_Status* status) {
 
 extern "C" {
 
-CS_Sink CS_CreateCvSink(const char* name, CS_Status* status) {
-  return cs::CreateCvSink(name, status);
+CS_Sink CS_CreateCvSink(const char* name, enum CS_PixelFormat pixelFormat,
+                        CS_Status* status) {
+  return cs::CreateCvSink(
+      name, static_cast<VideoMode::PixelFormat>(pixelFormat), status);
 }
 
-CS_Sink CS_CreateCvSinkCallback(const char* name, void* data,
+CS_Sink CS_CreateCvSinkCallback(const char* name,
+                                enum CS_PixelFormat pixelFormat, void* data,
                                 void (*processFrame)(void* data, uint64_t time),
                                 CS_Status* status) {
   return cs::CreateCvSinkCallback(
-      name, [=](uint64_t time) { processFrame(data, time); }, status);
+      name, static_cast<VideoMode::PixelFormat>(pixelFormat),
+      [=](uint64_t time) { processFrame(data, time); }, status);
 }
 
 void CS_SetSinkDescription(CS_Sink sink, const char* description,

diff --git a/cscore/src/main/native/cpp/CvSinkImpl.h b/cscore/src/main/native/cpp/CvSinkImpl.h
@@ -25,9 +25,9 @@ class SourceImpl;
 class CvSinkImpl : public SinkImpl {
  public:
   CvSinkImpl(std::string_view name, wpi::Logger& logger, Notifier& notifier,
-             Telemetry& telemetry);
+             Telemetry& telemetry, VideoMode::PixelFormat pixelFormat);
   CvSinkImpl(std::string_view name, wpi::Logger& logger, Notifier& notifier,
-             Telemetry& telemetry,
+             Telemetry& telemetry, VideoMode::PixelFormat pixelFormat,
              std::function<void(uint64_t time)> processFrame);
   ~CvSinkImpl() override;
 
@@ -42,6 +42,7 @@ class CvSinkImpl : public SinkImpl {
   std::atomic_bool m_active;  // set to false to terminate threads
   std::thread m_thread;
   std::function<void(uint64_t time)> m_processFrame;
+  VideoMode::PixelFormat m_pixelFormat;
 };
 
 }  // namespace cs

diff --git a/cscore/src/main/native/cpp/Frame.cpp b/cscore/src/main/native/cpp/Frame.cpp
@@ -709,8 +709,9 @@ Image* Frame::GetImageImpl(int width, int height,
   return ConvertImpl(cur, pixelFormat, requiredJpegQuality, defaultJpegQuality);
 }
 
-bool Frame::GetCv(cv::Mat& image, int width, int height) {
-  Image* rawImage = GetImage(width, height, VideoMode::kBGR);
+bool Frame::GetCv(cv::Mat& image, int width, int height,
+                  VideoMode::PixelFormat pixelFormat) {
+  Image* rawImage = GetImage(width, height, pixelFormat);
   if (!rawImage) {
     return false;
   }

diff --git a/cscore/src/main/native/cpp/Frame.h b/cscore/src/main/native/cpp/Frame.h
@@ -219,10 +219,11 @@ class Frame {
                         defaultQuality);
   }
 
-  bool GetCv(cv::Mat& image) {
-    return GetCv(image, GetOriginalWidth(), GetOriginalHeight());
+  bool GetCv(cv::Mat& image, VideoMode::PixelFormat pixelFormat) {
+    return GetCv(image, GetOriginalWidth(), GetOriginalHeight(), pixelFormat);
   }
-  bool GetCv(cv::Mat& image, int width, int height);
+  bool GetCv(cv::Mat& image, int width, int height,
+             VideoMode::PixelFormat pixelFormat);
 
  private:
   Image* ConvertImpl(Image* image, VideoMode::PixelFormat pixelFormat,

diff --git a/cscore/src/main/native/cpp/jni/CameraServerJNI.cpp b/cscore/src/main/native/cpp/jni/CameraServerJNI.cpp
@@ -1392,18 +1392,20 @@ Java_edu_wpi_first_cscore_CameraServerJNI_createMjpegServer
 /*
  * Class:     edu_wpi_first_cscore_CameraServerCvJNI
  * Method:    createCvSink
- * Signature: (Ljava/lang/String;)I
+ * Signature: (Ljava/lang/String;I)I
  */
 JNIEXPORT jint JNICALL
 Java_edu_wpi_first_cscore_CameraServerCvJNI_createCvSink
-  (JNIEnv* env, jclass, jstring name)
+  (JNIEnv* env, jclass, jstring name, jint pixelFormat)
 {
   if (!name) {
     nullPointerEx.Throw(env, "name cannot be null");
     return 0;
   }
   CS_Status status = 0;
-  auto val = cs::CreateCvSink(JStringRef{env, name}.str(), &status);
+  auto val = cs::CreateCvSink(
+      JStringRef{env, name}.str(),
+      static_cast<cs::VideoMode::PixelFormat>(pixelFormat), &status);
   CheckStatus(env, status);
   return val;
 }

diff --git a/cscore/src/main/native/include/cscore_c.h b/cscore/src/main/native/include/cscore_c.h
@@ -382,8 +382,10 @@ void CS_SetSourceEnumPropertyChoices(CS_Source source, CS_Property property,
  */
 CS_Sink CS_CreateMjpegServer(const char* name, const char* listenAddress,
                              int port, CS_Status* status);
-CS_Sink CS_CreateCvSink(const char* name, CS_Status* status);
-CS_Sink CS_CreateCvSinkCallback(const char* name, void* data,
+CS_Sink CS_CreateCvSink(const char* name, enum CS_PixelFormat pixelFormat,
+                        CS_Status* status);
+CS_Sink CS_CreateCvSinkCallback(const char* name,
+                                enum CS_PixelFormat pixelFormat, void* data,
                                 void (*processFrame)(void* data, uint64_t time),
                                 CS_Status* status);
 /** @} */

diff --git a/cscore/src/main/native/include/cscore_cpp.h b/cscore/src/main/native/include/cscore_cpp.h
@@ -316,8 +316,10 @@ void SetSourceEnumPropertyChoices(CS_Source source, CS_Property property,
  */
 CS_Sink CreateMjpegServer(std::string_view name, std::string_view listenAddress,
                           int port, CS_Status* status);
-CS_Sink CreateCvSink(std::string_view name, CS_Status* status);
+CS_Sink CreateCvSink(std::string_view name, VideoMode::PixelFormat pixelFormat,
+                     CS_Status* status);
 CS_Sink CreateCvSinkCallback(std::string_view name,
+                             VideoMode::PixelFormat pixelFormat,
                              std::function<void(uint64_t time)> processFrame,
                              CS_Status* status);
 

diff --git a/cscore/src/main/native/include/cscore_cv.h b/cscore/src/main/native/include/cscore_cv.h
@@ -127,8 +127,10 @@ class CvSink : public ImageSink {
    * image.
    *
    * @param name Source name (arbitrary unique identifier)
+   * @param pixelFormat Source pixel format
    */
-  explicit CvSink(std::string_view name);
+  explicit CvSink(std::string_view name, VideoMode::PixelFormat pixelFormat =
+                                             VideoMode::PixelFormat::kBGR);
 
   /**
    * Create a sink for accepting OpenCV images in a separate thread.
@@ -141,9 +143,10 @@ class CvSink : public ImageSink {
    *        time=0 if an error occurred.  processFrame should call GetImage()
    *        or GetError() as needed, but should not call (except in very
    *        unusual circumstances) WaitForImage().
+   * @param pixelFormat Source pixel format
    */
-  CvSink(std::string_view name,
-         std::function<void(uint64_t time)> processFrame);
+  CvSink(std::string_view name, std::function<void(uint64_t time)> processFrame,
+         VideoMode::PixelFormat pixelFormat = VideoMode::PixelFormat::kBGR);
 
   /**
    * Wait for the next frame and get the image.
@@ -184,13 +187,15 @@ inline void CvSource::PutFrame(cv::Mat& image) {
   PutSourceFrame(m_handle, image, &m_status);
 }
 
-inline CvSink::CvSink(std::string_view name) {
-  m_handle = CreateCvSink(name, &m_status);
+inline CvSink::CvSink(std::string_view name,
+                      VideoMode::PixelFormat pixelFormat) {
+  m_handle = CreateCvSink(name, pixelFormat, &m_status);
 }
 
 inline CvSink::CvSink(std::string_view name,
-                      std::function<void(uint64_t time)> processFrame) {
-  m_handle = CreateCvSinkCallback(name, processFrame, &m_status);
+                      std::function<void(uint64_t time)> processFrame,
+                      VideoMode::PixelFormat pixelFormat) {
+  m_handle = CreateCvSinkCallback(name, pixelFormat, processFrame, &m_status);
 }
 
 inline uint64_t CvSink::GrabFrame(cv::Mat& image, double timeout) const {