Merge branch 'main' into workflow_partition

Author: NicolasHug

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -203,6 +203,18 @@ VideoDecoder::BatchDecodedOutput::BatchDecodedOutput(
   frames = allocateEmptyHWCTensor(height, width, options.device, numFrames);
 }
 
+bool VideoDecoder::SwsContextKey::operator==(
+    const VideoDecoder::SwsContextKey& other) {
+  return decodedWidth == other.decodedWidth && decodedHeight == decodedHeight &&
+      decodedFormat == other.decodedFormat &&
+      outputWidth == other.outputWidth && outputHeight == other.outputHeight;
+}
+
+bool VideoDecoder::SwsContextKey::operator!=(
+    const VideoDecoder::SwsContextKey& other) {
+  return !(*this == other);
+}
+
 VideoDecoder::VideoDecoder() {}
 
 void VideoDecoder::initializeDecoder() {
@@ -1339,7 +1351,14 @@ int VideoDecoder::convertFrameToBufferUsingSwsScale(
 
   int expectedOutputHeight = outputTensor.sizes()[0];
   int expectedOutputWidth = outputTensor.sizes()[1];
-  if (activeStream.swsContext.get() == nullptr) {
+  auto curFrameSwsContextKey = SwsContextKey{
+      frame->width,
+      frame->height,
+      frameFormat,
+      expectedOutputWidth,
+      expectedOutputHeight};
+  if (activeStream.swsContext.get() == nullptr ||
+      activeStream.swsContextKey != curFrameSwsContextKey) {
     SwsContext* swsContext = sws_getContext(
         frame->width,
         frame->height,
@@ -1373,6 +1392,7 @@ int VideoDecoder::convertFrameToBufferUsingSwsScale(
         brightness,
         contrast,
         saturation);
+    activeStream.swsContextKey = curFrameSwsContextKey;
     activeStream.swsContext.reset(swsContext);
   }
   SwsContext* swsContext = activeStream.swsContext.get();

src/torchcodec/decoders/_core/VideoDecoder.h

@@ -317,6 +317,15 @@ class VideoDecoder {
     AVFilterContext* sourceContext = nullptr;
     AVFilterContext* sinkContext = nullptr;
   };
+  struct SwsContextKey {
+    int decodedWidth;
+    int decodedHeight;
+    AVPixelFormat decodedFormat;
+    int outputWidth;
+    int outputHeight;
+    bool operator==(const SwsContextKey&);
+    bool operator!=(const SwsContextKey&);
+  };
   // Stores information for each stream.
   struct StreamInfo {
     int streamIndex = -1;
@@ -337,6 +346,7 @@ class VideoDecoder {
     ColorConversionLibrary colorConversionLibrary = FILTERGRAPH;
     std::vector<FrameInfo> keyFrames;
     std::vector<FrameInfo> allFrames;
+    SwsContextKey swsContextKey;
     UniqueSwsContext swsContext;
   };
   VideoDecoder();