001// Copyright (c) FIRST and other WPILib contributors.
002// Open Source Software; you can modify and/or share it under the terms of
003// the WPILib BSD license file in the root directory of this project.
004
005package edu.wpi.first.cscore;
006
007import edu.wpi.first.util.PixelFormat;
008import org.opencv.core.Mat;
009
010/**
011 * A source that represents a video camera. These sources require the WPILib OpenCV builds. For an
012 * alternate OpenCV, see the documentation how to build your own with RawSource.
013 */
014public class CvSource extends ImageSource {
015  /**
016   * Create an OpenCV source.
017   *
018   * @param name Source name (arbitrary unique identifier)
019   * @param mode Video mode being generated
020   */
021  public CvSource(String name, VideoMode mode) {
022    super(
023        CameraServerJNI.createRawSource(
024            name, true, mode.pixelFormat.getValue(), mode.width, mode.height, mode.fps));
025    OpenCvLoader.forceStaticLoad();
026  }
027
028  /**
029   * Create an OpenCV source.
030   *
031   * @param name Source name (arbitrary unique identifier)
032   * @param pixelFormat Pixel format
033   * @param width width
034   * @param height height
035   * @param fps fps
036   */
037  public CvSource(String name, PixelFormat pixelFormat, int width, int height, int fps) {
038    super(CameraServerJNI.createRawSource(name, true, pixelFormat.getValue(), width, height, fps));
039    OpenCvLoader.forceStaticLoad();
040  }
041
042  /**
043   * Put an OpenCV image and notify sinks
044   *
045   * <p>The image format is guessed from the number of channels. The channel mapping is as follows.
046   * 1: kGray 2: kYUYV 3: BGR 4: BGRA Any other channel numbers will throw an error. If your image
047   * is an in alternate format, use the overload that takes a PixelFormat.
048   *
049   * @param image OpenCV Image
050   */
051  public void putFrame(Mat image) {
052    // We only support 8 bit channels
053    boolean cleanupRequired = false;
054    Mat finalImage;
055    if (image.depth() == 0) {
056      finalImage = image;
057    } else {
058      finalImage = new Mat();
059      image.convertTo(finalImage, 0);
060      cleanupRequired = true;
061    }
062
063    try {
064      int channels = finalImage.channels();
065      PixelFormat format =
066          switch (channels) {
067            case 1 -> PixelFormat.kGray; // 1 channel is assumed Grayscale
068            case 2 -> PixelFormat.kYUYV; // 2 channels is assumed YUYV
069            case 3 -> PixelFormat.kBGR; // 3 channels is assumed BGR
070            case 4 -> PixelFormat.kBGRA; // 4 channels is assumed BGRA
071            default ->
072                throw new VideoException(
073                    "Unable to get pixel format for " + channels + " channels");
074          };
075
076      putFrame(finalImage, format, true);
077    } finally {
078      if (cleanupRequired) {
079        finalImage.release();
080      }
081    }
082  }
083
084  /**
085   * Put an OpenCV image and notify sinks.
086   *
087   * <p>The format of the Mat must match the PixelFormat. You will corrupt memory if they dont. With
088   * skipVerification false, we will verify the number of channels matches the pixel format. If
089   * skipVerification is true, this step is skipped and is passed straight through.
090   *
091   * @param image OpenCV image
092   * @param format The pixel format of the image
093   * @param skipVerification skip verifying pixel format
094   */
095  public void putFrame(Mat image, PixelFormat format, boolean skipVerification) {
096    // We only support 8-bit images, convert if necessary
097    boolean cleanupRequired = false;
098    Mat finalImage;
099    if (image.depth() == 0) {
100      finalImage = image;
101    } else {
102      finalImage = new Mat();
103      image.convertTo(finalImage, 0);
104      cleanupRequired = true;
105    }
106
107    try {
108      if (!skipVerification) {
109        verifyFormat(finalImage, format);
110      }
111      long step = image.step1() * image.elemSize1();
112      CameraServerJNI.putRawSourceFrameData(
113          m_handle,
114          finalImage.dataAddr(),
115          (int) finalImage.total() * finalImage.channels(),
116          finalImage.width(),
117          finalImage.height(),
118          (int) step,
119          format.getValue());
120
121    } finally {
122      if (cleanupRequired) {
123        finalImage.release();
124      }
125    }
126  }
127
128  private void verifyFormat(Mat image, PixelFormat pixelFormat) {
129    int channels = image.channels();
130    switch (pixelFormat) {
131      case kBGR:
132        if (channels == 3) {
133          return;
134        }
135        break;
136      case kBGRA:
137        if (channels == 4) {
138          return;
139        }
140        break;
141      case kGray:
142        if (channels == 1) {
143          return;
144        }
145        break;
146      case kRGB565:
147        if (channels == 2) {
148          return;
149        }
150        break;
151      case kUYVY:
152        if (channels == 2) {
153          return;
154        }
155        break;
156      case kY16:
157        if (channels == 2) {
158          return;
159        }
160        break;
161      case kYUYV:
162        if (channels == 2) {
163          return;
164        }
165        break;
166      case kMJPEG:
167        if (channels == 1) {
168          return;
169        }
170        break;
171      default:
172        break;
173    }
174  }
175}