Handle swscale correctly

Author: scotts

src/torchcodec/_core/CpuDeviceInterface.cpp

@@ -65,12 +65,12 @@ void CpuDeviceInterface::initialize(
   // choose color conversion library publicly; we only use this ability
   // internally.
 
-  // If any transforms are not swscale compatible, then we can't use swscale.
-  bool areTransformsSwScaleCompatible = true;
-  for (const auto& transform : transforms) {
-    areTransformsSwScaleCompatible =
-        areTransformsSwScaleCompatible && transform->isSwScaleCompatible();
-  }
+  // We can only use swscale when we have a single resize transform. Note that
+  // this means swscale will not support the case of having several,
+  // back-to-base resizes. There's no strong reason to even do that, but if
+  // someone does, it's more correct to implement that with filtergraph.
+  bool areTransformsSwScaleCompatible = transforms.empty() ||
+      (transforms.size() == 1 && transforms[0]->isResize());
 
   // swscale requires widths to be multiples of 32:
   // https://stackoverflow.com/questions/74351955/turn-off-sw-scale-conversion-to-planar-yuv-32-byte-alignment-requirements
@@ -95,8 +95,14 @@ void CpuDeviceInterface::initialize(
       (userRequestedSwScale || isWidthSwScaleCompatible)) {
     colorConversionLibrary_ = ColorConversionLibrary::SWSCALE;
 
-    // SCOTT NEXT TODO: set swsFlags_
-
+    // We established above that if the transforms are swscale compatible and
+    // non-empty, then they must have only one transforms, and that transform is
+    // ResizeTransform.
+    if (!transforms.empty()) {
+      auto resize = dynamic_cast<ResizeTransform*>(transforms[0].get());
+      TORCH_CHECK(resize != nullptr, "ResizeTransform expected but not found!")
+      swsFlags_ = resize->getSwsFlags();
+    }
   } else {
     colorConversionLibrary_ = ColorConversionLibrary::FILTERGRAPH;
 

src/torchcodec/_core/CpuDeviceInterface.h

@@ -72,8 +72,8 @@ class CpuDeviceInterface : public DeviceInterface {
   FrameDims outputDims_;
 
   // If we use swscale for resizing, the flags control the resizing algorithm.
-  // We exclusively get the value from the ResizeTransform.
-  int swsFlags_ = 0;
+  // We default to bilinear. Users can override this with a ResizeTransform.
+  int swsFlags_ = SWS_BILINEAR;
 
   // The copy filter just copies the input to the output. Computationally, it
   // should be a no-op. If we get no user-provided transforms, we will use the

src/torchcodec/_core/Transform.cpp

@@ -56,7 +56,7 @@ std::optional<FrameDims> ResizeTransform::getOutputFrameDims() const {
   return FrameDims(width_, height_);
 }
 
-bool ResizeTransform::isSwScaleCompatible() const {
+bool ResizeTransform::isResize() const {
   return true;
 }
 

src/torchcodec/_core/Transform.h

@@ -28,7 +28,9 @@ class Transform {
     return std::nullopt;
   }
 
-  virtual bool isSwScaleCompatible() const {
+  // The ResizeTransform is special, because it is the only transform that
+  // swscale can handle.
+  virtual bool isResize() const {
     return false;
   }
 };
@@ -47,7 +49,7 @@ class ResizeTransform : public Transform {
 
   std::string getFilterGraphCpu() const override;
   std::optional<FrameDims> getOutputFrameDims() const override;
-  bool isSwScaleCompatible() const override;
+  bool isResize() const override;
 
   int getSwsFlags() const;