allwpilib/src/Frame.cpp

/*----------------------------------------------------------------------------*/
/* Copyright (c) FIRST 2016. All Rights Reserved.                             */
/* Open Source Software - may be modified and shared by FRC teams. The code   */
/* must be accompanied by the FIRST BSD license file in the root directory of */
/* the project.                                                               */
/*----------------------------------------------------------------------------*/

#include "Frame.h"

#include "opencv2/core/core.hpp"
#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/highgui/highgui.hpp"

#include "Log.h"
#include "SourceImpl.h"

using namespace cs;

Frame::Frame(SourceImpl& source, llvm::StringRef error, Time time)
    : m_impl{source.AllocFrameImpl().release()} {
  m_impl->refcount = 1;
  m_impl->error = error;
  m_impl->time = time;
}

Frame::Frame(SourceImpl& source, std::unique_ptr<Image> image, Time time)
    : m_impl{source.AllocFrameImpl().release()} {
  m_impl->refcount = 1;
  m_impl->error.resize(0);
  m_impl->time = time;
  m_impl->images.push_back(image.release());
}

Image* Frame::GetNearestImage(int width, int height) const {
  if (!m_impl) return nullptr;
  std::lock_guard<std::recursive_mutex> lock(m_impl->mutex);
  Image* found = nullptr;

  // Ideally we want the smallest image at least width/height in size
  for (auto i : m_impl->images) {
    if (i->IsLarger(width, height) && (!found || (i->IsSmaller(*found))))
      found = i;
  }
  if (found) return found;

  // Find the largest image (will be less than width/height)
  for (auto i : m_impl->images) {
    if (!found || (i->IsLarger(*found))) found = i;
  }
  if (found) return found;

  // Shouldn't reach this, but just in case...
  return m_impl->images.empty() ? nullptr : m_impl->images[0];
}

Image* Frame::GetNearestImage(int width, int height,
                              VideoMode::PixelFormat pixelFormat) const {
  if (!m_impl) return nullptr;
  std::lock_guard<std::recursive_mutex> lock(m_impl->mutex);
  Image* found = nullptr;

  // We want the smallest image at least width/height (or the next largest),
  // but the primary search order is in order of conversion cost.
  // If we don't find exactly what we want, we prefer non-JPEG source images
  // (because JPEG source images require a decompression step).
  // While the searching takes a little time, it pales in comparison to the
  // image processing to come, so it's worth spending a little extra time
  // looking for the most efficient conversion.

  // 1) Same width, height, and pixelFormat (e.g. exactly what we want)
  for (auto i : m_impl->images) {
    if (i->Is(width, height, pixelFormat)) return i;
  }

  // 2) Same width, height, different (but non-JPEG) pixelFormat (color conv)
  // 2a) If we want JPEG output, prefer BGR over other pixel formats
  if (pixelFormat == VideoMode::kMJPEG) {
    for (auto i : m_impl->images) {
      if (i->Is(width, height, VideoMode::kBGR)) return i;
    }
  }

  for (auto i : m_impl->images) {
    if (i->Is(width, height) && i->pixelFormat != VideoMode::kMJPEG) return i;
  }

  // 3) Different width, height, same pixelFormat (only if non-JPEG) (resample)
  if (pixelFormat != VideoMode::kMJPEG) {
    // 3a) Smallest image at least width/height in size
    for (auto i : m_impl->images) {
      if (i->IsLarger(width, height) && i->pixelFormat == pixelFormat &&
          (!found || (i->IsSmaller(*found))))
        found = i;
    }
    if (found) return found;

    // 3b) Largest image (less than width/height)
    for (auto i : m_impl->images) {
      if (i->pixelFormat == pixelFormat && (!found || (i->IsLarger(*found))))
        found = i;
    }
    if (found) return found;
  }

  // 4) Different width, height, different (but non-JPEG) pixelFormat
  //    (color conversion + resample)
  // 4a) Smallest image at least width/height in size
  for (auto i : m_impl->images) {
    if (i->IsLarger(width, height) && i->pixelFormat != VideoMode::kMJPEG &&
        (!found || (i->IsSmaller(*found))))
      found = i;
  }
  if (found) return found;

  // 4b) Largest image (less than width/height)
  for (auto i : m_impl->images) {
    if (i->pixelFormat != VideoMode::kMJPEG &&
        (!found || (i->IsLarger(*found))))
      found = i;
  }
  if (found) return found;

  // 5) Same width, height, JPEG pixelFormat (decompression)
  for (auto i : m_impl->images) {
    if (i->Is(width, height, VideoMode::kMJPEG)) return i;
  }

  // 6) Different width, height, JPEG pixelFormat (decompression)
  // 6a) Smallest image at least width/height in size
  for (auto i : m_impl->images) {
    if (i->IsLarger(width, height) && i->pixelFormat == VideoMode::kMJPEG &&
        (!found || (i->IsSmaller(*found))))
      found = i;
  }
  if (found) return found;

  // 6b) Largest image (less than width/height)
  for (auto i : m_impl->images) {
    if (i->pixelFormat != VideoMode::kMJPEG &&
        (!found || (i->IsLarger(*found))))
      found = i;
  }
  if (found) return found;

  // Shouldn't reach this, but just in case...
  return m_impl->images.empty() ? nullptr : m_impl->images[0];
}

Image* Frame::Convert(Image* image, VideoMode::PixelFormat pixelFormat,
                      int jpegQuality) {
  if (!image || image->pixelFormat == pixelFormat) return image;
  Image* cur = image;

  // If the source image is a JPEG, we need to decode it before we can do
  // anything else with it.  Note that if the destination format is JPEG, we
  // still need to do this (unless it was already a JPEG, in which case we
  // would have returned above).
  if (cur->pixelFormat == VideoMode::kMJPEG) {
    cur = ConvertMJPEGToBGR(cur);
    if (pixelFormat == VideoMode::kBGR) return cur;
  }

  // Color convert; if ultimate destination is JPEG, we need to convert to BGR
  switch (pixelFormat) {
    case VideoMode::kRGB565:
      // If source is YUYV, need to convert to BGR first
      if (cur->pixelFormat == VideoMode::kYUYV) {
        // Check to see if BGR version already exists...
        if (Image* newImage =
                GetExistingImage(cur->width, cur->height, VideoMode::kBGR))
          cur = newImage;
        else
          cur = ConvertYUYVToBGR(cur);
      }
      return ConvertBGRToRGB565(cur);
    case VideoMode::kBGR:
    case VideoMode::kMJPEG:
      if (cur->pixelFormat == VideoMode::kYUYV)
        cur = ConvertYUYVToBGR(cur);
      else if (cur->pixelFormat == VideoMode::kRGB565)
        cur = ConvertRGB565ToBGR(cur);
      break;
    case VideoMode::kYUYV:
    default:
      return nullptr;  // Unsupported
  }

  // Compress if destination is JPEG
  if (pixelFormat == VideoMode::kMJPEG)
    cur = ConvertBGRToMJPEG(cur, jpegQuality);

  return cur;
}

Image* Frame::ConvertMJPEGToBGR(Image* image) {
  if (!image || image->pixelFormat != VideoMode::kMJPEG) return nullptr;

  // Allocate an BGR image
  auto newImage =
      m_impl->source.AllocImage(VideoMode::kBGR, image->width, image->height,
                                image->width * image->height * 3);

  // Decode
  cv::Mat newMat = newImage->AsMat();
  cv::imdecode(image->AsInputArray(), cv::IMREAD_COLOR, &newMat);

  // Save the result
  Image* rv = newImage.release();
  if (m_impl) {
    std::lock_guard<std::recursive_mutex> lock(m_impl->mutex);
    m_impl->images.push_back(rv);
  }
  return rv;
}

Image* Frame::ConvertYUYVToBGR(Image* image) {
  if (!image || image->pixelFormat != VideoMode::kYUYV) return nullptr;

  // Allocate a BGR image
  auto newImage =
      m_impl->source.AllocImage(VideoMode::kBGR, image->width, image->height,
                                image->width * image->height * 3);

  // Convert
  cv::cvtColor(image->AsMat(), newImage->AsMat(), cv::COLOR_YUV2BGR_YUYV);

  // Save the result
  Image* rv = newImage.release();
  if (m_impl) {
    std::lock_guard<std::recursive_mutex> lock(m_impl->mutex);
    m_impl->images.push_back(rv);
  }
  return rv;
}

Image* Frame::ConvertBGRToRGB565(Image* image) {
  if (!image || image->pixelFormat != VideoMode::kBGR) return nullptr;

  // Allocate a RGB565 image
  auto newImage =
      m_impl->source.AllocImage(VideoMode::kRGB565, image->width, image->height,
                                image->width * image->height * 2);

  // Convert
  cv::cvtColor(image->AsMat(), newImage->AsMat(), cv::COLOR_RGB2BGR565);

  // Save the result
  Image* rv = newImage.release();
  if (m_impl) {
    std::lock_guard<std::recursive_mutex> lock(m_impl->mutex);
    m_impl->images.push_back(rv);
  }
  return rv;
}

Image* Frame::ConvertRGB565ToBGR(Image* image) {
  if (!image || image->pixelFormat != VideoMode::kRGB565) return nullptr;

  // Allocate a BGR image
  auto newImage =
      m_impl->source.AllocImage(VideoMode::kBGR, image->width, image->height,
                                image->width * image->height * 3);

  // Convert
  cv::cvtColor(image->AsMat(), newImage->AsMat(), cv::COLOR_BGR5652RGB);

  // Save the result
  Image* rv = newImage.release();
  if (m_impl) {
    std::lock_guard<std::recursive_mutex> lock(m_impl->mutex);
    m_impl->images.push_back(rv);
  }
  return rv;
}

Image* Frame::ConvertBGRToMJPEG(Image* image, int quality) {
  if (!image || image->pixelFormat != VideoMode::kBGR) return nullptr;
  if (!m_impl) return nullptr;
  std::lock_guard<std::recursive_mutex> lock(m_impl->mutex);

  // Allocate a JPEG image.  We don't actually know what the resulting size
  // will be; while the destination will automatically grow, doing so will
  // cause an extra malloc, so we don't want to be too conservative here.
  // Per Wikipedia, Q=100 on a sample image results in 8.25 bits per pixel,
  // this is a little bit more conservative in assuming 50% space savings over
  // the equivalent BGR image.
  auto newImage =
      m_impl->source.AllocImage(VideoMode::kMJPEG, image->width, image->height,
                                image->width * image->height * 1.5);

  // Compress
  if (m_impl->compressionParams.empty()) {
    m_impl->compressionParams.push_back(CV_IMWRITE_JPEG_QUALITY);
    m_impl->compressionParams.push_back(quality);
  } else {
    m_impl->compressionParams[1] = quality;
  }
  cv::imencode(".jpg", image->AsMat(), newImage->vec(),
               m_impl->compressionParams);

  // Save the result
  Image* rv = newImage.release();
  m_impl->images.push_back(rv);
  return rv;
}

Image* Frame::GetImage(int width, int height,
                       VideoMode::PixelFormat pixelFormat, int jpegQuality) {
  if (!m_impl) return nullptr;
  std::lock_guard<std::recursive_mutex> lock(m_impl->mutex);
  Image* cur = GetNearestImage(width, height, pixelFormat);
  if (!cur || cur->Is(width, height, pixelFormat)) return cur;

  DEBUG4("converting image from "
         << cur->width << "x" << cur->height << " type " << cur->pixelFormat
         << " to " << width << "x" << height << " type " << pixelFormat);

  // If the source image is a JPEG, we need to decode it before we can do
  // anything else with it.  Note that if the destination format is JPEG, we
  // still need to do this (unless the width/height were the same, in which
  // case we already returned the existing JPEG above).
  if (cur->pixelFormat == VideoMode::kMJPEG) cur = ConvertMJPEGToBGR(cur);

  // Resize
  if (!cur->Is(width, height)) {
    // Allocate an image.
    auto newImage = m_impl->source.AllocImage(
        cur->pixelFormat, width, height,
        width * height * (cur->size() / (cur->width * cur->height)));

    // Resize
    cv::Mat newMat = newImage->AsMat();
    cv::resize(cur->AsMat(), newMat, newMat.size(), 0, 0);

    // Save the result
    cur = newImage.release();
    m_impl->images.push_back(cur);
  }

  // Convert to output format
  return Convert(cur, pixelFormat, jpegQuality);
}

bool Frame::GetCv(cv::Mat& image, int width, int height) {
  Image* rawImage = GetImage(width, height, VideoMode::kBGR);
  if (!rawImage) return false;
  rawImage->AsMat().copyTo(image);
  return true;
}

void Frame::ReleaseFrame() {
  for (auto image : m_impl->images)
    m_impl->source.ReleaseImage(std::unique_ptr<Image>(image));
  m_impl->images.clear();
  m_impl->source.ReleaseFrameImpl(std::unique_ptr<Impl>(m_impl));
  m_impl = nullptr;
}