Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 35 additions & 10 deletions torchvision/csrc/io/decoder/video_sampler.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -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,
Expand All @@ -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;
}

Expand Down