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 -> throw new VideoException(
072                "Unable to get pixel format for " + channels + " channels");
073          };
074
075      putFrame(finalImage, format, true);
076    } finally {
077      if (cleanupRequired) {
078        finalImage.release();
079      }
080    }
081  }
082
083  /**
084   * Put an OpenCV image and notify sinks.
085   *
086   * <p>The format of the Mat must match the PixelFormat. You will corrupt memory if they dont. With
087   * skipVerification false, we will verify the number of channels matches the pixel format. If
088   * skipVerification is true, this step is skipped and is passed straight through.
089   *
090   * @param image OpenCV image
091   * @param format The pixel format of the image
092   * @param skipVerification skip verifying pixel format
093   */
094  public void putFrame(Mat image, PixelFormat format, boolean skipVerification) {
095    // We only support 8-bit images, convert if necessary
096    boolean cleanupRequired = false;
097    Mat finalImage;
098    if (image.depth() == 0) {
099      finalImage = image;
100    } else {
101      finalImage = new Mat();
102      image.convertTo(finalImage, 0);
103      cleanupRequired = true;
104    }
105
106    try {
107      if (!skipVerification) {
108        verifyFormat(finalImage, format);
109      }
110      long step = image.step1() * image.elemSize1();
111      CameraServerJNI.putRawSourceFrameData(
112          m_handle,
113          finalImage.dataAddr(),
114          (int) finalImage.total() * finalImage.channels(),
115          finalImage.width(),
116          finalImage.height(),
117          (int) step,
118          format.getValue());
119
120    } finally {
121      if (cleanupRequired) {
122        finalImage.release();
123      }
124    }
125  }
126
127  private void verifyFormat(Mat image, PixelFormat pixelFormat) {
128    int channels = image.channels();
129    switch (pixelFormat) {
130      case kBGR:
131        if (channels == 3) {
132          return;
133        }
134        break;
135      case kBGRA:
136        if (channels == 4) {
137          return;
138        }
139        break;
140      case kGray:
141        if (channels == 1) {
142          return;
143        }
144        break;
145      case kRGB565:
146        if (channels == 2) {
147          return;
148        }
149        break;
150      case kUYVY:
151        if (channels == 2) {
152          return;
153        }
154        break;
155      case kY16:
156        if (channels == 2) {
157          return;
158        }
159        break;
160      case kYUYV:
161        if (channels == 2) {
162          return;
163        }
164        break;
165      case kMJPEG:
166        if (channels == 1) {
167          return;
168        }
169        break;
170      default:
171        break;
172    }
173  }
174}