Skip to content

Commit 179f4e7

Browse files
committed
Add BICUBIC and LANCZOS interpolation modes to ResizeTransform
This enables FFmpeg's bicubic and lanczos scaling filters for higher-quality video resizing. The existing infrastructure already supports multiple modes; this change exposes them through the full E2E path. Changes: - Transform.h: Added BICUBIC and LANCZOS to InterpolationMode enum - Transform.cpp: Added cases for bicubic and lanczos in toFilterGraphInterpolation() and toSwsInterpolation() functions - custom_ops.cpp: Updated makeResizeTransform() to parse optional interpolation mode from transform spec string (e.g., 'resize, 256, 256, bicubic') - _decoder_transforms.py: Added 'interpolation' parameter to Resize class, updated _from_torchvision() to map TorchVision interpolation modes - test_transform_ops.py: Updated tests to allow BICUBIC/LANCZOS, added tests for new interpolation modes This allows users to use TorchCodec's Resize with bicubic interpolation: TorchCodecResize(size=(256, 256), interpolation='bicubic') Or use TorchVision's Resize with bicubic and it will be converted: v2.Resize(size=(256, 256), interpolation=v2.InterpolationMode.BICUBIC)
1 parent e88bee9 commit 179f4e7

5 files changed

Lines changed: 299 additions & 18 deletions

File tree

src/torchcodec/_core/Transform.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ std::string toFilterGraphInterpolation(
1717
switch (mode) {
1818
case ResizeTransform::InterpolationMode::BILINEAR:
1919
return "bilinear";
20+
case ResizeTransform::InterpolationMode::BICUBIC:
21+
return "bicubic";
2022
default:
2123
TORCH_CHECK(
2224
false,
@@ -29,6 +31,8 @@ int toSwsInterpolation(ResizeTransform::InterpolationMode mode) {
2931
switch (mode) {
3032
case ResizeTransform::InterpolationMode::BILINEAR:
3133
return SWS_BILINEAR;
34+
case ResizeTransform::InterpolationMode::BICUBIC:
35+
return SWS_BICUBIC;
3236
default:
3337
TORCH_CHECK(
3438
false,

src/torchcodec/_core/Transform.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ class Transform {
4747

4848
class ResizeTransform : public Transform {
4949
public:
50-
enum class InterpolationMode { BILINEAR };
50+
enum class InterpolationMode { BILINEAR, BICUBIC };
5151

5252
explicit ResizeTransform(const FrameDims& dims)
5353
: outputDims_(dims), interpolationMode_(InterpolationMode::BILINEAR) {}

src/torchcodec/_core/custom_ops.cpp

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -254,18 +254,39 @@ int checkedToNonNegativeInt(const std::string& str) {
254254

255255
// Resize transform specs take the form:
256256
//
257-
// "resize, <height>, <width>"
257+
// "resize, <height>, <width>" or "resize, <height>, <width>, <interpolation>"
258258
//
259-
// Where "resize" is the string literal and <height> and <width> are positive
260-
// integers.
259+
// Where "resize" is the string literal, <height> and <width> are positive
260+
// integers, and <interpolation> is an optional string (bilinear, bicubic, or
261+
// lanczos). If <interpolation> is not specified, it defaults to bilinear.
261262
Transform* makeResizeTransform(
262263
const std::vector<std::string>& resizeTransformSpec) {
263264
TORCH_CHECK(
264-
resizeTransformSpec.size() == 3,
265-
"resizeTransformSpec must have 3 elements including its name");
265+
resizeTransformSpec.size() >= 3 && resizeTransformSpec.size() <= 4,
266+
"resizeTransformSpec must have 3 or 4 elements including its name");
266267
int height = checkedToPositiveInt(resizeTransformSpec[1]);
267268
int width = checkedToPositiveInt(resizeTransformSpec[2]);
268-
return new ResizeTransform(FrameDims(height, width));
269+
270+
auto interpolationMode = ResizeTransform::InterpolationMode::BILINEAR;
271+
if (resizeTransformSpec.size() == 4) {
272+
const std::string& modeStr = resizeTransformSpec[3];
273+
// Trim leading/trailing whitespace
274+
auto trimmed = modeStr;
275+
trimmed.erase(0, trimmed.find_first_not_of(" \t"));
276+
trimmed.erase(trimmed.find_last_not_of(" \t") + 1);
277+
278+
if (trimmed == "bilinear") {
279+
interpolationMode = ResizeTransform::InterpolationMode::BILINEAR;
280+
} else if (trimmed == "bicubic") {
281+
interpolationMode = ResizeTransform::InterpolationMode::BICUBIC;
282+
} else {
283+
TORCH_CHECK(
284+
false,
285+
"Unknown interpolation mode: '" + trimmed +
286+
"'. Supported modes: bilinear, bicubic");
287+
}
288+
}
289+
return new ResizeTransform(FrameDims(height, width), interpolationMode);
269290
}
270291

271292
// Crop transform specs take the form:

src/torchcodec/transforms/_decoder_transforms.py

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -92,23 +92,33 @@ class Resize(DecoderTransform):
9292
"""Resize the decoded frame to a given size.
9393
9494
Complementary TorchVision transform: :class:`~torchvision.transforms.v2.Resize`.
95-
Interpolation is always bilinear. Anti-aliasing is always on.
95+
Anti-aliasing is always on.
9696
9797
Args:
9898
size (Sequence[int]): Desired output size. Must be a sequence of
9999
the form (height, width).
100+
interpolation (str): Interpolation mode. Must be one of "bilinear"
101+
or "bicubic". Default is "bilinear".
100102
"""
101103

102-
def __init__(self, size: Sequence[int]):
104+
_VALID_INTERPOLATIONS = {"bilinear", "bicubic"}
105+
106+
def __init__(self, size: Sequence[int], interpolation: str = "bilinear"):
103107
if len(size) != 2:
104108
raise ValueError(
105109
"Resize transform must have a (height, width) "
106110
f"pair for the size, got {size}."
107111
)
112+
if interpolation not in self._VALID_INTERPOLATIONS:
113+
raise ValueError(
114+
f"Invalid interpolation mode: '{interpolation}'. "
115+
f"Must be one of {self._VALID_INTERPOLATIONS}."
116+
)
108117
self.size = size
118+
self.interpolation = interpolation
109119

110120
def _make_transform_spec(self, input_dims: tuple[int | None, int | None]) -> str:
111-
return f"resize, {self.size[0]}, {self.size[1]}"
121+
return f"resize, {self.size[0]}, {self.size[1]}, {self.interpolation}"
112122

113123
def _get_output_dims(self) -> tuple[int | None, int | None] | None:
114124
return (self.size[0], self.size[1])
@@ -119,9 +129,16 @@ def _from_torchvision(cls, tv_resize: nn.Module):
119129

120130
assert isinstance(tv_resize, v2.Resize)
121131

122-
if tv_resize.interpolation is not v2.InterpolationMode.BILINEAR:
132+
# Map TorchVision interpolation modes to TorchCodec string modes
133+
interpolation_map = {
134+
v2.InterpolationMode.BILINEAR: "bilinear",
135+
v2.InterpolationMode.BICUBIC: "bicubic",
136+
}
137+
138+
if tv_resize.interpolation not in interpolation_map:
123139
raise ValueError(
124-
"TorchVision Resize transform must use bilinear interpolation."
140+
f"TorchVision Resize transform must use bilinear or bicubic "
141+
f"interpolation. Got: {tv_resize.interpolation}"
125142
)
126143
if tv_resize.antialias is False:
127144
raise ValueError(
@@ -134,7 +151,9 @@ def _from_torchvision(cls, tv_resize: nn.Module):
134151
"TorchVision Resize transform must have a (height, width) "
135152
f"pair for the size, got {tv_resize.size}."
136153
)
137-
return cls(size=tv_resize.size)
154+
155+
interpolation = interpolation_map[tv_resize.interpolation]
156+
return cls(size=tv_resize.size, interpolation=interpolation)
138157

139158

140159
class CenterCrop(DecoderTransform):

0 commit comments

Comments
 (0)