From 29d55b15ccb84ef021f4d94c80498bd1b5c889d1 Mon Sep 17 00:00:00 2001
From: Sergiy Bilobrov <sergiy@fb.com>
Date: Wed, 18 May 2022 19:29:02 -0700
Subject: [PATCH] Handle images with AV_PIX_FMT_PAL8 pixel format in decoder
 callback

Summary: provider callback uses video sampler to copy/scale input frames using sws_scale function, added option to use av_image_copy to copy image planes in case if sws_getContext fails (required to support AV_PIX_FMT_PAL8 pixel format)

Reviewed By: ovoietsa

Differential Revision: D36466749

fbshipit-source-id: 1674f98a362db835d147800c24bdda6f82d2f8ff
---
 torchvision/csrc/io/decoder/video_sampler.cpp | 45 ++++++++++++++-----
 1 file changed, 35 insertions(+), 10 deletions(-)

diff --git a/torchvision/csrc/io/decoder/video_sampler.cpp b/torchvision/csrc/io/decoder/video_sampler.cpp
index 7816d70b706..0a604e43607 100644
--- a/torchvision/csrc/io/decoder/video_sampler.cpp
+++ b/torchvision/csrc/io/decoder/video_sampler.cpp
@@ -65,12 +65,29 @@ int transformImage(
   if ((result = preparePlanes(outFormat, out, planes, lines)) < 0) {
     return result;
   }
-  // NOTE: srcY stride always 0: this is a parameter of YUV format
-  if ((result = sws_scale(
-           context, srcSlice, srcStride, 0, inFormat.height, planes, lines)) <
-      0) {
-    LOG(ERROR) << "sws_scale failed, err: " << Util::generateErrorDesc(result);
-    return result;
+  if (context) {
+    // NOTE: srcY stride always 0: this is a parameter of YUV format
+    if ((result = sws_scale(
+            context, srcSlice, srcStride, 0, inFormat.height, planes, lines)) <
+        0) {
+      LOG(ERROR) << "sws_scale failed, err: " << Util::generateErrorDesc(result);
+      return result;
+    }
+  } else if (inFormat.width == outFormat.width &&
+    inFormat.height == outFormat.height &&
+    inFormat.format == outFormat.format) {
+    // Copy planes without using sws_scale if sws_getContext failed.
+    av_image_copy(
+      planes,
+      lines,
+      (const uint8_t**)srcSlice,
+      srcStride,
+      (AVPixelFormat)inFormat.format,
+      inFormat.width,
+      inFormat.height);
+  } else {
+    LOG(ERROR) << "Invalid scale context format " << inFormat.format;
+    return AVERROR(EINVAL);
   }
   return 0;
 }
@@ -159,6 +176,9 @@ bool VideoSampler::init(const SamplerParameters& params) {
           << params.out.video.minDimension << ", cropImage "
           << params.out.video.cropImage;
 
+  // set output format
+  params_ = params;
+
   scaleContext_ = sws_getContext(
       params.in.video.width,
       params.in.video.height,
@@ -170,10 +190,15 @@ bool VideoSampler::init(const SamplerParameters& params) {
       nullptr,
       nullptr,
       nullptr);
-
-  // set output format
-  params_ = params;
-
+  // sws_getContext might fail if in/out format == AV_PIX_FMT_PAL8 (png format)
+  // Return true if input and output formats/width/height are identical
+  // Check scaleContext_ for nullptr in transformImage to copy planes directly
+
+  if (params.in.video.width == scaleFormat_.width &&
+      params.in.video.height == scaleFormat_.height &&
+      params.in.video.format == scaleFormat_.format) {
+    return true;
+  }
   return scaleContext_ != nullptr;
 }