Skip to content

Commit

Permalink
vapoursynth: Respect YCbCr Matrix header
Browse files Browse the repository at this point in the history 
This probably leaks some memory in some of the error cases, but the next
commit will replace everything with proper RAII wrappers.
  • Loading branch information
@arch1t3cht
arch1t3cht committed Oct 25, 2024
1 parent be2e6cb commit 2b855a3
Showing 1 changed file with 107 additions and 61 deletions.
168 changes: 107 additions & 61 deletions src/video_provider_vs.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,17 +45,19 @@ namespace {
class VapourSynthVideoProvider: public VideoProvider {
VapourSynthWrapper vs;
VSScript *script = nullptr;
VSNode *node = nullptr;
VSNode *source_node = nullptr;
VSNode *prepared_node = nullptr;
const VSVideoInfo *vi = nullptr;

double dar = 0;
agi::vfr::Framerate fps;
std::vector<int> keyframes;
std::string colorspace;
std::string real_colorspace;
int video_cs = -1; // Reported or guessed color matrix of first frame
int video_cr = -1; // Reported or guessed color range of first frame
bool has_audio = false;

const VSFrame *GetVSFrame(int n);
const VSFrame *GetVSFrame(VSNode *node, int n);
void SetResizeArg(VSMap *args, const VSMap *props, const char *arg_name, const char *prop_name, int64_t deflt, int64_t unspecified = -1);

public:
Expand All @@ -64,46 +66,28 @@ class VapourSynthVideoProvider: public VideoProvider {

void GetFrame(int n, VideoFrame &frame) override;

void SetColorSpace(std::string const& matrix) override { }
void SetColorSpace(std::string const& matrix) override;

int GetFrameCount() const override { return vi->numFrames; }
agi::vfr::Framerate GetFPS() const override { return fps; }
int GetWidth() const override { return vi->width; }
int GetHeight() const override { return vi->height; }
double GetDAR() const override { return dar; }
std::vector<int> GetKeyFrames() const override { return keyframes; }
std::string GetColorSpace() const override { return GetRealColorSpace(); }
std::string GetRealColorSpace() const override { return colorspace == "Unknown" ? "None" : colorspace; }
std::string GetColorSpace() const override { return colorspace; }
std::string GetRealColorSpace() const override {
std::string result = ColorMatrix::colormatrix_description(video_cs, video_cr);
if (result == "") {
return "None";
}
return result;
}
bool HasAudio() const override { return has_audio; }
bool WantsCaching() const override { return true; }
std::string GetDecoderName() const override { return "VapourSynth"; }
bool ShouldSetVideoProperties() const override { return colorspace != "Unknown"; }
};

std::string colormatrix_description(int colorFamily, int colorRange, int matrix) {
if (colorFamily != cfYUV) {
return "None";
}
// Assuming TV for unspecified
std::string str = colorRange == VSC_RANGE_FULL ? "PC" : "TV";

switch (matrix) {
case VSC_MATRIX_RGB:
return "None";
case VSC_MATRIX_BT709:
return str + ".709";
case VSC_MATRIX_FCC:
return str + ".FCC";
case VSC_MATRIX_BT470_BG:
case VSC_MATRIX_ST170_M:
return str + ".601";
case VSC_MATRIX_ST240_M:
return str + ".240M";
default:
return "Unknown"; // Will return "None" in GetColorSpace
}
}

VapourSynthVideoProvider::VapourSynthVideoProvider(agi::fs::path const& filename, std::string const& colormatrix, agi::BackgroundRunner *br) try { try {
std::lock_guard<std::mutex> lock(vs.GetMutex());

Expand Down Expand Up @@ -137,14 +121,14 @@ VapourSynthVideoProvider::VapourSynthVideoProvider(agi::fs::path const& filename
std::string msg = agi::format("Error executing VapourSynth script: %s", vs.GetScriptAPI()->getError(script));
throw VapourSynthError(msg);
}
node = vs.GetScriptAPI()->getOutputNode(script, 0);
if (node == nullptr)
source_node = vs.GetScriptAPI()->getOutputNode(script, 0);
if (source_node == nullptr)
throw VapourSynthError("No output node set");

if (vs.GetAPI()->getNodeType(node) != mtVideo) {
if (vs.GetAPI()->getNodeType(source_node) != mtVideo) {
throw VapourSynthError("Output node isn't a video node");
}
vi = vs.GetAPI()->getVideoInfo(node);
vi = vs.GetAPI()->getVideoInfo(source_node);
if (vi == nullptr)
throw VapourSynthError("Couldn't get video info");
if (!vsh::isConstantVideoFormat(vi))
Expand Down Expand Up @@ -228,7 +212,7 @@ VapourSynthVideoProvider::VapourSynthVideoProvider(agi::fs::path const& filename
vs.GetAPI()->freeMap(clipinfo);

// Find the first frame Of the video to get some info
const VSFrame *frame = GetVSFrame(0);
const VSFrame *frame = GetVSFrame(source_node, 0);

const VSMap *props = vs.GetAPI()->getFramePropertiesRO(frame);
if (props == nullptr)
Expand All @@ -239,13 +223,76 @@ VapourSynthVideoProvider::VapourSynthVideoProvider(agi::fs::path const& filename
dar = double(vi->width * sarn) / (vi->height * sard);
}

int64_t range = vs.GetAPI()->mapGetInt(props, "_ColorRange", 0, &err1);
int64_t matrix = vs.GetAPI()->mapGetInt(props, "_Matrix", 0, &err2);
colorspace = colormatrix_description(vi->format.colorFamily, err1 == 0 ? range : -1, err2 == 0 ? matrix : -1);
int video_cr_vs = vs.GetAPI()->mapGetInt(props, "_ColorRange", 0, &err1);
switch (video_cr_vs) {
case VSC_RANGE_FULL:
video_cr = AGI_CR_JPEG;
case VSC_RANGE_LIMITED:
video_cr = AGI_CR_MPEG;
default:
video_cr = AGI_CR_UNSPECIFIED;
}
video_cs = vs.GetAPI()->mapGetInt(props, "_Matrix", 0, &err2);
ColorMatrix::guess_colorspace(video_cs, video_cr, vi->width, vi->height);

vs.GetAPI()->freeFrame(frame);

SetColorSpace(colormatrix);
} catch (VapourSynthError const& err) { // for try inside of function. We need both here since we need to catch errors from the VapourSynthWrap constructor.
if (prepared_node != nullptr)
vs.GetAPI()->freeNode(prepared_node);
if (source_node != nullptr)
vs.GetAPI()->freeNode(source_node);
if (script != nullptr)
vs.GetScriptAPI()->freeScript(script);
throw err;
}
}
catch (VapourSynthError const& err) { // for the entire constructor
throw VideoProviderError(agi::format("VapourSynth error: %s", err.GetMessage()));
}

void VapourSynthVideoProvider::SetColorSpace(std::string const& matrix) {
if (vi->format.colorFamily != cfRGB || vi->format.bitsPerSample != 8) {
if (matrix == colorspace && prepared_node != nullptr) {
return;
}
if (prepared_node != nullptr) {
vs.GetAPI()->freeNode(prepared_node);
}

VSNode *intermediary = nullptr;

auto [force_cs, force_cr] = ColorMatrix::parse_colormatrix(matrix);
if (force_cs != AGI_CS_UNSPECIFIED && force_cr != AGI_CR_UNSPECIFIED) {
// Override the _Matrix and _Range frame props to force the color space
VSPlugin *std = vs.GetAPI()->getPluginByID(VSH_STD_PLUGIN_ID, vs.GetScriptAPI()->getCore(script));
if (std == nullptr)
throw VapourSynthError("Couldn't find std plugin");

VSMap *args = vs.GetAPI()->createMap();
if (args == nullptr)
throw VapourSynthError("Failed to create argument map");

vs.GetAPI()->mapSetNode(args, "clip", source_node, maAppend);
vs.GetAPI()->mapSetInt(args, "_Matrix", force_cs, maAppend);
vs.GetAPI()->mapSetInt(args, "_ColorRange", force_cr == AGI_CR_JPEG ? VSC_RANGE_FULL : VSC_RANGE_LIMITED, maAppend);

VSMap *result = vs.GetAPI()->invoke(std, "SetFrameProps", args);
vs.GetAPI()->freeMap(args);
const char *error = vs.GetAPI()->mapGetError(result);
if (error) {
vs.GetAPI()->freeMap(result);
throw VideoOpenError(agi::format("Failed set color space frame props: %s", error));
}
int err;
intermediary = vs.GetAPI()->mapGetNode(result, "clip", 0, &err);
vs.GetAPI()->freeMap(result);
if (err) {
throw VideoOpenError("Failed to get SetFrameProps output node");
}
}

// Convert to RGB24 format
VSPlugin *resize = vs.GetAPI()->getPluginByID(VSH_RESIZE_PLUGIN_ID, vs.GetScriptAPI()->getCore(script));
if (resize == nullptr)
Expand All @@ -255,51 +302,47 @@ VapourSynthVideoProvider::VapourSynthVideoProvider(agi::fs::path const& filename
if (args == nullptr)
throw VapourSynthError("Failed to create argument map");

vs.GetAPI()->mapSetNode(args, "clip", node, maAppend);
vs.GetAPI()->mapSetNode(args, "clip", intermediary == nullptr ? source_node : intermediary, maAppend);
vs.GetAPI()->mapSetInt(args, "format", pfRGB24, maAppend);

// Set defaults for the colorspace parameters.
// If the video node has frame props (like if the video is tagged with
// some color space), these will override these arguments.
vs.GetAPI()->mapSetInt(args, "matrix_in", VSC_MATRIX_BT709, maAppend);
vs.GetAPI()->mapSetInt(args, "range_in", 0, maAppend);
vs.GetAPI()->mapSetInt(args, "matrix_in", video_cs, maAppend);
vs.GetAPI()->mapSetInt(args, "range_in", video_cr == AGI_CR_JPEG, maAppend);
vs.GetAPI()->mapSetInt(args, "chromaloc_in", VSC_CHROMA_LEFT, maAppend);

VSMap *result = vs.GetAPI()->invoke(resize, "Bicubic", args);
vs.GetAPI()->freeMap(args);
const char *error = vs.GetAPI()->mapGetError(result);
if (error) {
vs.GetAPI()->freeMap(result);
vs.GetAPI()->freeNode(node);
vs.GetScriptAPI()->freeScript(script);
throw VideoProviderError(agi::format("Failed to convert to RGB24: %s", error));
throw VideoOpenError(agi::format("Failed to convert to RGB24: %s", error));
}
int err;
vs.GetAPI()->freeNode(node);
node = vs.GetAPI()->mapGetNode(result, "clip", 0, &err);
prepared_node = vs.GetAPI()->mapGetNode(result, "clip", 0, &err);
vs.GetAPI()->freeMap(result);
if (err) {
vs.GetScriptAPI()->freeScript(script);
throw VideoProviderError("Failed to get resize output node");
throw VideoOpenError("Failed to get resize output node");
}

// Finally, try to get the first frame again, so if the filter does crash, it happens before loading finishes
const VSFrame *rgbframe = GetVSFrame(0);
const VSFrame *rgbframe = GetVSFrame(prepared_node, 0);
vs.GetAPI()->freeFrame(rgbframe);

if (intermediary != nullptr) {
vs.GetAPI()->freeNode(intermediary);
}
} else {
vs.GetAPI()->freeNode(prepared_node);
prepared_node = vs.GetAPI()->addNodeRef(source_node);
}
} catch (VapourSynthError const& err) { // for try inside of function. We need both here since we need to catch errors from the VapourSynthWrap constructor.
if (node != nullptr)
vs.GetAPI()->freeNode(node);
if (script != nullptr)
vs.GetScriptAPI()->freeScript(script);
throw err;
}
}
catch (VapourSynthError const& err) { // for the entire constructor
throw VideoProviderError(agi::format("VapourSynth error: %s", err.GetMessage()));
colorspace = matrix;
}

const VSFrame *VapourSynthVideoProvider::GetVSFrame(int n) {
const VSFrame *VapourSynthVideoProvider::GetVSFrame(VSNode *node, int n) {
char errorMsg[1024];
const VSFrame *frame = vs.GetAPI()->getFrame(n, node, errorMsg, sizeof(errorMsg));
if (frame == nullptr) {
Expand All @@ -311,7 +354,7 @@ const VSFrame *VapourSynthVideoProvider::GetVSFrame(int n) {
void VapourSynthVideoProvider::GetFrame(int n, VideoFrame &out) {
std::lock_guard<std::mutex> lock(vs.GetMutex());

const VSFrame *frame = GetVSFrame(n);
const VSFrame *frame = GetVSFrame(prepared_node, n);

const VSVideoFormat *format = vs.GetAPI()->getVideoFrameFormat(frame);
if (format->colorFamily != cfRGB || format->numPlanes != 3 || format->bitsPerSample != 8 || format->subSamplingH != 0 || format->subSamplingW != 0) {
Expand Down Expand Up @@ -348,8 +391,11 @@ void VapourSynthVideoProvider::GetFrame(int n, VideoFrame &out) {
}

VapourSynthVideoProvider::~VapourSynthVideoProvider() {
if (node != nullptr) {
vs.GetAPI()->freeNode(node);
if (prepared_node != nullptr) {
vs.GetAPI()->freeNode(prepared_node);
}
if (source_node != nullptr) {
vs.GetAPI()->freeNode(source_node);
}
if (script != nullptr) {
vs.GetScriptAPI()->freeScript(script);
Expand Down

0 comments on commit 2b855a3

Please sign in to comment.