diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 14cf00fb..d298a5c4 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -98,10 +98,12 @@ set(OPENSHOT_CV_SOURCES CVStabilization.cpp ClipProcessingJobs.cpp CVObjectDetection.cpp + CVObjectMask.cpp TrackedObjectBBox.cpp effects/Stabilizer.cpp effects/Tracker.cpp effects/ObjectDetection.cpp + effects/ObjectMask.cpp effects/Outline.cpp ./sort_filter/sort.cpp ./sort_filter/Hungarian.cpp diff --git a/src/CVObjectMask.cpp b/src/CVObjectMask.cpp new file mode 100644 index 00000000..7e5d9f26 --- /dev/null +++ b/src/CVObjectMask.cpp @@ -0,0 +1,939 @@ +/** + * @file + * @brief Source file for CVObjectMask class + * @author Jonathan Thomas + * + * @ref License + */ + +// Copyright (c) 2026 OpenShot Studios, LLC +// +// SPDX-License-Identifier: LGPL-3.0-or-later + +#include "CVObjectMask.h" + +#include "Exceptions.h" +#include "objdetectdata.pb.h" + +#include +#include +#include +#include +#include +#include +#include + +#include + +using namespace openshot; +using google::protobuf::util::TimeUtil; + +namespace { + +std::string LoadONNXModel(const std::string& modelPath, cv::dnn::Net* net) +{ + try { + cv::dnn::Net loadedNet = cv::dnn::readNetFromONNX(modelPath); + if (net) + *net = loadedNet; + return ""; + } catch (const cv::Exception& e) { + return std::string("Failed to load ONNX model: ") + e.what(); + } catch (const std::exception& e) { + return std::string("Failed to load ONNX model: ") + e.what(); + } +} + +std::vector EncodeBinaryMaskRLE(const cv::Mat& mask) +{ + std::vector rle; + if (mask.empty()) + return rle; + + uint8_t current = 0; + uint32_t count = 0; + for (int y = 0; y < mask.rows; ++y) { + const uint8_t* row = mask.ptr(y); + for (int x = 0; x < mask.cols; ++x) { + const uint8_t value = row[x] ? 1 : 0; + if (value == current) { + ++count; + } else { + rle.push_back(count); + current = value; + count = 1; + } + } + } + rle.push_back(count); + return rle; +} + +struct SamPreprocessResult { + cv::Mat blob; + float scale = 1.0f; + int resizedWidth = 0; + int resizedHeight = 0; +}; + +SamPreprocessResult MakeSamBlob(const cv::Mat& bgr, int modelSize) +{ + SamPreprocessResult result; + result.scale = static_cast(modelSize) / static_cast(std::max(bgr.cols, bgr.rows)); + result.resizedWidth = static_cast(bgr.cols * result.scale + 0.5f); + result.resizedHeight = static_cast(bgr.rows * result.scale + 0.5f); + + cv::Mat resized; + cv::resize(bgr, resized, cv::Size(result.resizedWidth, result.resizedHeight), 0, 0, cv::INTER_LINEAR); + + const int shape[] = {1, 3, modelSize, modelSize}; + result.blob = cv::Mat(4, shape, CV_32F, cv::Scalar(0.0f)); + float* dst = result.blob.ptr(); + + const float mean[] = {123.675f, 116.28f, 103.53f}; + const float stddev[] = {58.395f, 57.12f, 57.375f}; + for (int y = 0; y < resized.rows; ++y) { + const cv::Vec3b* row = resized.ptr(y); + for (int x = 0; x < resized.cols; ++x) { + const float rgb[] = { + static_cast(row[x][2]), + static_cast(row[x][1]), + static_cast(row[x][0]), + }; + for (int c = 0; c < 3; ++c) + dst[(c * modelSize + y) * modelSize + x] = (rgb[c] - mean[c]) / stddev[c]; + } + } + + return result; +} + +cv::Rect_ NormalizedBoundingBox(const cv::Mat& mask) +{ + std::vector points; + cv::findNonZero(mask, points); + if (points.empty()) + return {}; + + cv::Rect rect = cv::boundingRect(points); + return cv::Rect_( + rect.x / static_cast(mask.cols), + rect.y / static_cast(mask.rows), + rect.width / static_cast(mask.cols), + rect.height / static_cast(mask.rows)); +} + +cv::Mat LowMaskToFrameMask(const cv::Mat& lowMask, const SamPreprocessResult& prep, + const cv::Size& frameSize, int modelSize, float maskThreshold) +{ + cv::Mat paddedMask; + cv::resize(lowMask, paddedMask, cv::Size(modelSize, modelSize), 0, 0, cv::INTER_LINEAR); + + cv::Mat cropped = paddedMask(cv::Rect(0, 0, prep.resizedWidth, prep.resizedHeight)); + cv::Mat fullSize; + cv::resize(cropped, fullSize, frameSize, 0, 0, cv::INTER_LINEAR); + + cv::Mat binary; + cv::threshold(fullSize, binary, maskThreshold, 255.0, cv::THRESH_BINARY); + binary.convertTo(binary, CV_8U); + return binary; +} + +CVObjectMaskFrameData FrameDataFromMask(const cv::Mat& mask, size_t frameId, float score) +{ + CVObjectMaskFrameData frameData; + frameData.frameId = frameId; + frameData.objectId = 1; + if (mask.empty()) + return frameData; + + frameData.score = score; + frameData.width = mask.cols; + frameData.height = mask.rows; + frameData.rle = EncodeBinaryMaskRLE(mask); + frameData.box = NormalizedBoundingBox(mask); + return frameData; +} + +cv::Point2f JsonPoint(const Json::Value& value) +{ + if (!value.isObject() || value["x"].isNull() || value["y"].isNull()) + return cv::Point2f(-1.0f, -1.0f); + return cv::Point2f(value["x"].asFloat(), value["y"].asFloat()); +} + +bool IsValidPoint(const cv::Point2f& point) +{ + return point.x >= 0.0f && point.y >= 0.0f; +} + +void AppendJsonPoints(const Json::Value& values, std::vector& points) +{ + if (!values.isArray()) + return; + for (const auto& value : values) { + cv::Point2f point = JsonPoint(value); + if (IsValidPoint(point)) + points.push_back(point); + } +} + +size_t JsonFrameNumber(const std::string& frameName) +{ + try { + return static_cast(std::max(0, std::stoi(frameName))); + } catch (...) { + return 0; + } +} + +void ApplyRectJson(const Json::Value& rect, CVObjectMaskPromptSet& prompts) +{ + if (!rect.isObject() || rect["x1"].isNull() || rect["y1"].isNull() || + rect["x2"].isNull() || rect["y2"].isNull()) { + return; + } + + prompts.rectTopLeft.x = std::min(rect["x1"].asFloat(), rect["x2"].asFloat()); + prompts.rectTopLeft.y = std::min(rect["y1"].asFloat(), rect["y2"].asFloat()); + prompts.rectBottomRight.x = std::max(rect["x1"].asFloat(), rect["x2"].asFloat()); + prompts.rectBottomRight.y = std::max(rect["y1"].asFloat(), rect["y2"].asFloat()); + prompts.hasRect = IsValidPoint(prompts.rectTopLeft) && IsValidPoint(prompts.rectBottomRight); +} + +void AppendNegativeRectCenters(const Json::Value& values, CVObjectMaskPromptSet& prompts) +{ + if (!values.isArray()) + return; + for (const auto& rect : values) { + if (!rect.isObject() || rect["x1"].isNull() || rect["y1"].isNull() || + rect["x2"].isNull() || rect["y2"].isNull()) { + continue; + } + cv::Point2f center( + (rect["x1"].asFloat() + rect["x2"].asFloat()) / 2.0f, + (rect["y1"].asFloat() + rect["y2"].asFloat()) / 2.0f); + if (IsValidPoint(center)) + prompts.negativePoints.push_back(center); + } +} + +CVObjectMaskPromptSet PromptSetFromJson(const Json::Value& framePayload) +{ + CVObjectMaskPromptSet prompts; + AppendJsonPoints(framePayload["positive_points"], prompts.positivePoints); + AppendJsonPoints(framePayload["negative_points"], prompts.negativePoints); + if (framePayload["positive_rects"].isArray() && framePayload["positive_rects"].size() > 0) + ApplyRectJson(framePayload["positive_rects"][0], prompts); + AppendNegativeRectCenters(framePayload["negative_rects"], prompts); + return prompts; +} + +cv::Mat MakeXMemImageBlob(const cv::Mat& bgr) +{ + cv::Mat resized; + cv::resize(bgr, resized, cv::Size(640, 480), 0, 0, cv::INTER_LINEAR); + + const int shape[] = {1, 3, 480, 640}; + cv::Mat blob(4, shape, CV_32F); + float* dst = blob.ptr(); + const float mean[] = {0.485f, 0.456f, 0.406f}; + const float stddev[] = {0.229f, 0.224f, 0.225f}; + + for (int y = 0; y < resized.rows; ++y) { + const cv::Vec3b* row = resized.ptr(y); + for (int x = 0; x < resized.cols; ++x) { + const float rgb[] = { + static_cast(row[x][2]) / 255.0f, + static_cast(row[x][1]) / 255.0f, + static_cast(row[x][0]) / 255.0f, + }; + for (int c = 0; c < 3; ++c) + dst[(c * 480 + y) * 640 + x] = (rgb[c] - mean[c]) / stddev[c]; + } + } + + return blob; +} + +cv::Mat MakeScalarBlob(float value) +{ + const int shape[] = {1}; + return cv::Mat(1, shape, CV_32F, cv::Scalar(value)); +} + +cv::Mat MakeXMemMaskBlob(const cv::Mat& mask) +{ + cv::Mat resized; + cv::resize(mask, resized, cv::Size(640, 480), 0, 0, cv::INTER_NEAREST); + const int shape[] = {1, 1, 480, 640}; + cv::Mat blob(4, shape, CV_32F, cv::Scalar(0.0f)); + float* dst = blob.ptr(); + for (int y = 0; y < resized.rows; ++y) { + const uint8_t* row = resized.ptr(y); + for (int x = 0; x < resized.cols; ++x) + dst[y * 640 + x] = row[x] ? 1.0f : 0.0f; + } + return blob; +} + +cv::Mat BinaryMaskFromXMemProb(const cv::Mat& prob) +{ + cv::Mat mask(480, 640, CV_8U, cv::Scalar(0)); + const float* src = prob.ptr(); + for (int y = 0; y < mask.rows; ++y) { + uint8_t* row = mask.ptr(y); + for (int x = 0; x < mask.cols; ++x) + row[x] = src[y * 640 + x] >= 0.5f ? 255 : 0; + } + return mask; +} + +cv::Mat AggregateXMemForegroundProb(const cv::Mat& rawProb) +{ + const int shape[] = {1, 1, 480, 640}; + cv::Mat output(4, shape, CV_32F); + const float* src = rawProb.ptr(); + float* dst = output.ptr(); + for (int i = 0; i < 480 * 640; ++i) { + const float fg = std::min(1.0f - 1e-7f, std::max(1e-7f, src[i])); + const float bg = std::min(1.0f - 1e-7f, std::max(1e-7f, 1.0f - fg)); + const float bgLogit = std::log(bg / (1.0f - bg)); + const float fgLogit = std::log(fg / (1.0f - fg)); + const float maxLogit = std::max(bgLogit, fgLogit); + const float bgExp = std::exp(bgLogit - maxLogit); + const float fgExp = std::exp(fgLogit - maxLogit); + dst[i] = fgExp / (bgExp + fgExp); + } + return output; +} + +cv::Mat FlattenFeature(const cv::Mat& feature, int channels) +{ + return cv::Mat(channels, 30 * 40, CV_32F, const_cast(feature.ptr())).clone(); +} + +cv::Mat FlattenShrinkage(const cv::Mat& shrinkage) +{ + return cv::Mat(1, 30 * 40, CV_32F, const_cast(shrinkage.ptr())).clone(); +} + +cv::Mat MatFromReadout(const cv::Mat& readout) +{ + const int shape[] = {1, 1, 512, 30, 40}; + cv::Mat output(5, shape, CV_32F); + std::memcpy(output.ptr(), readout.ptr(), sizeof(float) * readout.total()); + return output; +} + +class XMemPropagator { +private: + struct MemoryFrame { + cv::Mat key; + cv::Mat shrinkage; + cv::Mat value; + }; + + cv::dnn::Net encodeKey; + cv::dnn::Net encodeValue; + cv::dnn::Net decode; + cv::Mat memoryKey; + cv::Mat memoryShrinkage; + cv::Mat memoryValue; + cv::Mat hidden; + std::deque memoryFrames; + int frameIndex = 0; + int lastMemoryFrame = -1000000; + int memEvery = 5; + int maxMemoryFrames = 10; + + void RebuildMemory() + { + std::vector keys; + std::vector shrinkages; + std::vector values; + for (const auto& frame : memoryFrames) { + keys.push_back(frame.key); + shrinkages.push_back(frame.shrinkage); + values.push_back(frame.value); + } + if (keys.empty()) { + memoryKey.release(); + memoryShrinkage.release(); + memoryValue.release(); + return; + } + cv::hconcat(keys, memoryKey); + cv::hconcat(shrinkages, memoryShrinkage); + cv::hconcat(values, memoryValue); + } + + cv::Mat MatchMemory(const cv::Mat& queryKey, const cv::Mat& selection) + { + cv::Mat query = FlattenFeature(queryKey, 64); + cv::Mat querySelection = FlattenFeature(selection, 64); + cv::Mat weightedQuery = query.mul(querySelection); + + cv::Mat twoAb; + cv::gemm(memoryKey, weightedQuery, 2.0, cv::Mat(), 0.0, twoAb, cv::GEMM_1_T); + + cv::Mat aSq; + cv::gemm(memoryKey.mul(memoryKey), querySelection, 1.0, cv::Mat(), 0.0, aSq, cv::GEMM_1_T); + + cv::Mat querySquared = query.mul(query).mul(querySelection); + std::vector bSq(query.cols, 0.0f); + for (int q = 0; q < query.cols; ++q) { + double sum = 0.0; + for (int c = 0; c < query.rows; ++c) + sum += querySquared.at(c, q); + bSq[q] = static_cast(sum); + } + + cv::Mat similarity = twoAb - aSq; + const float invSqrtKeyDim = 1.0f / std::sqrt(64.0f); + for (int n = 0; n < similarity.rows; ++n) { + float* row = similarity.ptr(n); + const float shrinkage = memoryShrinkage.at(0, n); + for (int q = 0; q < similarity.cols; ++q) + row[q] = (row[q] - bSq[q]) * shrinkage * invSqrtKeyDim; + } + + const int topK = std::min(30, similarity.rows); + for (int q = 0; q < similarity.cols; ++q) { + std::vector indices(similarity.rows); + std::iota(indices.begin(), indices.end(), 0); + std::partial_sort(indices.begin(), indices.begin() + topK, indices.end(), + [&](int a, int b) { return similarity.at(a, q) > similarity.at(b, q); }); + + float maxValue = similarity.at(indices[0], q); + for (int k = 1; k < topK; ++k) + maxValue = std::max(maxValue, similarity.at(indices[k], q)); + + double sum = 0.0; + std::vector values(topK, 0.0f); + for (int k = 0; k < topK; ++k) { + values[k] = std::exp(similarity.at(indices[k], q) - maxValue); + sum += values[k]; + } + if (sum <= 0.0) + continue; + for (int n = 0; n < similarity.rows; ++n) + similarity.at(n, q) = 0.0f; + for (int k = 0; k < topK; ++k) + similarity.at(indices[k], q) = static_cast(values[k] / sum); + } + + cv::Mat readout; + cv::gemm(memoryValue, similarity, 1.0, cv::Mat(), 0.0, readout); + return MatFromReadout(readout); + } + + void AddMemory(const cv::Mat& key, const cv::Mat& shrinkage, const cv::Mat& value) + { + MemoryFrame frame; + frame.key = FlattenFeature(key, 64); + frame.shrinkage = FlattenShrinkage(shrinkage); + frame.value = FlattenFeature(value, 512); + memoryFrames.push_back(frame); + while (static_cast(memoryFrames.size()) > maxMemoryFrames) + memoryFrames.pop_front(); + RebuildMemory(); + } + + void EnsureHidden() + { + if (!hidden.empty()) + return; + const int shape[] = {1, 1, 64, 30, 40}; + hidden = cv::Mat(5, shape, CV_32F, cv::Scalar(0.0f)); + } + +public: + void Load(const std::string& encodeKeyPath, const std::string& encodeValuePath, const std::string& decodePath) + { + encodeKey = cv::dnn::readNetFromONNX(encodeKeyPath); + encodeValue = cv::dnn::readNetFromONNX(encodeValuePath); + decode = cv::dnn::readNetFromONNX(decodePath); + } + + void SetDevice(const std::string& processingDevice) + { + if (processingDevice == "GPU") { + try { + const std::vector targets = cv::dnn::getAvailableTargets(cv::dnn::DNN_BACKEND_CUDA); + if (std::find(targets.begin(), targets.end(), cv::dnn::DNN_TARGET_CUDA) != targets.end()) { + encodeKey.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA); + encodeKey.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA); + encodeValue.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA); + encodeValue.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA); + decode.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA); + decode.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA); + return; + } + } catch (const cv::Exception&) { + } + } + + encodeKey.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV); + encodeKey.setPreferableTarget(cv::dnn::DNN_TARGET_CPU); + encodeValue.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV); + encodeValue.setPreferableTarget(cv::dnn::DNN_TARGET_CPU); + decode.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV); + decode.setPreferableTarget(cv::dnn::DNN_TARGET_CPU); + } + + void Reset() + { + memoryFrames.clear(); + memoryKey.release(); + memoryShrinkage.release(); + memoryValue.release(); + hidden.release(); + frameIndex = 0; + lastMemoryFrame = -1000000; + } + + bool HasMemory() const + { + return !memoryFrames.empty(); + } + + cv::Mat Step(const cv::Mat& frame, const cv::Mat& seedMask = cv::Mat()) + { + cv::Mat image = MakeXMemImageBlob(frame); + + encodeKey.setInput(image, "image"); + encodeKey.setInput(MakeScalarBlob(1.0f), "need_sk"); + encodeKey.setInput(MakeScalarBlob(1.0f), "need_ek"); + std::vector keyOutputs; + encodeKey.forward(keyOutputs, std::vector{"key", "shrinkage", "selection", "f16", "f8", "f4"}); + cv::Mat key = keyOutputs[0]; + cv::Mat shrinkage = keyOutputs[1]; + cv::Mat f16 = keyOutputs[3]; + cv::Mat f8 = keyOutputs[4]; + cv::Mat f4 = keyOutputs[5]; + + EnsureHidden(); + + cv::Mat modelMask; + if (!seedMask.empty()) { + modelMask = MakeXMemMaskBlob(seedMask); + } else if (HasMemory()) { + cv::Mat memoryReadout = MatchMemory(key, keyOutputs[2]); + decode.setInput(f16, "f16"); + decode.setInput(f8, "f8"); + decode.setInput(f4, "f4"); + decode.setInput(hidden, "h16_in"); + decode.setInput(memoryReadout, "memory_readout"); + decode.setInput(MakeScalarBlob(1.0f), "h_out"); + std::vector decodeOutputs; + decode.forward(decodeOutputs, std::vector{"h16_out", "logits", "prob"}); + hidden = decodeOutputs[0].clone(); + modelMask = AggregateXMemForegroundProb(decodeOutputs[2]); + } else { + ++frameIndex; + return cv::Mat(); + } + + const bool isMemoryFrame = !seedMask.empty() || frameIndex - lastMemoryFrame >= memEvery; + if (isMemoryFrame) { + const int othersShape[] = {1, 1, 480, 640}; + cv::Mat others(4, othersShape, CV_32F, cv::Scalar(0.0f)); + encodeValue.setInput(image, "image"); + encodeValue.setInput(f16, "f16"); + encodeValue.setInput(hidden, "h16_in"); + encodeValue.setInput(modelMask, "masks"); + encodeValue.setInput(others, "others"); + encodeValue.setInput(MakeScalarBlob(1.0f), "is_deep_update"); + std::vector valueOutputs; + encodeValue.forward(valueOutputs, std::vector{"g16", "h16_out"}); + hidden = valueOutputs[1].clone(); + AddMemory(key, shrinkage, valueOutputs[0]); + lastMemoryFrame = frameIndex; + } + + cv::Mat outputMask = BinaryMaskFromXMemProb(modelMask); + ++frameIndex; + return outputMask; + } +}; + +} + +CVObjectMask::CVObjectMask(std::string processInfoJson, ProcessingController& controller) + : processingController(&controller) +{ + SetJson(processInfoJson); +} + +std::string CVObjectMask::ValidateONNXModels(std::string encoderPath, std::string decoderPath) +{ + std::string error = LoadONNXModel(encoderPath, nullptr); + if (!error.empty()) + return error; + return LoadONNXModel(decoderPath, nullptr); +} + +void CVObjectMask::SetProcessingDevice() +{ + if (processingDevice == "GPU") { + try { + const std::vector targets = cv::dnn::getAvailableTargets(cv::dnn::DNN_BACKEND_CUDA); + if (std::find(targets.begin(), targets.end(), cv::dnn::DNN_TARGET_CUDA) != targets.end()) { + encoder.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA); + encoder.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA); + decoder.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA); + decoder.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA); + return; + } + } catch (const cv::Exception&) { + } + processingDevice = "CPU"; + } + + encoder.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV); + encoder.setPreferableTarget(cv::dnn::DNN_TARGET_CPU); + decoder.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV); + decoder.setPreferableTarget(cv::dnn::DNN_TARGET_CPU); +} + +void CVObjectMask::maskClip(openshot::Clip& video, size_t _start, size_t _end, bool process_interval) +{ + start = _start; + end = _end; + + video.Open(); + processingController->SetError(false, ""); + + if (encoderModelPath.empty() || decoderModelPath.empty()) { + processingController->SetError(true, "Missing path to EdgeSAM encoder or decoder ONNX model file"); + error = true; + return; + } + if (protobufDataPath.empty()) { + processingController->SetError(true, "Missing path to object mask protobuf data file"); + error = true; + return; + } + if (promptKeyframes.empty()) { + processingController->SetError(true, "Missing positive prompt point for Object Mask preprocessing"); + error = true; + return; + } + + std::string loadError = LoadONNXModel(encoderModelPath, &encoder); + if (!loadError.empty()) { + processingController->SetError(true, loadError); + error = true; + return; + } + loadError = LoadONNXModel(decoderModelPath, &decoder); + if (!loadError.empty()) { + processingController->SetError(true, loadError); + error = true; + return; + } + SetProcessingDevice(); + + if (xmemEncodeKeyModelPath.empty() && !xmemModelDir.empty()) + xmemEncodeKeyModelPath = xmemModelDir + "/XMem-encode_key.onnx"; + if (xmemEncodeValueModelPath.empty() && !xmemModelDir.empty()) + xmemEncodeValueModelPath = xmemModelDir + "/XMem-encode_value-m1.onnx"; + if (xmemDecodeModelPath.empty() && !xmemModelDir.empty()) + xmemDecodeModelPath = xmemModelDir + "/XMem-decode-m1.onnx"; + if (xmemEncodeKeyModelPath.empty() || xmemEncodeValueModelPath.empty() || xmemDecodeModelPath.empty()) { + processingController->SetError(true, "Missing path to XMem ONNX model files"); + error = true; + return; + } + + XMemPropagator xmem; + try { + xmem.Load(xmemEncodeKeyModelPath, xmemEncodeValueModelPath, xmemDecodeModelPath); + xmem.SetDevice(processingDevice); + } catch (const cv::Exception& e) { + processingController->SetError(true, std::string("Failed to load XMem ONNX models: ") + e.what()); + error = true; + return; + } catch (const std::exception& e) { + processingController->SetError(true, std::string("Failed to load XMem ONNX models: ") + e.what()); + error = true; + return; + } + + if (!process_interval || end <= 1 || end - start == 0) { + start = static_cast(video.Start() * video.Reader()->info.fps.ToFloat()); + end = static_cast(video.End() * video.Reader()->info.fps.ToFloat()); + } + if (end < start) + end = start; + + CVObjectMaskPromptSet activePrompts; + auto promptBeforeStart = promptKeyframes.upper_bound(start); + if (promptBeforeStart != promptKeyframes.begin()) { + --promptBeforeStart; + activePrompts = promptBeforeStart->second; + } + auto firstPromptAtOrAfterStart = promptKeyframes.lower_bound(start); + + for (size_t frameNumber = start; frameNumber <= end; ++frameNumber) { + if (processingController->ShouldStop()) + return; + + std::shared_ptr frame = video.GetFrame(frameNumber); + if (!frame) + continue; + + auto promptIt = promptKeyframes.find(frameNumber); + bool isPromptKeyframe = promptIt != promptKeyframes.end(); + if (promptIt != promptKeyframes.end()) { + activePrompts = promptIt->second; + xmem.Reset(); + } else if (!activePrompts.HasPositivePrompt()) { + if (firstPromptAtOrAfterStart != promptKeyframes.end() && frameNumber >= firstPromptAtOrAfterStart->first) { + activePrompts = firstPromptAtOrAfterStart->second; + isPromptKeyframe = true; + xmem.Reset(); + } else { + CVObjectMaskFrameData emptyFrame; + emptyFrame.frameId = frameNumber; + masksData[frameNumber] = emptyFrame; + continue; + } + } + + const cv::Mat frameImage = frame->GetImageCV(); + cv::Mat seedMask; + if (isPromptKeyframe || !xmem.HasMemory()) { + seedMask = CreateEdgeSAMSeedMask(frameImage, activePrompts); + if (seedMask.empty()) { + CVObjectMaskFrameData emptyFrame; + emptyFrame.frameId = frameNumber; + masksData[frameNumber] = emptyFrame; + continue; + } + if (!isPromptKeyframe) + xmem.Reset(); + } + + cv::Mat propagatedMask; + try { + propagatedMask = xmem.Step(frameImage, seedMask); + } catch (const cv::Exception& e) { + processingController->SetError(true, std::string("Failed to propagate Object Mask with XMem: ") + e.what()); + error = true; + return; + } + + cv::Mat outputMask; + if (!propagatedMask.empty()) + cv::resize(propagatedMask, outputMask, frameImage.size(), 0, 0, cv::INTER_NEAREST); + masksData[frameNumber] = FrameDataFromMask(outputMask, frameNumber, 1.0f); + + const size_t range = std::max(1, end - start); + processingController->SetProgress(uint(100 * (frameNumber - start) / range)); + } +} + +cv::Mat CVObjectMask::CreateEdgeSAMSeedMask(const cv::Mat& frame, const CVObjectMaskPromptSet& prompts) +{ + SamPreprocessResult prep = MakeSamBlob(frame, modelSize); + encoder.setInput(prep.blob, "image"); + cv::Mat embeddings = encoder.forward("image_embeddings"); + + const int coordsShape[] = {1, promptSlots, 2}; + const int labelsShape[] = {1, promptSlots}; + cv::Mat pointCoords(3, coordsShape, CV_32F, cv::Scalar(0.0f)); + cv::Mat pointLabels(2, labelsShape, CV_32F, cv::Scalar(-1.0f)); + + int promptIndex = 0; + if (prompts.hasRect && promptSlots >= 2) { + float* coords = pointCoords.ptr(); + float* labels = pointLabels.ptr(); + coords[0] = prompts.rectTopLeft.x * prep.scale; + coords[1] = prompts.rectTopLeft.y * prep.scale; + labels[0] = 2.0f; + coords[2] = prompts.rectBottomRight.x * prep.scale; + coords[3] = prompts.rectBottomRight.y * prep.scale; + labels[1] = 3.0f; + promptIndex = 2; + } + for (const auto& point : prompts.positivePoints) { + if (promptIndex >= promptSlots) + break; + pointCoords.ptr()[promptIndex * 2] = point.x * prep.scale; + pointCoords.ptr()[promptIndex * 2 + 1] = point.y * prep.scale; + pointLabels.ptr()[promptIndex] = 1.0f; + ++promptIndex; + } + for (const auto& point : prompts.negativePoints) { + if (promptIndex >= promptSlots) + break; + pointCoords.ptr()[promptIndex * 2] = point.x * prep.scale; + pointCoords.ptr()[promptIndex * 2 + 1] = point.y * prep.scale; + pointLabels.ptr()[promptIndex] = 0.0f; + ++promptIndex; + } + + decoder.setInput(embeddings, "image_embeddings"); + decoder.setInput(pointCoords, "point_coords"); + decoder.setInput(pointLabels, "point_labels"); + + std::vector outputs; + decoder.forward(outputs, std::vector{"scores", "masks"}); + if (outputs.size() != 2) + return cv::Mat(); + + const float* scores = outputs[0].ptr(); + const int maskCount = static_cast(outputs[0].total()); + int bestScoreMask = 0; + for (int i = 1; i < maskCount; ++i) { + if (scores[i] > scores[bestScoreMask]) + bestScoreMask = i; + } + + cv::Mat lowMask(maskSize, maskSize, CV_32F, outputs[1].ptr(0, bestScoreMask)); + return LowMaskToFrameMask(lowMask, prep, frame.size(), modelSize, maskThreshold); +} + +bool CVObjectMask::SaveObjMaskData() +{ + if (protobufDataPath.empty()) { + std::cerr << "Missing path to object mask protobuf data file." << std::endl; + return false; + } + if (error) + return false; + + pb_objdetect::ObjDetect objMessage; + objMessage.add_classnames()->assign("object mask"); + + for (const auto& frameData : masksData) + AddFrameDataToProto(objMessage.add_frame(), frameData.second); + + *objMessage.mutable_last_updated() = TimeUtil::SecondsToTimestamp(time(NULL)); + + std::fstream output(protobufDataPath, std::ios::out | std::ios::trunc | std::ios::binary); + if (!objMessage.SerializeToOstream(&output)) { + std::cerr << "Failed to write object mask protobuf message." << std::endl; + return false; + } + + google::protobuf::ShutdownProtobufLibrary(); + return true; +} + +void CVObjectMask::AddFrameDataToProto(pb_objdetect::Frame* pbFrameData, const CVObjectMaskFrameData& frameData) +{ + pbFrameData->set_id(frameData.frameId); + if (!frameData.HasMask()) + return; + + pb_objdetect::Frame_Box* box = pbFrameData->add_bounding_box(); + box->set_x(frameData.box.x); + box->set_y(frameData.box.y); + box->set_w(frameData.box.width); + box->set_h(frameData.box.height); + box->set_classid(0); + box->set_confidence(frameData.score); + box->set_objectid(frameData.objectId); + + pb_objdetect::Frame_Box_Mask* mask = box->mutable_mask(); + mask->set_width(frameData.width); + mask->set_height(frameData.height); + for (uint32_t count : frameData.rle) + mask->add_rle(count); +} + +void CVObjectMask::SetJson(const std::string value) +{ + try { + SetJsonValue(openshot::stringToJson(value)); + } catch (const std::exception&) { + std::cout << "JSON is invalid (missing keys or invalid data types)" << std::endl; + } +} + +void CVObjectMask::SetJsonValue(const Json::Value root) +{ + if (!root["protobuf_data_path"].isNull()) + protobufDataPath = root["protobuf_data_path"].asString(); + if (!root["encoder_model"].isNull()) + encoderModelPath = root["encoder_model"].asString(); + if (!root["encoder_model_path"].isNull()) + encoderModelPath = root["encoder_model_path"].asString(); + if (!root["decoder_model"].isNull()) + decoderModelPath = root["decoder_model"].asString(); + if (!root["decoder_model_path"].isNull()) + decoderModelPath = root["decoder_model_path"].asString(); + if (!root["xmem_model_dir"].isNull()) + xmemModelDir = root["xmem_model_dir"].asString(); + if (!root["xmem_encode_key_model"].isNull()) + xmemEncodeKeyModelPath = root["xmem_encode_key_model"].asString(); + if (!root["xmem_encode_key_model_path"].isNull()) + xmemEncodeKeyModelPath = root["xmem_encode_key_model_path"].asString(); + if (!root["xmem_encode_value_model"].isNull()) + xmemEncodeValueModelPath = root["xmem_encode_value_model"].asString(); + if (!root["xmem_encode_value_model_path"].isNull()) + xmemEncodeValueModelPath = root["xmem_encode_value_model_path"].asString(); + if (!root["xmem_decode_model"].isNull()) + xmemDecodeModelPath = root["xmem_decode_model"].asString(); + if (!root["xmem_decode_model_path"].isNull()) + xmemDecodeModelPath = root["xmem_decode_model_path"].asString(); + if (!root["processing-device"].isNull()) + processingDevice = root["processing-device"].asString(); + if (!root["processing_device"].isNull()) + processingDevice = root["processing_device"].asString(); + if (!root["prompt_slots"].isNull()) + promptSlots = std::max(1, root["prompt_slots"].asInt()); + if (!root["mask_threshold"].isNull()) + maskThreshold = root["mask_threshold"].asFloat(); + if (!root["model_size"].isNull()) + modelSize = root["model_size"].asInt(); + if (!root["mask_size"].isNull()) + maskSize = root["mask_size"].asInt(); + + promptKeyframes.clear(); + if (!root["object_mask_selection"].isNull()) { + const Json::Value& selection = root["object_mask_selection"]; + const Json::Value& frames = selection["frames"]; + if (frames.isObject()) { + for (const auto& frameName : frames.getMemberNames()) { + const size_t frameNumber = JsonFrameNumber(frameName); + if (frameNumber == 0) + continue; + CVObjectMaskPromptSet prompts = PromptSetFromJson(frames[frameName]); + if (prompts.HasPositivePrompt()) + promptKeyframes[frameNumber] = prompts; + } + } + } + + CVObjectMaskPromptSet legacyPrompts; + if (!root["positive_points"].isNull()) + AppendJsonPoints(root["positive_points"], legacyPrompts.positivePoints); + if (!root["negative_points"].isNull()) + AppendJsonPoints(root["negative_points"], legacyPrompts.negativePoints); + + if (!root["positive_x"].isNull() && !root["positive_y"].isNull()) { + cv::Point2f point(root["positive_x"].asFloat(), root["positive_y"].asFloat()); + if (IsValidPoint(point) && legacyPrompts.positivePoints.empty()) + legacyPrompts.positivePoints.push_back(point); + } + if (!root["negative_x"].isNull() && !root["negative_y"].isNull()) { + cv::Point2f point(root["negative_x"].asFloat(), root["negative_y"].asFloat()); + if (IsValidPoint(point) && legacyPrompts.negativePoints.empty()) + legacyPrompts.negativePoints.push_back(point); + } + if (!root["rect_x1"].isNull() && !root["rect_y1"].isNull() && + !root["rect_x2"].isNull() && !root["rect_y2"].isNull()) { + Json::Value rect; + rect["x1"] = root["rect_x1"]; + rect["y1"] = root["rect_y1"]; + rect["x2"] = root["rect_x2"]; + rect["y2"] = root["rect_y2"]; + ApplyRectJson(rect, legacyPrompts); + } + if (legacyPrompts.HasPositivePrompt() && promptKeyframes.empty()) + promptKeyframes[1] = legacyPrompts; +} diff --git a/src/CVObjectMask.h b/src/CVObjectMask.h new file mode 100644 index 00000000..ba39fe69 --- /dev/null +++ b/src/CVObjectMask.h @@ -0,0 +1,102 @@ +/** + * @file + * @brief Header file for CVObjectMask class + * @author Jonathan Thomas + * + * @ref License + */ + +// Copyright (c) 2026 OpenShot Studios, LLC +// +// SPDX-License-Identifier: LGPL-3.0-or-later + +#pragma once + +#define int64 opencv_broken_int +#define uint64 opencv_broken_uint +#include +#include +#include +#undef uint64 +#undef int64 + +#include "Clip.h" +#include "Json.h" +#include "ProcessingController.h" + +namespace pb_objdetect { + class Frame; +} + +namespace openshot +{ + struct CVObjectMaskFrameData { + size_t frameId = 0; + cv::Rect_ box; + float score = 0.0f; + int objectId = 1; + int width = 0; + int height = 0; + std::vector rle; + + bool HasMask() const { return width > 0 && height > 0 && !rle.empty(); } + }; + + struct CVObjectMaskPromptSet { + std::vector positivePoints; + std::vector negativePoints; + cv::Point2f rectTopLeft = cv::Point2f(-1.0f, -1.0f); + cv::Point2f rectBottomRight = cv::Point2f(-1.0f, -1.0f); + bool hasRect = false; + + bool HasPositivePrompt() const { return hasRect || !positivePoints.empty(); } + }; + + /** + * @brief Preprocess a clip into EdgeSAM object masks stored in the object-detection protobuf format. + */ + class CVObjectMask + { + private: + cv::dnn::Net encoder; + cv::dnn::Net decoder; + + std::string encoderModelPath; + std::string decoderModelPath; + std::string xmemModelDir; + std::string xmemEncodeKeyModelPath; + std::string xmemEncodeValueModelPath; + std::string xmemDecodeModelPath; + std::string protobufDataPath; + std::string processingDevice = "CPU"; + + std::map promptKeyframes; + int promptSlots = 10; + float maskThreshold = 0.0f; + int modelSize = 1024; + int maskSize = 256; + + size_t start = 0; + size_t end = 0; + bool error = false; + + ProcessingController* processingController; + + void SetProcessingDevice(); + cv::Mat CreateEdgeSAMSeedMask(const cv::Mat& frame, const CVObjectMaskPromptSet& prompts); + void AddFrameDataToProto(pb_objdetect::Frame* pbFrameData, const CVObjectMaskFrameData& frameData); + + public: + std::map masksData; + + CVObjectMask(std::string processInfoJson, ProcessingController& processingController); + + static std::string ValidateONNXModels(std::string encoderPath, std::string decoderPath); + + void maskClip(openshot::Clip& video, size_t start = 0, size_t end = 0, bool process_interval = false); + bool SaveObjMaskData(); + + void SetJson(const std::string value); + void SetJsonValue(const Json::Value root); + }; +} diff --git a/src/ClipProcessingJobs.cpp b/src/ClipProcessingJobs.cpp index 19a237ab..ed34e540 100644 --- a/src/ClipProcessingJobs.cpp +++ b/src/ClipProcessingJobs.cpp @@ -41,6 +41,9 @@ void ClipProcessingJobs::processClip(Clip& clip, std::string json){ if(processingType == "ObjectDetection"){ t = std::thread(&ClipProcessingJobs::detectObjectsClip, this, std::ref(clip), std::ref(this->processingController)); } + if(processingType == "ObjectMask"){ + t = std::thread(&ClipProcessingJobs::maskObjectClip, this, std::ref(clip), std::ref(this->processingController)); + } } // Apply object tracking to clip @@ -87,6 +90,21 @@ void ClipProcessingJobs::detectObjectsClip(Clip& clip, ProcessingController& con } } +// Apply object segmentation mask to clip +void ClipProcessingJobs::maskObjectClip(Clip& clip, ProcessingController& controller){ + CVObjectMask objectMask(processInfoJson, controller); + objectMask.maskClip(clip); + + if(controller.ShouldStop()){ + controller.SetFinished(true); + return; + } + else{ + objectMask.SaveObjMaskData(); + controller.SetFinished(true); + } +} + void ClipProcessingJobs::stabilizeClip(Clip& clip, ProcessingController& controller){ // create CVStabilization object CVStabilization stabilizer(processInfoJson, controller); diff --git a/src/ClipProcessingJobs.h b/src/ClipProcessingJobs.h index dff27b26..1f38b521 100644 --- a/src/ClipProcessingJobs.h +++ b/src/ClipProcessingJobs.h @@ -22,6 +22,7 @@ #include "CVStabilization.h" #include "CVTracker.h" #include "CVObjectDetection.h" + #include "CVObjectMask.h" #endif #include @@ -51,6 +52,8 @@ class ClipProcessingJobs{ void stabilizeClip(Clip& clip, ProcessingController& controller); // Apply object detection to clip void detectObjectsClip(Clip& clip, ProcessingController& controller); + // Apply object segmentation mask to clip + void maskObjectClip(Clip& clip, ProcessingController& controller); public: diff --git a/src/EffectInfo.cpp b/src/EffectInfo.cpp index 6281812b..87e417dc 100644 --- a/src/EffectInfo.cpp +++ b/src/EffectInfo.cpp @@ -149,6 +149,9 @@ EffectBase* EffectInfo::CreateEffect(std::string effect_type) { else if(effect_type == "ObjectDetection") return new ObjectDetection(); + + else if(effect_type == "ObjectMask") + return new ObjectMask(); #endif return NULL; @@ -205,6 +208,7 @@ Json::Value EffectInfo::JsonValue() { root.append(Stabilizer().JsonInfo()); root.append(Tracker().JsonInfo()); root.append(ObjectDetection().JsonInfo()); + root.append(ObjectMask().JsonInfo()); #endif // return JsonValue diff --git a/src/Effects.h b/src/Effects.h index e3f8e7eb..af19ed7e 100644 --- a/src/Effects.h +++ b/src/Effects.h @@ -58,6 +58,7 @@ #ifdef USE_OPENCV #include "effects/Outline.h" #include "effects/ObjectDetection.h" +#include "effects/ObjectMask.h" #include "effects/Tracker.h" #include "effects/Stabilizer.h" #endif diff --git a/src/effects/ObjectMask.cpp b/src/effects/ObjectMask.cpp new file mode 100644 index 00000000..477ce2e8 --- /dev/null +++ b/src/effects/ObjectMask.cpp @@ -0,0 +1,327 @@ +/** + * @file + * @brief Source file for Object Mask effect class + * @author Jonathan Thomas + * + * @ref License + */ + +// Copyright (c) 2026 OpenShot Studios, LLC +// +// SPDX-License-Identifier: LGPL-3.0-or-later + +#include "effects/ObjectMask.h" + +#include "Exceptions.h" +#include "Frame.h" +#include "objdetectdata.pb.h" + +#define int64 opencv_broken_int +#define uint64 opencv_broken_uint +#include +#include +#undef uint64 +#undef int64 + +#include +#include +#include + +#include +#include + +using namespace openshot; + +namespace { + +QImage AlphaMaskImageFromRLE(const ObjectMaskFrameData& mask) +{ + QImage image(mask.width, mask.height, QImage::Format_ARGB32_Premultiplied); + image.fill(Qt::transparent); + if (!mask.HasData()) + return image; + + QRgb* data = reinterpret_cast(image.bits()); + const int total = mask.width * mask.height; + int offset = 0; + bool value = false; + for (uint32_t count : mask.rle) { + const int end = std::min(total, offset + static_cast(count)); + if (value) + std::fill(data + offset, data + end, qRgba(255, 255, 255, 255)); + offset = end; + value = !value; + if (offset >= total) + break; + } + return image; +} + +cv::Mat BinaryMaskFromImage(const QImage& image) +{ + QImage rgba = image.convertToFormat(QImage::Format_RGBA8888); + cv::Mat binary(rgba.height(), rgba.width(), CV_8UC1, cv::Scalar(0)); + for (int y = 0; y < rgba.height(); ++y) { + const uchar* source = rgba.constScanLine(y); + uchar* target = binary.ptr(y); + for (int x = 0; x < rgba.width(); ++x) + target[x] = source[x * 4 + 3] > 0 ? 255 : 0; + } + return binary; +} + +QImage StrokeImageFromMask(const QImage& alphaMask, int width) +{ + QImage result(alphaMask.size(), QImage::Format_ARGB32_Premultiplied); + result.fill(Qt::transparent); + if (width <= 0) + return result; + + cv::Mat binary = BinaryMaskFromImage(alphaMask); + cv::Mat dilated; + const int kernelSize = std::max(1, width * 2 + 1); + cv::Mat kernel = cv::getStructuringElement(cv::MORPH_ELLIPSE, cv::Size(kernelSize, kernelSize)); + cv::dilate(binary, dilated, kernel); + cv::Mat edge = dilated - binary; + + for (int y = 0; y < edge.rows; ++y) { + const uchar* edgeRow = edge.ptr(y); + QRgb* target = reinterpret_cast(result.scanLine(y)); + for (int x = 0; x < edge.cols; ++x) { + if (edgeRow[x]) + target[x] = qRgba(255, 255, 255, 255); + } + } + return result; +} + +} + +ObjectMask::ObjectMask() + : draw_mask(1.0) + , mask_color(83, 160, 237, 255) + , mask_alpha(120.0 / 255.0) + , stroke_color(255, 255, 255, 255) + , stroke_alpha(1.0) + , stroke_width(3.0) +{ + init_effect_details(); +} + +void ObjectMask::init_effect_details() +{ + InitEffectInfo(); + info.class_name = "ObjectMask"; + info.name = "Object Mask"; + info.description = "Create and draw a segmentation mask for a prompted object."; + info.has_audio = false; + info.has_video = true; + info.has_tracked_object = true; +} + +std::shared_ptr ObjectMask::GetFrame(std::shared_ptr frame, int64_t frame_number) +{ + std::shared_ptr frame_image = frame->GetImage(); + if (!frame_image || frame_image->isNull() || draw_mask.GetValue(frame_number) != 1) + return frame; + + auto mask_it = masksData.find(frame_number); + if (mask_it == masksData.end() || !mask_it->second.HasData()) + return frame; + + QImage alpha_mask = AlphaMaskImageFromRLE(mask_it->second) + .scaled(frame_image->size(), Qt::IgnoreAspectRatio, Qt::SmoothTransformation); + std::vector mask_rgba = mask_color.GetColorRGBA(frame_number); + QColor overlay_color(mask_rgba[0], mask_rgba[1], mask_rgba[2], 255 * mask_alpha.GetValue(frame_number)); + + QImage overlay(frame_image->size(), QImage::Format_ARGB32_Premultiplied); + overlay.fill(Qt::transparent); + QPainter overlay_painter(&overlay); + overlay_painter.setCompositionMode(QPainter::CompositionMode_Source); + overlay_painter.fillRect(overlay.rect(), overlay_color); + overlay_painter.setCompositionMode(QPainter::CompositionMode_DestinationIn); + overlay_painter.drawImage(0, 0, alpha_mask); + overlay_painter.end(); + + QPainter painter(frame_image.get()); + painter.drawImage(0, 0, overlay); + + const int strokeWidth = static_cast(std::round(stroke_width.GetValue(frame_number))); + if (strokeWidth > 0 && stroke_alpha.GetValue(frame_number) > 0.0) { + QImage stroke_mask = StrokeImageFromMask(alpha_mask, strokeWidth); + std::vector stroke_rgba = stroke_color.GetColorRGBA(frame_number); + QColor stroke_qcolor(stroke_rgba[0], stroke_rgba[1], stroke_rgba[2], 255 * stroke_alpha.GetValue(frame_number)); + + QImage stroke_overlay(frame_image->size(), QImage::Format_ARGB32_Premultiplied); + stroke_overlay.fill(Qt::transparent); + QPainter stroke_painter(&stroke_overlay); + stroke_painter.setCompositionMode(QPainter::CompositionMode_Source); + stroke_painter.fillRect(stroke_overlay.rect(), stroke_qcolor); + stroke_painter.setCompositionMode(QPainter::CompositionMode_DestinationIn); + stroke_painter.drawImage(0, 0, stroke_mask); + stroke_painter.end(); + painter.drawImage(0, 0, stroke_overlay); + } + painter.end(); + + return frame; +} + +bool ObjectMask::LoadObjMaskData(std::string inputFilePath) +{ + pb_objdetect::ObjDetect objMessage; + std::fstream input(inputFilePath, std::ios::in | std::ios::binary); + if (!objMessage.ParseFromIstream(&input)) + return false; + + masksData.clear(); + trackedObjects.clear(); + + auto trackedObject = std::make_shared(83, 160, 237, 255); + trackedObject->Id(Id().empty() ? "Object Mask" : Id() + "-1"); + trackedObject->ParentClip(this->ParentClip()); + trackedObject->draw_box = Keyframe(0.0); + trackedObject->draw_text = Keyframe(0.0); + trackedObject->draw_mask = draw_mask; + trackedObject->mask_alpha = mask_alpha; + trackedObject->mask_color = mask_color; + trackedObject->stroke = stroke_color; + trackedObject->stroke_alpha = stroke_alpha; + trackedObject->stroke_width = stroke_width; + + for (int frameIndex = 0; frameIndex < objMessage.frame_size(); ++frameIndex) { + const auto& pbFrame = objMessage.frame(frameIndex); + if (pbFrame.bounding_box_size() <= 0) + continue; + + const auto& box = pbFrame.bounding_box(0); + ObjectMaskFrameData mask; + mask.box = BBox(box.x() + box.w() / 2.0f, box.y() + box.h() / 2.0f, box.w(), box.h(), 0.0f); + mask.score = box.confidence(); + if (box.has_mask()) { + mask.width = box.mask().width(); + mask.height = box.mask().height(); + for (int rleIndex = 0; rleIndex < box.mask().rle_size(); ++rleIndex) + mask.rle.push_back(box.mask().rle(rleIndex)); + } + masksData[pbFrame.id()] = mask; + + trackedObject->AddBox(pbFrame.id(), mask.box.cx, mask.box.cy, mask.box.width, mask.box.height, 0.0f); + if (mask.HasData()) { + ObjectMaskData trackedMask; + trackedMask.width = mask.width; + trackedMask.height = mask.height; + trackedMask.rle = mask.rle; + trackedObject->AddMask(pbFrame.id(), trackedMask); + } + } + + if (!masksData.empty()) + trackedObjects[1] = trackedObject; + + google::protobuf::ShutdownProtobufLibrary(); + return true; +} + +std::shared_ptr ObjectMask::TrackedObjectMask(std::shared_ptr target_image, int64_t frame_number) const +{ + if (!target_image || target_image->isNull()) + return {}; + + auto mask_it = masksData.find(frame_number); + if (mask_it == masksData.end() || !mask_it->second.HasData()) + return {}; + + QImage alpha_mask = AlphaMaskImageFromRLE(mask_it->second) + .scaled(target_image->size(), Qt::IgnoreAspectRatio, Qt::SmoothTransformation); + + auto mask_image = std::make_shared( + target_image->width(), target_image->height(), QImage::Format_RGBA8888_Premultiplied); + mask_image->fill(QColor(0, 0, 0, 255)); + QPainter painter(mask_image.get()); + painter.drawImage(0, 0, alpha_mask); + painter.end(); + return mask_image; +} + +std::string ObjectMask::Json() const +{ + return JsonValue().toStyledString(); +} + +Json::Value ObjectMask::JsonValue() const +{ + Json::Value root = EffectBase::JsonValue(); + root["type"] = info.class_name; + root["protobuf_data_path"] = protobuf_data_path; + root["draw_mask"] = draw_mask.JsonValue(); + root["mask_color"] = mask_color.JsonValue(); + root["mask_alpha"] = mask_alpha.JsonValue(); + root["stroke_color"] = stroke_color.JsonValue(); + root["stroke_alpha"] = stroke_alpha.JsonValue(); + root["stroke_width"] = stroke_width.JsonValue(); + return root; +} + +void ObjectMask::SetJson(const std::string value) +{ + try { + SetJsonValue(openshot::stringToJson(value)); + } catch (const std::exception&) { + throw InvalidJSON("JSON is invalid (missing keys or invalid data types)"); + } +} + +void ObjectMask::SetJsonValue(const Json::Value root) +{ + EffectBase::SetJsonValue(root); + + if (!root["draw_mask"].isNull()) + draw_mask.SetJsonValue(root["draw_mask"]); + if (!root["mask_color"].isNull()) + mask_color.SetJsonValue(root["mask_color"]); + if (!root["mask_alpha"].isNull()) + mask_alpha.SetJsonValue(root["mask_alpha"]); + if (!root["stroke_color"].isNull()) + stroke_color.SetJsonValue(root["stroke_color"]); + if (!root["stroke"].isNull()) + stroke_color.SetJsonValue(root["stroke"]); + if (!root["stroke_alpha"].isNull()) + stroke_alpha.SetJsonValue(root["stroke_alpha"]); + if (!root["stroke_width"].isNull()) + stroke_width.SetJsonValue(root["stroke_width"]); + + if (!root["protobuf_data_path"].isNull()) { + std::string new_path = root["protobuf_data_path"].asString(); + if (protobuf_data_path != new_path || masksData.empty()) { + protobuf_data_path = new_path; + if (!LoadObjMaskData(protobuf_data_path)) + throw InvalidFile("Invalid object mask protobuf data path", ""); + } + } +} + +std::string ObjectMask::PropertiesJSON(int64_t requested_frame) const +{ + Json::Value root = BasePropertiesJSON(requested_frame); + root["protobuf_data_path"] = add_property_json("Object Mask Data", 0.0, "string", protobuf_data_path, NULL, -1, -1, false, requested_frame); + + root["draw_mask"] = add_property_json("Draw Mask", draw_mask.GetValue(requested_frame), "int", "", &draw_mask, 0, 1, false, requested_frame); + root["draw_mask"]["choices"].append(add_property_choice_json("Yes", true, draw_mask.GetValue(requested_frame))); + root["draw_mask"]["choices"].append(add_property_choice_json("No", false, draw_mask.GetValue(requested_frame))); + + root["mask_color"] = add_property_json("Mask Color", 0.0, "color", "", NULL, 0, 255, false, requested_frame); + root["mask_color"]["red"] = add_property_json("Red", mask_color.red.GetValue(requested_frame), "float", "", &mask_color.red, 0, 255, false, requested_frame); + root["mask_color"]["blue"] = add_property_json("Blue", mask_color.blue.GetValue(requested_frame), "float", "", &mask_color.blue, 0, 255, false, requested_frame); + root["mask_color"]["green"] = add_property_json("Green", mask_color.green.GetValue(requested_frame), "float", "", &mask_color.green, 0, 255, false, requested_frame); + root["mask_alpha"] = add_property_json("Mask Alpha", mask_alpha.GetValue(requested_frame), "float", "", &mask_alpha, 0.0, 1.0, false, requested_frame); + + root["stroke_color"] = add_property_json("Stroke Color", 0.0, "color", "", NULL, 0, 255, false, requested_frame); + root["stroke_color"]["red"] = add_property_json("Red", stroke_color.red.GetValue(requested_frame), "float", "", &stroke_color.red, 0, 255, false, requested_frame); + root["stroke_color"]["blue"] = add_property_json("Blue", stroke_color.blue.GetValue(requested_frame), "float", "", &stroke_color.blue, 0, 255, false, requested_frame); + root["stroke_color"]["green"] = add_property_json("Green", stroke_color.green.GetValue(requested_frame), "float", "", &stroke_color.green, 0, 255, false, requested_frame); + root["stroke_alpha"] = add_property_json("Stroke Alpha", stroke_alpha.GetValue(requested_frame), "float", "", &stroke_alpha, 0.0, 1.0, false, requested_frame); + root["stroke_width"] = add_property_json("Stroke Width", stroke_width.GetValue(requested_frame), "int", "", &stroke_width, 0, 50, false, requested_frame); + + return root.toStyledString(); +} diff --git a/src/effects/ObjectMask.h b/src/effects/ObjectMask.h new file mode 100644 index 00000000..3114af40 --- /dev/null +++ b/src/effects/ObjectMask.h @@ -0,0 +1,73 @@ +/** + * @file + * @brief Header file for Object Mask effect class + * @author Jonathan Thomas + * + * @ref License + */ + +// Copyright (c) 2026 OpenShot Studios, LLC +// +// SPDX-License-Identifier: LGPL-3.0-or-later + +#ifndef OPENSHOT_OBJECT_MASK_EFFECT_H +#define OPENSHOT_OBJECT_MASK_EFFECT_H + +#include "Color.h" +#include "EffectBase.h" +#include "KeyFrame.h" +#include "TrackedObjectBBox.h" + +#include + +namespace openshot +{ + class Frame; + + struct ObjectMaskFrameData { + int width = 0; + int height = 0; + std::vector rle; + BBox box; + float score = 0.0f; + + bool HasData() const { return width > 0 && height > 0 && !rle.empty(); } + }; + + /** + * @brief Display and expose a preprocessed segmentation mask for an object. + */ + class ObjectMask : public EffectBase + { + private: + std::string protobuf_data_path; + std::map masksData; + + void init_effect_details(); + + public: + Keyframe draw_mask; + Color mask_color; + Keyframe mask_alpha; + Color stroke_color; + Keyframe stroke_alpha; + Keyframe stroke_width; + + ObjectMask(); + + std::shared_ptr GetFrame(std::shared_ptr frame, int64_t frame_number) override; + std::shared_ptr GetFrame(int64_t frame_number) override { return GetFrame(std::make_shared(), frame_number); } + + bool LoadObjMaskData(std::string inputFilePath); + std::shared_ptr TrackedObjectMask(std::shared_ptr target_image, int64_t frame_number) const override; + + std::string Json() const override; + void SetJson(const std::string value) override; + Json::Value JsonValue() const override; + void SetJsonValue(const Json::Value root) override; + + std::string PropertiesJSON(int64_t requested_frame) const override; + }; +} + +#endif diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 0167c9fa..4f915b30 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -91,6 +91,7 @@ if($CACHE{HAVE_OPENCV}) CVTracker CVStabilizer CVOutline + ObjectMask # CVObjectDetection ) endif() diff --git a/tests/ObjectMask.cpp b/tests/ObjectMask.cpp new file mode 100644 index 00000000..f774aa8d --- /dev/null +++ b/tests/ObjectMask.cpp @@ -0,0 +1,153 @@ +/** + * @file + * @brief Unit tests for Object Mask effect + * @author Jonathan Thomas + * + * @ref License + */ + +// Copyright (c) 2026 OpenShot Studios, LLC +// +// SPDX-License-Identifier: LGPL-3.0-or-later + +#include "openshot_catch.h" + +#include "EffectInfo.h" +#include "Frame.h" +#include "Json.h" +#include "effects/ObjectMask.h" + +#include +#include + +#include +#include +#include +#include +#include + +using namespace openshot; + +static std::shared_ptr make_object_mask_frame(int64_t number, int width, int height) { + auto frame = std::make_shared(number, width, height, "#000000"); + frame->GetImage()->fill(QColor(64, 64, 64, 255)); + return frame; +} + +static std::string temp_object_mask_path() { + char path[] = "/tmp/libopenshot_object_mask_XXXXXX"; + int fd = mkstemp(path); + REQUIRE(fd != -1); + close(fd); + std::remove(path); + return std::string(path) + ".data"; +} + +static void append_varint(std::string& output, uint64_t value) { + while (value >= 0x80) { + output.push_back(static_cast((value & 0x7f) | 0x80)); + value >>= 7; + } + output.push_back(static_cast(value)); +} + +static void append_fixed32_float(std::string& output, float value) { + uint32_t bits; + std::memcpy(&bits, &value, sizeof(float)); + for (int i = 0; i < 4; ++i) + output.push_back(static_cast((bits >> (8 * i)) & 0xff)); +} + +static void append_length_delimited(std::string& output, uint32_t field_number, const std::string& value) { + append_varint(output, (field_number << 3) | 2); + append_varint(output, static_cast(value.size())); + output.append(value); +} + +static std::string create_object_mask_data() { + const std::string path = temp_object_mask_path(); + + std::string mask; + append_varint(mask, 8); + append_varint(mask, 4); + append_varint(mask, 16); + append_varint(mask, 4); + for (uint32_t count : {0u, 6u, 10u}) { + append_varint(mask, 24); + append_varint(mask, count); + } + + std::string box; + append_varint(box, 13); + append_fixed32_float(box, 0.0f); + append_varint(box, 21); + append_fixed32_float(box, 0.0f); + append_varint(box, 29); + append_fixed32_float(box, 1.0f); + append_varint(box, 37); + append_fixed32_float(box, 1.0f); + append_varint(box, 40); + append_varint(box, 0); + append_varint(box, 53); + append_fixed32_float(box, 0.95f); + append_varint(box, 56); + append_varint(box, 1); + append_length_delimited(box, 8, mask); + + std::string frame; + append_varint(frame, 8); + append_varint(frame, 1); + append_length_delimited(frame, 2, box); + + std::string data; + append_length_delimited(data, 1, frame); + append_length_delimited(data, 3, "object mask"); + + std::ofstream output(path, std::ios::out | std::ios::binary); + output.write(data.data(), static_cast(data.size())); + REQUIRE(output.good()); + return path; +} + +TEST_CASE("ObjectMask effect is registered", "[effect][object_mask]") { + std::unique_ptr effect(EffectInfo().CreateEffect("ObjectMask")); + REQUIRE(effect != nullptr); + CHECK(effect->info.name == "Object Mask"); + CHECK(effect->info.has_video); + CHECK(effect->info.has_tracked_object); +} + +TEST_CASE("ObjectMask loads protobuf masks and exposes style controls", "[effect][object_mask]") { + const std::string protobuf_path = create_object_mask_data(); + + ObjectMask effect; + Json::Value config; + config["protobuf_data_path"] = protobuf_path; + config["mask_alpha"] = Keyframe(0.5).JsonValue(); + config["stroke_width"] = Keyframe(2.0).JsonValue(); + effect.SetJsonValue(config); + + Json::Value properties = stringToJson(effect.PropertiesJSON(1)); + CHECK(properties["draw_mask"]["name"].asString() == "Draw Mask"); + CHECK(properties["mask_color"]["name"].asString() == "Mask Color"); + CHECK(properties["mask_alpha"]["value"].asDouble() == Approx(0.5)); + CHECK(properties["stroke_color"]["name"].asString() == "Stroke Color"); + CHECK(properties["stroke_alpha"]["name"].asString() == "Stroke Alpha"); + CHECK(properties["stroke_width"]["value"].asDouble() == Approx(2.0)); + + auto generated_mask = effect.TrackedObjectMask(std::make_shared(4, 4, QImage::Format_RGBA8888_Premultiplied), 1); + REQUIRE(generated_mask != nullptr); + CHECK(generated_mask->pixelColor(0, 0).red() == 255); + CHECK(generated_mask->pixelColor(3, 3).red() == 0); + + Json::Value no_stroke; + no_stroke["stroke_width"] = Keyframe(0.0).JsonValue(); + effect.SetJsonValue(no_stroke); + + auto frame = make_object_mask_frame(1, 4, 4); + auto output = effect.GetFrame(frame, 1)->GetImage(); + CHECK(output->pixelColor(0, 0) != QColor(64, 64, 64, 255)); + CHECK(output->pixelColor(3, 3) == QColor(64, 64, 64, 255)); + + std::remove(protobuf_path.c_str()); +}