torchcodec 0.10.0__cp312-cp312-manylinux_2_28_x86_64.whl

torchcodec/_core/FFMPEGCommon.cpp

@@ -0,0 +1,684 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include "FFMPEGCommon.h"
+
+#include <c10/util/Exception.h>
+
+extern "C" {
+#include <libavfilter/avfilter.h>
+#include <libavfilter/buffersink.h>
+}
+
+namespace facebook::torchcodec {
+
+AutoAVPacket::AutoAVPacket() : avPacket_(av_packet_alloc()) {
+  TORCH_CHECK(avPacket_ != nullptr, "Couldn't allocate avPacket.");
+}
+
+AutoAVPacket::~AutoAVPacket() {
+  av_packet_free(&avPacket_);
+}
+
+ReferenceAVPacket::ReferenceAVPacket(AutoAVPacket& shared)
+    : avPacket_(shared.avPacket_) {}
+
+ReferenceAVPacket::~ReferenceAVPacket() {
+  av_packet_unref(avPacket_);
+}
+
+AVPacket* ReferenceAVPacket::get() {
+  return avPacket_;
+}
+
+AVPacket* ReferenceAVPacket::operator->() {
+  return avPacket_;
+}
+
+AVCodecOnlyUseForCallingAVFindBestStream
+makeAVCodecOnlyUseForCallingAVFindBestStream(const AVCodec* codec) {
+#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(59, 18, 100)
+  return const_cast<AVCodec*>(codec);
+#else
+  return codec;
+#endif
+}
+
+std::string getFFMPEGErrorStringFromErrorCode(int errorCode) {
+  char errorBuffer[AV_ERROR_MAX_STRING_SIZE] = {0};
+  av_strerror(errorCode, errorBuffer, AV_ERROR_MAX_STRING_SIZE);
+  return std::string(errorBuffer);
+}
+
+int64_t getDuration(const UniqueAVFrame& avFrame) {
+#if LIBAVUTIL_VERSION_MAJOR < 58
+  return avFrame->pkt_duration;
+#else
+  return avFrame->duration;
+#endif
+}
+
+void setDuration(const UniqueAVFrame& avFrame, int64_t duration) {
+#if LIBAVUTIL_VERSION_MAJOR < 58
+  avFrame->pkt_duration = duration;
+#else
+  avFrame->duration = duration;
+#endif
+}
+
+const int* getSupportedSampleRates(const AVCodec& avCodec) {
+  const int* supportedSampleRates = nullptr;
+#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100) // FFmpeg >= 7.1
+  int numSampleRates = 0;
+  int ret = avcodec_get_supported_config(
+      nullptr,
+      &avCodec,
+      AV_CODEC_CONFIG_SAMPLE_RATE,
+      0,
+      reinterpret_cast<const void**>(&supportedSampleRates),
+      &numSampleRates);
+  if (ret < 0 || supportedSampleRates == nullptr) {
+    // Return nullptr to skip validation in validateSampleRate.
+    return nullptr;
+  }
+#else
+  supportedSampleRates = avCodec.supported_samplerates;
+#endif
+  return supportedSampleRates;
+}
+
+const AVPixelFormat* getSupportedPixelFormats(const AVCodec& avCodec) {
+  const AVPixelFormat* supportedPixelFormats = nullptr;
+#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100) // FFmpeg >= 7.1
+  int numPixelFormats = 0;
+  int ret = avcodec_get_supported_config(
+      nullptr,
+      &avCodec,
+      AV_CODEC_CONFIG_PIX_FORMAT,
+      0,
+      reinterpret_cast<const void**>(&supportedPixelFormats),
+      &numPixelFormats);
+  if (ret < 0 || supportedPixelFormats == nullptr) {
+    TORCH_CHECK(false, "Couldn't get supported pixel formats from encoder.");
+  }
+#else
+  supportedPixelFormats = avCodec.pix_fmts;
+#endif
+  return supportedPixelFormats;
+}
+
+const AVSampleFormat* getSupportedOutputSampleFormats(const AVCodec& avCodec) {
+  const AVSampleFormat* supportedSampleFormats = nullptr;
+#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100) // FFmpeg >= 7.1
+  int numSampleFormats = 0;
+  int ret = avcodec_get_supported_config(
+      nullptr,
+      &avCodec,
+      AV_CODEC_CONFIG_SAMPLE_FORMAT,
+      0,
+      reinterpret_cast<const void**>(&supportedSampleFormats),
+      &numSampleFormats);
+  if (ret < 0 || supportedSampleFormats == nullptr) {
+    // Return nullptr to use default output format in
+    // findBestOutputSampleFormat.
+    return nullptr;
+  }
+#else
+  supportedSampleFormats = avCodec.sample_fmts;
+#endif
+  return supportedSampleFormats;
+}
+
+int getNumChannels(const UniqueAVFrame& avFrame) {
+#if LIBAVFILTER_VERSION_MAJOR > 8 || \
+    (LIBAVFILTER_VERSION_MAJOR == 8 && LIBAVFILTER_VERSION_MINOR >= 44)
+  return avFrame->ch_layout.nb_channels;
+#else
+  int numChannels = av_get_channel_layout_nb_channels(avFrame->channel_layout);
+  // Handle FFmpeg 4 bug where channel_layout and numChannels are 0 or unset
+  // Set values based on avFrame->channels which appears to be correct
+  // to allow successful initialization of SwrContext
+  if (numChannels == 0 && avFrame->channels > 0) {
+    avFrame->channel_layout = av_get_default_channel_layout(avFrame->channels);
+    numChannels = avFrame->channels;
+  }
+  return numChannels;
+#endif
+}
+
+int getNumChannels(const SharedAVCodecContext& avCodecContext) {
+#if LIBAVFILTER_VERSION_MAJOR > 8 || \
+    (LIBAVFILTER_VERSION_MAJOR == 8 && LIBAVFILTER_VERSION_MINOR >= 44)
+  return avCodecContext->ch_layout.nb_channels;
+#else
+  return avCodecContext->channels;
+#endif
+}
+
+int getNumChannels(const AVCodecParameters* codecpar) {
+  TORCH_CHECK(codecpar != nullptr, "codecpar is null")
+#if LIBAVFILTER_VERSION_MAJOR > 8 || \
+    (LIBAVFILTER_VERSION_MAJOR == 8 && LIBAVFILTER_VERSION_MINOR >= 44)
+  return codecpar->ch_layout.nb_channels;
+#else
+  return codecpar->channels;
+#endif
+}
+
+void setDefaultChannelLayout(
+    UniqueAVCodecContext& avCodecContext,
+    int numChannels) {
+#if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4
+  AVChannelLayout channel_layout;
+  av_channel_layout_default(&channel_layout, numChannels);
+  avCodecContext->ch_layout = channel_layout;
+#else
+  uint64_t channel_layout = av_get_default_channel_layout(numChannels);
+  avCodecContext->channel_layout = channel_layout;
+  avCodecContext->channels = numChannels;
+#endif
+}
+
+void setDefaultChannelLayout(UniqueAVFrame& avFrame, int numChannels) {
+#if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4
+  AVChannelLayout channel_layout;
+  av_channel_layout_default(&channel_layout, numChannels);
+  avFrame->ch_layout = channel_layout;
+#else
+  uint64_t channel_layout = av_get_default_channel_layout(numChannels);
+  avFrame->channel_layout = channel_layout;
+  avFrame->channels = numChannels;
+#endif
+}
+
+void validateNumChannels(const AVCodec& avCodec, int numChannels) {
+#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100) // FFmpeg >= 7.1
+  std::stringstream supportedNumChannels;
+  const AVChannelLayout* supportedLayouts = nullptr;
+  int numLayouts = 0;
+  int ret = avcodec_get_supported_config(
+      nullptr,
+      &avCodec,
+      AV_CODEC_CONFIG_CHANNEL_LAYOUT,
+      0,
+      reinterpret_cast<const void**>(&supportedLayouts),
+      &numLayouts);
+  if (ret < 0 || supportedLayouts == nullptr) {
+    // If we can't validate, we must assume it'll be fine. If not, FFmpeg will
+    // eventually raise.
+    return;
+  }
+  for (int i = 0; i < numLayouts; ++i) {
+    if (i > 0) {
+      supportedNumChannels << ", ";
+    }
+    supportedNumChannels << supportedLayouts[i].nb_channels;
+    if (numChannels == supportedLayouts[i].nb_channels) {
+      return;
+    }
+  }
+#elif LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4
+  if (avCodec.ch_layouts == nullptr) {
+    // If we can't validate, we must assume it'll be fine. If not, FFmpeg will
+    // eventually raise.
+    return;
+  }
+  // FFmpeg doc indicate that the ch_layouts array is terminated by a zeroed
+  // layout, so checking for nb_channels == 0 should indicate its end.
+  for (auto i = 0; avCodec.ch_layouts[i].nb_channels != 0; ++i) {
+    if (numChannels == avCodec.ch_layouts[i].nb_channels) {
+      return;
+    }
+  }
+  // At this point it seems that the encoder doesn't support the requested
+  // number of channels, so we error out.
+  std::stringstream supportedNumChannels;
+  for (auto i = 0; avCodec.ch_layouts[i].nb_channels != 0; ++i) {
+    if (i > 0) {
+      supportedNumChannels << ", ";
+    }
+    supportedNumChannels << avCodec.ch_layouts[i].nb_channels;
+  }
+#else // FFmpeg <= 4
+  if (avCodec.channel_layouts == nullptr) {
+    // can't validate, same as above.
+    return;
+  }
+  for (auto i = 0; avCodec.channel_layouts[i] != 0; ++i) {
+    if (numChannels ==
+        av_get_channel_layout_nb_channels(avCodec.channel_layouts[i])) {
+      return;
+    }
+  }
+  // At this point it seems that the encoder doesn't support the requested
+  // number of channels, so we error out.
+  std::stringstream supportedNumChannels;
+  for (auto i = 0; avCodec.channel_layouts[i] != 0; ++i) {
+    if (i > 0) {
+      supportedNumChannels << ", ";
+    }
+    supportedNumChannels << av_get_channel_layout_nb_channels(
+        avCodec.channel_layouts[i]);
+  }
+#endif
+  TORCH_CHECK(
+      false,
+      "Desired number of channels (",
+      numChannels,
+      ") is not supported by the ",
+      "encoder. Supported number of channels are: ",
+      supportedNumChannels.str(),
+      ".");
+}
+
+namespace {
+#if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4
+
+// Returns:
+// - the srcAVFrame's channel layout if srcAVFrame has outNumChannels
+// - the default channel layout with outNumChannels otherwise.
+AVChannelLayout getOutputChannelLayout(
+    int outNumChannels,
+    const UniqueAVFrame& srcAVFrame) {
+  AVChannelLayout outLayout;
+  if (outNumChannels == getNumChannels(srcAVFrame)) {
+    outLayout = srcAVFrame->ch_layout;
+  } else {
+    av_channel_layout_default(&outLayout, outNumChannels);
+  }
+  return outLayout;
+}
+
+#else
+
+// Same as above
+int64_t getOutputChannelLayout(
+    int outNumChannels,
+    const UniqueAVFrame& srcAVFrame) {
+  int64_t outLayout;
+  if (outNumChannels == getNumChannels(srcAVFrame)) {
+    outLayout = srcAVFrame->channel_layout;
+  } else {
+    outLayout = av_get_default_channel_layout(outNumChannels);
+  }
+  return outLayout;
+}
+#endif
+} // namespace
+
+// Sets dstAVFrame' channel layout to getOutputChannelLayout(): see doc above
+void setChannelLayout(
+    UniqueAVFrame& dstAVFrame,
+    const UniqueAVFrame& srcAVFrame,
+    int outNumChannels) {
+#if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4
+  AVChannelLayout outLayout =
+      getOutputChannelLayout(outNumChannels, srcAVFrame);
+  auto status = av_channel_layout_copy(&dstAVFrame->ch_layout, &outLayout);
+  TORCH_CHECK(
+      status == AVSUCCESS,
+      "Couldn't copy channel layout to avFrame: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+#else
+  dstAVFrame->channel_layout =
+      getOutputChannelLayout(outNumChannels, srcAVFrame);
+  dstAVFrame->channels = outNumChannels;
+#endif
+}
+
+UniqueAVFrame allocateAVFrame(
+    int numSamples,
+    int sampleRate,
+    int numChannels,
+    AVSampleFormat sampleFormat) {
+  auto avFrame = UniqueAVFrame(av_frame_alloc());
+  TORCH_CHECK(avFrame != nullptr, "Couldn't allocate AVFrame.");
+
+  avFrame->nb_samples = numSamples;
+  avFrame->sample_rate = sampleRate;
+  setDefaultChannelLayout(avFrame, numChannels);
+  avFrame->format = sampleFormat;
+  auto status = av_frame_get_buffer(avFrame.get(), 0);
+
+  TORCH_CHECK(
+      status == AVSUCCESS,
+      "Couldn't allocate avFrame's buffers: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+
+  status = av_frame_make_writable(avFrame.get());
+  TORCH_CHECK(
+      status == AVSUCCESS,
+      "Couldn't make AVFrame writable: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+  return avFrame;
+}
+
+SwrContext* createSwrContext(
+    AVSampleFormat srcSampleFormat,
+    AVSampleFormat outSampleFormat,
+    int srcSampleRate,
+    int outSampleRate,
+    const UniqueAVFrame& srcAVFrame,
+    int outNumChannels) {
+  SwrContext* swrContext = nullptr;
+  int status = AVSUCCESS;
+#if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4
+  AVChannelLayout outLayout =
+      getOutputChannelLayout(outNumChannels, srcAVFrame);
+  status = swr_alloc_set_opts2(
+      &swrContext,
+      &outLayout,
+      outSampleFormat,
+      outSampleRate,
+      &srcAVFrame->ch_layout,
+      srcSampleFormat,
+      srcSampleRate,
+      0,
+      nullptr);
+
+  TORCH_CHECK(
+      status == AVSUCCESS,
+      "Couldn't create SwrContext: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+#else
+  int64_t outLayout = getOutputChannelLayout(outNumChannels, srcAVFrame);
+  swrContext = swr_alloc_set_opts(
+      nullptr,
+      outLayout,
+      outSampleFormat,
+      outSampleRate,
+      srcAVFrame->channel_layout,
+      srcSampleFormat,
+      srcSampleRate,
+      0,
+      nullptr);
+#endif
+
+  TORCH_CHECK(swrContext != nullptr, "Couldn't create swrContext");
+  status = swr_init(swrContext);
+  TORCH_CHECK(
+      status == AVSUCCESS,
+      "Couldn't initialize SwrContext: ",
+      getFFMPEGErrorStringFromErrorCode(status),
+      ". If the error says 'Invalid argument', it's likely that you are using "
+      "a buggy FFmpeg version. FFmpeg4 is known to fail here in some "
+      "valid scenarios. Try to upgrade FFmpeg?");
+  return swrContext;
+}
+
+AVFilterContext* createAVFilterContextWithOptions(
+    AVFilterGraph* filterGraph,
+    const AVFilter* buffer,
+    const enum AVPixelFormat outputFormat) {
+  AVFilterContext* avFilterContext = nullptr;
+  const char* filterName = "out";
+
+  enum AVPixelFormat pixFmts[] = {outputFormat, AV_PIX_FMT_NONE};
+
+// av_opt_set_int_list was replaced by av_opt_set_array() in FFmpeg 8.
+#if LIBAVUTIL_VERSION_MAJOR >= 60 // FFmpeg >= 8
+  // Output options like pixel_formats must be set before filter init
+  avFilterContext =
+      avfilter_graph_alloc_filter(filterGraph, buffer, filterName);
+  TORCH_CHECK(
+      avFilterContext != nullptr, "Failed to allocate buffer filter context.");
+
+  // When setting pix_fmts, only the first element is used, so nb_elems = 1
+  // AV_PIX_FMT_NONE acts as a terminator for the array in av_opt_set_int_list
+  int status = av_opt_set_array(
+      avFilterContext,
+      "pixel_formats",
+      AV_OPT_SEARCH_CHILDREN,
+      0, // start_elem
+      1, // nb_elems
+      AV_OPT_TYPE_PIXEL_FMT,
+      pixFmts);
+  TORCH_CHECK(
+      status >= 0,
+      "Failed to set pixel format for buffer filter: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+
+  status = avfilter_init_str(avFilterContext, nullptr);
+  TORCH_CHECK(
+      status >= 0,
+      "Failed to initialize buffer filter: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+#else // FFmpeg <= 7
+  // For older FFmpeg versions, create filter and then set options
+  int status = avfilter_graph_create_filter(
+      &avFilterContext, buffer, filterName, nullptr, nullptr, filterGraph);
+  TORCH_CHECK(
+      status >= 0,
+      "Failed to create buffer filter: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+
+  status = av_opt_set_int_list(
+      avFilterContext,
+      "pix_fmts",
+      pixFmts,
+      AV_PIX_FMT_NONE,
+      AV_OPT_SEARCH_CHILDREN);
+  TORCH_CHECK(
+      status >= 0,
+      "Failed to set pixel formats for buffer filter: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+#endif
+
+  return avFilterContext;
+}
+
+UniqueAVFrame convertAudioAVFrameSamples(
+    const UniqueSwrContext& swrContext,
+    const UniqueAVFrame& srcAVFrame,
+    AVSampleFormat outSampleFormat,
+    int outSampleRate,
+    int outNumChannels) {
+  UniqueAVFrame convertedAVFrame(av_frame_alloc());
+  TORCH_CHECK(
+      convertedAVFrame,
+      "Could not allocate frame for sample format conversion.");
+
+  convertedAVFrame->pts = srcAVFrame->pts;
+  convertedAVFrame->format = static_cast<int>(outSampleFormat);
+
+  convertedAVFrame->sample_rate = outSampleRate;
+  int srcSampleRate = srcAVFrame->sample_rate;
+  if (srcSampleRate != outSampleRate) {
+    // Note that this is an upper bound on the number of output samples.
+    // `swr_convert()` will likely not fill convertedAVFrame with that many
+    // samples if sample rate conversion is needed. It will buffer the last few
+    // ones because those require future samples. That's also why we reset
+    // nb_samples after the call to `swr_convert()`.
+    // We could also use `swr_get_out_samples()` to determine the number of
+    // output samples, but empirically `av_rescale_rnd()` seems to provide a
+    // tighter bound.
+    convertedAVFrame->nb_samples = av_rescale_rnd(
+        swr_get_delay(swrContext.get(), srcSampleRate) + srcAVFrame->nb_samples,
+        outSampleRate,
+        srcSampleRate,
+        AV_ROUND_UP);
+  } else {
+    convertedAVFrame->nb_samples = srcAVFrame->nb_samples;
+  }
+
+  setChannelLayout(convertedAVFrame, srcAVFrame, outNumChannels);
+
+  auto status = av_frame_get_buffer(convertedAVFrame.get(), 0);
+  TORCH_CHECK(
+      status == AVSUCCESS,
+      "Could not allocate frame buffers for sample format conversion: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+
+  auto numConvertedSamples = swr_convert(
+      swrContext.get(),
+      convertedAVFrame->data,
+      convertedAVFrame->nb_samples,
+      static_cast<const uint8_t**>(
+          const_cast<const uint8_t**>(srcAVFrame->data)),
+      srcAVFrame->nb_samples);
+  // numConvertedSamples can be 0 if we're downsampling by a great factor and
+  // the first frame doesn't contain a lot of samples. It should be handled
+  // properly by the caller.
+  TORCH_CHECK(
+      numConvertedSamples >= 0,
+      "Error in swr_convert: ",
+      getFFMPEGErrorStringFromErrorCode(numConvertedSamples));
+
+  // See comment above about nb_samples
+  convertedAVFrame->nb_samples = numConvertedSamples;
+
+  return convertedAVFrame;
+}
+
+void setFFmpegLogLevel() {
+  auto logLevel = AV_LOG_QUIET;
+  const char* logLevelEnvPtr = std::getenv("TORCHCODEC_FFMPEG_LOG_LEVEL");
+  if (logLevelEnvPtr != nullptr) {
+    std::string logLevelEnv(logLevelEnvPtr);
+    if (logLevelEnv == "QUIET") {
+      logLevel = AV_LOG_QUIET;
+    } else if (logLevelEnv == "PANIC") {
+      logLevel = AV_LOG_PANIC;
+    } else if (logLevelEnv == "FATAL") {
+      logLevel = AV_LOG_FATAL;
+    } else if (logLevelEnv == "ERROR") {
+      logLevel = AV_LOG_ERROR;
+    } else if (logLevelEnv == "WARNING") {
+      logLevel = AV_LOG_WARNING;
+    } else if (logLevelEnv == "INFO") {
+      logLevel = AV_LOG_INFO;
+    } else if (logLevelEnv == "VERBOSE") {
+      logLevel = AV_LOG_VERBOSE;
+    } else if (logLevelEnv == "DEBUG") {
+      logLevel = AV_LOG_DEBUG;
+    } else if (logLevelEnv == "TRACE") {
+      logLevel = AV_LOG_TRACE;
+    } else {
+      TORCH_CHECK(
+          false,
+          "Invalid TORCHCODEC_FFMPEG_LOG_LEVEL: ",
+          logLevelEnv,
+          ". Use e.g. 'QUIET', 'PANIC', 'VERBOSE', etc.");
+    }
+  }
+  av_log_set_level(logLevel);
+}
+
+AVIOContext* avioAllocContext(
+    uint8_t* buffer,
+    int buffer_size,
+    int write_flag,
+    void* opaque,
+    AVIOReadFunction read_packet,
+    AVIOWriteFunction write_packet,
+    AVIOSeekFunction seek) {
+  return avio_alloc_context(
+      buffer,
+      buffer_size,
+      write_flag,
+      opaque,
+      read_packet,
+// The buf parameter of the write function is not const before FFmpeg 7.
+#if LIBAVFILTER_VERSION_MAJOR >= 10 // FFmpeg >= 7
+      write_packet,
+#else
+      reinterpret_cast<AVIOWriteFunctionOld>(write_packet),
+#endif
+      seek);
+}
+
+double ptsToSeconds(int64_t pts, const AVRational& timeBase) {
+  // To perform the multiplication before the division, av_q2d is not used
+  return static_cast<double>(pts) * timeBase.num / timeBase.den;
+}
+
+int64_t secondsToClosestPts(double seconds, const AVRational& timeBase) {
+  return static_cast<int64_t>(
+      std::round(seconds * timeBase.den / timeBase.num));
+}
+
+int64_t computeSafeDuration(
+    const AVRational& frameRate,
+    const AVRational& timeBase) {
+  if (frameRate.num <= 0 || frameRate.den <= 0 || timeBase.num <= 0 ||
+      timeBase.den <= 0) {
+    return 0;
+  } else {
+    return (static_cast<int64_t>(frameRate.den) * timeBase.den) /
+        (static_cast<int64_t>(timeBase.num) * frameRate.num);
+  }
+}
+
+SwsFrameContext::SwsFrameContext(
+    int inputWidth,
+    int inputHeight,
+    AVPixelFormat inputFormat,
+    int outputWidth,
+    int outputHeight)
+    : inputWidth(inputWidth),
+      inputHeight(inputHeight),
+      inputFormat(inputFormat),
+      outputWidth(outputWidth),
+      outputHeight(outputHeight) {}
+
+bool SwsFrameContext::operator==(const SwsFrameContext& other) const {
+  return inputWidth == other.inputWidth && inputHeight == other.inputHeight &&
+      inputFormat == other.inputFormat && outputWidth == other.outputWidth &&
+      outputHeight == other.outputHeight;
+}
+
+bool SwsFrameContext::operator!=(const SwsFrameContext& other) const {
+  return !(*this == other);
+}
+
+UniqueSwsContext createSwsContext(
+    const SwsFrameContext& swsFrameContext,
+    AVColorSpace colorspace,
+    AVPixelFormat outputFormat,
+    int swsFlags) {
+  SwsContext* swsContext = sws_getContext(
+      swsFrameContext.inputWidth,
+      swsFrameContext.inputHeight,
+      swsFrameContext.inputFormat,
+      swsFrameContext.outputWidth,
+      swsFrameContext.outputHeight,
+      outputFormat,
+      swsFlags,
+      nullptr,
+      nullptr,
+      nullptr);
+  TORCH_CHECK(swsContext, "sws_getContext() returned nullptr");
+
+  int* invTable = nullptr;
+  int* table = nullptr;
+  int srcRange, dstRange, brightness, contrast, saturation;
+  int ret = sws_getColorspaceDetails(
+      swsContext,
+      &invTable,
+      &srcRange,
+      &table,
+      &dstRange,
+      &brightness,
+      &contrast,
+      &saturation);
+  TORCH_CHECK(ret != -1, "sws_getColorspaceDetails returned -1");
+
+  const int* colorspaceTable = sws_getCoefficients(colorspace);
+  ret = sws_setColorspaceDetails(
+      swsContext,
+      colorspaceTable,
+      srcRange,
+      colorspaceTable,
+      dstRange,
+      brightness,
+      contrast,
+      saturation);
+  TORCH_CHECK(ret != -1, "sws_setColorspaceDetails returned -1");
+
+  return UniqueSwsContext(swsContext);
+}
+
+} // namespace facebook::torchcodec