17#include "objdetectdata.pb.h"
20#define uint64 uint64_t
21#include <opencv2/core/ocl.hpp>
22#if CV_VERSION_MAJOR >= 5
23#include <opencv2/geometry/2d.hpp>
25#include <opencv2/imgproc.hpp>
40#include <google/protobuf/util/time_util.h>
43using google::protobuf::util::TimeUtil;
47std::string LoadONNXModel(
const std::string& modelPath, cv::dnn::Net* net)
50 cv::dnn::Net loadedNet = cv::dnn::readNetFromONNX(modelPath);
54 }
catch (
const cv::Exception& e) {
55 return std::string(
"Failed to load ONNX model: ") + e.what();
56 }
catch (
const std::exception& e) {
57 return std::string(
"Failed to load ONNX model: ") + e.what();
61std::vector<uint32_t> EncodeBinaryMaskRLE(
const cv::Mat& mask)
63 std::vector<uint32_t> rle;
69 for (
int y = 0; y < mask.rows; ++y) {
70 const uint8_t* row = mask.ptr<uint8_t>(y);
71 for (
int x = 0; x < mask.cols; ++x) {
72 const uint8_t value = row[x] ? 1 : 0;
73 if (value == current) {
86struct EfficientSamPreprocessResult {
92EfficientSamPreprocessResult MakeEfficientSamBlob(
const cv::Mat& bgr,
int modelSize)
94 EfficientSamPreprocessResult result;
95 result.scaleX =
static_cast<float>(modelSize) /
static_cast<float>(bgr.cols);
96 result.scaleY =
static_cast<float>(modelSize) /
static_cast<float>(bgr.rows);
99 cv::resize(bgr, resized, cv::Size(modelSize, modelSize), 0, 0, cv::INTER_LINEAR);
101 const int shape[] = {1, 3, modelSize, modelSize};
102 result.blob = cv::Mat(4, shape, CV_32F);
103 float* dst = result.blob.ptr<
float>();
105 for (
int y = 0; y < resized.rows; ++y) {
106 const cv::Vec3b* row = resized.ptr<cv::Vec3b>(y);
107 for (
int x = 0; x < resized.cols; ++x) {
108 const float rgb[] = {
109 static_cast<float>(row[x][2]) / 255.0f,
110 static_cast<float>(row[x][1]) / 255.0f,
111 static_cast<float>(row[x][0]) / 255.0f,
113 for (
int c = 0; c < 3; ++c)
114 dst[(c * modelSize + y) * modelSize + x] = rgb[c];
121cv::Rect_<float> NormalizedBoundingBox(
const cv::Mat& mask)
123 std::vector<cv::Point> points;
124 cv::findNonZero(mask, points);
128 cv::Rect rect = cv::boundingRect(points);
129 return cv::Rect_<float>(
130 rect.x /
static_cast<float>(mask.cols),
131 rect.y /
static_cast<float>(mask.rows),
132 rect.width /
static_cast<float>(mask.cols),
133 rect.height /
static_cast<float>(mask.rows));
136cv::Mat EfficientSamMaskToFrameMask(
const cv::Mat& modelMask,
const cv::Size& frameSize,
float maskThreshold)
139 cv::resize(modelMask, fullSize, frameSize, 0, 0, cv::INTER_LINEAR);
142 cv::threshold(fullSize, binary, maskThreshold, 255.0, cv::THRESH_BINARY);
143 if (cv::countNonZero(binary) == 0) {
144 double maxValue = 0.0;
145 cv::minMaxLoc(fullSize,
nullptr, &maxValue);
146 if (maxValue > 0.0) {
147 cv::threshold(fullSize, binary, maxValue * 0.5, 255.0, cv::THRESH_BINARY);
150 binary.convertTo(binary, CV_8U);
154cv::Mat MakeEfficientSamPromptBlob(
156 const EfficientSamPreprocessResult& prep,
158 std::vector<cv::Point>& backgroundPoints,
159 std::vector<cv::Rect>& backgroundRects)
161 const int coordsShape[] = {1, 1, promptSlots, 2};
162 cv::Mat pointCoords(4, coordsShape, CV_32F, cv::Scalar(0.0f));
164 float* coords = pointCoords.ptr<
float>();
167 if (promptIndex + 1 >= promptSlots)
169 coords[promptIndex * 2] = rect.x * prep.scaleX;
170 coords[promptIndex * 2 + 1] = rect.y * prep.scaleY;
172 coords[promptIndex * 2] = (rect.x + rect.width) * prep.scaleX;
173 coords[promptIndex * 2 + 1] = (rect.y + rect.height) * prep.scaleY;
177 if (promptIndex >= promptSlots)
179 coords[promptIndex * 2] = point.x * prep.scaleX;
180 coords[promptIndex * 2 + 1] = point.y * prep.scaleY;
184 backgroundPoints.emplace_back(
185 static_cast<int>(std::lround(point.x * prep.scaleX)),
186 static_cast<int>(std::lround(point.y * prep.scaleY)));
189 const int x1 =
static_cast<int>(std::floor(rect.x * prep.scaleX));
190 const int y1 =
static_cast<int>(std::floor(rect.y * prep.scaleY));
191 const int x2 =
static_cast<int>(std::ceil((rect.x + rect.width) * prep.scaleX));
192 const int y2 =
static_cast<int>(std::ceil((rect.y + rect.height) * prep.scaleY));
193 const int modelWidth = prep.blob.size[3];
194 const int modelHeight = prep.blob.size[2];
195 const int left = std::max(0, std::min(modelWidth - 1, x1));
196 const int top = std::max(0, std::min(modelHeight - 1, y1));
197 const int right = std::max(left + 1, std::min(modelWidth, x2));
198 const int bottom = std::max(top + 1, std::min(modelHeight, y2));
199 backgroundRects.emplace_back(left, top, right - left, bottom - top);
207 const int labelsShape[] = {1, 1, promptSlots, 1};
208 cv::Mat pointLabels(4, labelsShape, CV_32F, cv::Scalar(-1.0f));
210 float* labels = pointLabels.ptr<
float>();
212 for (
size_t i = 0; i < prompts.
positiveRects.size() && promptIndex + 1 < promptSlots; ++i) {
213 labels[promptIndex++] = 2.0f;
214 labels[promptIndex++] = 3.0f;
216 for (
size_t i = 0; i < prompts.
positivePoints.size() && promptIndex < promptSlots; ++i, ++promptIndex)
217 labels[promptIndex] = 1.0f;
222cv::Mat SelectEfficientSamMask(
const cv::Mat& outputMasks,
const cv::Mat& iouPredictions,
223 const std::vector<cv::Point>& backgroundPoints,
224 const std::vector<cv::Rect>& backgroundRects,
227 if (outputMasks.dims != 5 || iouPredictions.empty())
230 const int candidateCount = outputMasks.size[2];
231 const int maskHeight = outputMasks.size[3];
232 const int maskWidth = outputMasks.size[4];
233 const float* ious = iouPredictions.ptr<
float>();
235 const float* masks = outputMasks.ptr<
float>();
236 const size_t candidatePixels =
static_cast<size_t>(maskHeight) *
static_cast<size_t>(maskWidth);
238 float bestScore = -std::numeric_limits<float>::infinity();
239 for (
int candidate = 0; candidate < candidateCount; ++candidate) {
240 cv::Mat mask(maskHeight, maskWidth, CV_32F,
241 const_cast<float*
>(masks +
static_cast<size_t>(candidate) * candidatePixels));
243 int backgroundHits = 0;
244 for (
const cv::Point& point : backgroundPoints) {
245 const int x = std::max(0, std::min(maskWidth - 1, point.x));
246 const int y = std::max(0, std::min(maskHeight - 1, point.y));
247 if (mask.at<
float>(y, x) >= maskThreshold)
251 float rectOverlapPenalty = 0.0f;
252 for (
const cv::Rect& rect : backgroundRects) {
253 const cv::Rect clipped = rect & cv::Rect(0, 0, maskWidth, maskHeight);
254 const int area = clipped.area();
258 for (
int y = clipped.y; y < clipped.y + clipped.height; ++y) {
259 const float* row = mask.ptr<
float>(y);
260 for (
int x = clipped.x; x < clipped.x + clipped.width; ++x) {
261 if (row[x] >= maskThreshold)
265 rectOverlapPenalty +=
static_cast<float>(overlap) /
static_cast<float>(area);
268 const float pointPenalty = backgroundPoints.empty()
270 :
static_cast<float>(backgroundHits) /
static_cast<float>(backgroundPoints.size());
271 if (!backgroundRects.empty())
272 rectOverlapPenalty /=
static_cast<float>(backgroundRects.size());
274 const float score = ious[candidate] - (0.35f * pointPenalty) - (0.75f * rectOverlapPenalty);
275 if (bestMask.empty() || score > bestScore) {
277 bestMask = mask.clone();
291 frameData.
score = score;
292 frameData.
width = mask.cols;
293 frameData.
height = mask.rows;
294 frameData.
rle = EncodeBinaryMaskRLE(mask);
295 frameData.
box = NormalizedBoundingBox(mask);
299cv::Point2f JsonPoint(
const Json::Value& value)
301 if (!value.isObject() || value[
"x"].isNull() || value[
"y"].isNull())
302 return cv::Point2f(-1.0f, -1.0f);
303 return cv::Point2f(value[
"x"].asFloat(), value[
"y"].asFloat());
306bool IsValidPoint(
const cv::Point2f& point)
308 return point.x >= 0.0f && point.y >= 0.0f;
311void AppendJsonPoints(
const Json::Value& values, std::vector<cv::Point2f>& points)
313 if (!values.isArray())
315 for (
const auto& value : values) {
316 cv::Point2f point = JsonPoint(value);
317 if (IsValidPoint(point))
318 points.push_back(point);
322size_t JsonFrameNumber(
const std::string& frameName)
325 return static_cast<size_t>(std::max(0, std::stoi(frameName)));
331bool RectFromJson(
const Json::Value& rect, cv::Rect_<float>& output)
333 if (!rect.isObject() || rect[
"x1"].isNull() || rect[
"y1"].isNull() ||
334 rect[
"x2"].isNull() || rect[
"y2"].isNull()) {
338 const float x1 = std::min(rect[
"x1"].asFloat(), rect[
"x2"].asFloat());
339 const float y1 = std::min(rect[
"y1"].asFloat(), rect[
"y2"].asFloat());
340 const float x2 = std::max(rect[
"x1"].asFloat(), rect[
"x2"].asFloat());
341 const float y2 = std::max(rect[
"y1"].asFloat(), rect[
"y2"].asFloat());
342 cv::Point2f topLeft(x1, y1);
343 cv::Point2f bottomRight(x2, y2);
344 if (!IsValidPoint(topLeft) || !IsValidPoint(bottomRight) || x2 <= x1 || y2 <= y1)
347 output = cv::Rect_<float>(x1, y1, x2 - x1, y2 - y1);
351void AppendJsonRects(
const Json::Value& values, std::vector<cv::Rect_<float>>& rects)
353 if (!values.isArray())
355 for (
const auto& rect : values) {
356 cv::Rect_<float> parsed;
357 if (RectFromJson(rect, parsed))
358 rects.push_back(parsed);
365 AppendJsonPoints(framePayload[
"positive_points"], prompts.
positivePoints);
366 AppendJsonPoints(framePayload[
"negative_points"], prompts.
negativePoints);
367 AppendJsonRects(framePayload[
"positive_rects"], prompts.
positiveRects);
368 AppendJsonRects(framePayload[
"negative_rects"], prompts.
negativeRects);
372cv::Mat MakeBlob(
const std::vector<int>& shape,
float value = 0.0f)
374 cv::Mat output(
static_cast<int>(shape.size()), shape.data(), CV_32F);
379std::string SetNetDevice(cv::dnn::Net& net,
const std::string& processingDevice)
381 if (processingDevice ==
"CPU") {
382 net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
383 net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
387 if (processingDevice ==
"GPU" || processingDevice ==
"GPU_AUTO" || processingDevice ==
"GPU_CUDA") {
389 const std::vector<cv::dnn::Target> targets = cv::dnn::getAvailableTargets(cv::dnn::DNN_BACKEND_CUDA);
390 if (std::find(targets.begin(), targets.end(), cv::dnn::DNN_TARGET_CUDA) != targets.end()) {
391 net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
392 net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
395 }
catch (
const cv::Exception&) {
399 if (processingDevice ==
"GPU_OPENCL") {
401 const std::vector<cv::dnn::Target> targets = cv::dnn::getAvailableTargets(cv::dnn::DNN_BACKEND_OPENCV);
402 if (std::find(targets.begin(), targets.end(), cv::dnn::DNN_TARGET_OPENCL) != targets.end()) {
403 cv::ocl::setUseOpenCL(
true);
404 net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
405 net.setPreferableTarget(cv::dnn::DNN_TARGET_OPENCL);
408 }
catch (
const cv::Exception&) {
412 net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
413 net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
417class CutiePropagator {
419 static constexpr int memorySlots = 6;
420 int modelWidth = 640;
421 int modelHeight = 368;
422 int stride16Width = modelWidth / 16;
423 int stride16Height = modelHeight / 16;
432 struct LetterboxTransform {
433 cv::Size originalSize;
434 cv::Rect contentRect;
437 cv::dnn::Net encodeKey;
438 cv::dnn::Net encodeValue;
439 cv::dnn::Net memoryReadout;
443 cv::Mat objectMemory;
444 MemoryFrame permanentMemory;
445 bool hasPermanentMemory =
false;
446 std::deque<MemoryFrame> workingMemoryFrames;
448 int lastMemoryFrame = -1000000;
450 int maxMemoryFrames = memorySlots;
452 static bool ParseModelSize(
const std::string& modelPath,
int& width,
int& height)
454 size_t xPos = modelPath.find(
'x');
455 while (xPos != std::string::npos) {
456 size_t widthStart = xPos;
457 while (widthStart > 0 && std::isdigit(
static_cast<unsigned char>(modelPath[widthStart - 1])))
460 size_t heightEnd = xPos + 1;
461 while (heightEnd < modelPath.size() && std::isdigit(
static_cast<unsigned char>(modelPath[heightEnd])))
464 if (widthStart != xPos && heightEnd != xPos + 1) {
465 width = std::stoi(modelPath.substr(widthStart, xPos - widthStart));
466 height = std::stoi(modelPath.substr(xPos + 1, heightEnd - xPos - 1));
467 if (width > 0 && height > 0 && width % 16 == 0 && height % 16 == 0)
470 xPos = modelPath.find(
'x', xPos + 1);
475 void ConfigureModelSize(
const std::string& modelPath)
477 int width = modelWidth;
478 int height = modelHeight;
479 if (!ParseModelSize(modelPath, width, height))
482 modelHeight = height;
483 stride16Width = modelWidth / 16;
484 stride16Height = modelHeight / 16;
487 LetterboxTransform ComputeLetterbox(
const cv::Size& sourceSize)
const
489 LetterboxTransform transform;
490 transform.originalSize = sourceSize;
491 if (sourceSize.width <= 0 || sourceSize.height <= 0) {
492 transform.contentRect = cv::Rect(0, 0, modelWidth, modelHeight);
496 const float scaleX =
static_cast<float>(modelWidth) /
static_cast<float>(sourceSize.width);
497 const float scaleY =
static_cast<float>(modelHeight) /
static_cast<float>(sourceSize.height);
498 const float scale = std::min(scaleX, scaleY);
500 const int resizedWidth = std::max(1, std::min(
501 modelWidth,
static_cast<int>(std::lround(sourceSize.width * scale))));
502 const int resizedHeight = std::max(1, std::min(
503 modelHeight,
static_cast<int>(std::lround(sourceSize.height * scale))));
504 const int offsetX = (modelWidth - resizedWidth) / 2;
505 const int offsetY = (modelHeight - resizedHeight) / 2;
506 transform.contentRect = cv::Rect(offsetX, offsetY, resizedWidth, resizedHeight);
510 cv::Mat MakeImageBlob(
const cv::Mat& bgr,
const LetterboxTransform& transform)
const
513 cv::resize(bgr, resized, transform.contentRect.size(), 0, 0, cv::INTER_LINEAR);
514 cv::Mat canvas(modelHeight, modelWidth, bgr.type(), cv::Scalar::all(0));
515 resized.copyTo(canvas(transform.contentRect));
517 const int shape[] = {1, 3, modelHeight, modelWidth};
518 cv::Mat blob(4, shape, CV_32F);
519 float* dst = blob.ptr<
float>();
520 for (
int y = 0; y < canvas.rows; ++y) {
521 const cv::Vec3b* row = canvas.ptr<cv::Vec3b>(y);
522 for (
int x = 0; x < canvas.cols; ++x) {
523 dst[(0 * modelHeight + y) * modelWidth + x] =
static_cast<float>(row[x][2]) / 255.0f;
524 dst[(1 * modelHeight + y) * modelWidth + x] =
static_cast<float>(row[x][1]) / 255.0f;
525 dst[(2 * modelHeight + y) * modelWidth + x] =
static_cast<float>(row[x][0]) / 255.0f;
531 cv::Mat MakeMaskBlob(
const cv::Mat& mask,
const LetterboxTransform& transform)
const
534 cv::resize(mask, resized, transform.contentRect.size(), 0, 0, cv::INTER_NEAREST);
535 cv::Mat canvas(modelHeight, modelWidth, CV_8U, cv::Scalar(0));
536 resized.copyTo(canvas(transform.contentRect));
538 const int shape[] = {1, 1, modelHeight, modelWidth};
539 cv::Mat blob(4, shape, CV_32F, cv::Scalar(0.0f));
540 float* dst = blob.ptr<
float>();
541 for (
int y = 0; y < canvas.rows; ++y) {
542 const uint8_t* row = canvas.ptr<uint8_t>(y);
543 for (
int x = 0; x < canvas.cols; ++x)
544 dst[y * modelWidth + x] = row[x] ? 1.0f : 0.0f;
549 cv::Mat ForegroundFromProb(
const cv::Mat& prob)
const
551 const int shape[] = {1, 1, modelHeight, modelWidth};
552 cv::Mat foreground(4, shape, CV_32F);
553 const float* src = prob.ptr<
float>();
554 float* dst = foreground.ptr<
float>();
555 const int plane = modelWidth * modelHeight;
556 std::memcpy(dst, src + plane,
sizeof(
float) * plane);
560 cv::Mat BinaryMaskFromForeground(
const cv::Mat& foreground,
const LetterboxTransform& transform)
const
562 cv::Mat modelMask(modelHeight, modelWidth, CV_8U, cv::Scalar(0));
563 const float* src = foreground.ptr<
float>();
564 for (
int y = 0; y < modelMask.rows; ++y) {
565 uint8_t* row = modelMask.ptr<uint8_t>(y);
566 for (
int x = 0; x < modelMask.cols; ++x)
567 row[x] = src[y * modelWidth + x] >= 0.5f ? 255 : 0;
570 cv::Mat cropped = modelMask(transform.contentRect);
572 cv::resize(cropped, restored, transform.originalSize, 0, 0, cv::INTER_NEAREST);
576 cv::Mat ValidMaskFromLetterbox(
const LetterboxTransform& transform)
const
578 cv::Mat valid(stride16Height, stride16Width, CV_32F, cv::Scalar(0.0f));
579 for (
int y = 0; y < stride16Height; ++y) {
580 float* row = valid.ptr<
float>(y);
581 const int centerY = y * 16 + 8;
582 for (
int x = 0; x < stride16Width; ++x) {
583 const int centerX = x * 16 + 8;
584 if (transform.contentRect.contains(cv::Point(centerX, centerY)))
589 const int shape[] = {1, 1, stride16Height, stride16Width};
590 cv::Mat blob(4, shape, CV_32F);
591 std::memcpy(blob.ptr<
float>(), valid.ptr<
float>(),
sizeof(
float) * valid.total());
595 void CopyKeySlot(
const cv::Mat& src, cv::Mat& dst,
int slot,
int channels)
const
597 const float* in = src.ptr<
float>();
598 float* out = dst.ptr<
float>();
599 const int plane = stride16Width * stride16Height;
600 for (
int c = 0; c < channels; ++c) {
601 std::memcpy(out + (c * memorySlots + slot) * plane,
603 sizeof(
float) * plane);
607 void CopyValueSlot(
const cv::Mat& src, cv::Mat& dst,
int slot)
const
609 const float* in = src.ptr<
float>();
610 float* out = dst.ptr<
float>();
611 const int plane = stride16Width * stride16Height;
612 for (
int c = 0; c < 256; ++c) {
613 std::memcpy(out + (c * memorySlots + slot) * plane,
615 sizeof(
float) * plane);
619 cv::Mat MemoryKeyBlob()
const
621 cv::Mat output = MakeBlob({1, 64, memorySlots, stride16Height, stride16Width});
623 if (hasPermanentMemory)
624 CopyKeySlot(permanentMemory.key, output, slot++, 64);
626 index < static_cast<int>(workingMemoryFrames.size()) && slot < memorySlots;
628 CopyKeySlot(workingMemoryFrames[index].key, output, slot, 64);
632 cv::Mat MemoryShrinkageBlob()
const
634 cv::Mat output = MakeBlob({1, 1, memorySlots, stride16Height, stride16Width});
636 if (hasPermanentMemory)
637 CopyKeySlot(permanentMemory.shrinkage, output, slot++, 1);
639 index < static_cast<int>(workingMemoryFrames.size()) && slot < memorySlots;
641 CopyKeySlot(workingMemoryFrames[index].shrinkage, output, slot, 1);
645 cv::Mat MemoryValueBlob()
const
647 cv::Mat output = MakeBlob({1, 1, 256, memorySlots, stride16Height, stride16Width});
649 if (hasPermanentMemory)
650 CopyValueSlot(permanentMemory.value, output, slot++);
652 index < static_cast<int>(workingMemoryFrames.size()) && slot < memorySlots;
654 CopyValueSlot(workingMemoryFrames[index].value, output, slot);
658 cv::Mat MemoryValidBlob()
const
660 cv::Mat output = MakeBlob({1, 1, memorySlots, stride16Height, stride16Width});
661 float* data = output.ptr<
float>();
662 const int plane = stride16Width * stride16Height;
663 auto copyValidSlot = [&](
const cv::Mat& valid,
int slot) {
664 std::memcpy(data + slot * plane, valid.ptr<
float>(),
sizeof(
float) * plane);
668 if (hasPermanentMemory)
669 copyValidSlot(permanentMemory.valid, slot++);
671 index < static_cast<int>(workingMemoryFrames.size()) && slot < memorySlots;
673 copyValidSlot(workingMemoryFrames[index].valid, slot);
677 void AddMemory(
const cv::Mat& key,
const cv::Mat& shrinkage,
const cv::Mat& value,
678 const cv::Mat& valid,
bool asPermanent)
681 frame.key = key.clone();
682 frame.shrinkage = shrinkage.clone();
683 frame.value = value.clone();
684 frame.valid = valid.clone();
686 if (asPermanent || !hasPermanentMemory) {
687 permanentMemory = frame;
688 hasPermanentMemory =
true;
692 workingMemoryFrames.push_back(frame);
693 const int workingCapacity = std::max(0, maxMemoryFrames - 1);
694 while (
static_cast<int>(workingMemoryFrames.size()) > workingCapacity)
695 workingMemoryFrames.pop_front();
698 void AddObjectMemory(
const cv::Mat& value)
700 if (objectMemory.empty()) {
701 objectMemory = MakeBlob({1, 1, 1, 16, 257});
702 std::memcpy(objectMemory.ptr<
float>(), value.ptr<
float>(),
sizeof(
float) * value.total());
706 float* dst = objectMemory.ptr<
float>();
707 const float* src = value.ptr<
float>();
708 for (
size_t i = 0; i < value.total(); ++i)
713 void Load(
const std::string& encodeKeyPath,
const std::string& encodeValuePath,
714 const std::string& memoryReadoutPath,
const std::string& decodePath)
716 ConfigureModelSize(encodeKeyPath);
717 encodeKey = cv::dnn::readNetFromONNX(encodeKeyPath);
718 encodeValue = cv::dnn::readNetFromONNX(encodeValuePath);
719 memoryReadout = cv::dnn::readNetFromONNX(memoryReadoutPath);
720 decode = cv::dnn::readNetFromONNX(decodePath);
721 sensory = MakeBlob({1, 1, 256, stride16Height, stride16Width});
724 std::string SetDevice(
const std::string& processingDevice)
726 std::string selected = SetNetDevice(encodeKey, processingDevice);
727 const std::string valueDevice = SetNetDevice(encodeValue, processingDevice);
728 const std::string readoutDevice = SetNetDevice(memoryReadout, processingDevice);
729 const std::string decodeDevice = SetNetDevice(decode, processingDevice);
730 if (selected != valueDevice || selected != readoutDevice || selected != decodeDevice)
737 sensory = MakeBlob({1, 1, 256, stride16Height, stride16Width});
739 objectMemory.release();
740 permanentMemory = MemoryFrame();
741 hasPermanentMemory =
false;
742 workingMemoryFrames.clear();
744 lastMemoryFrame = -1000000;
747 bool HasMemory()
const
749 return hasPermanentMemory || !workingMemoryFrames.empty();
752 cv::Mat Step(
const cv::Mat& frame,
const cv::Mat& seedMask = cv::Mat())
754 const LetterboxTransform transform = ComputeLetterbox(frame.size());
755 const cv::Mat validMask = ValidMaskFromLetterbox(transform);
756 cv::Mat image = MakeImageBlob(frame, transform);
758 encodeKey.setInput(image,
"image");
759 std::vector<cv::Mat> keyOutputs;
760 encodeKey.forward(keyOutputs, std::vector<cv::String>{
"f16",
"f8",
"f4",
"pix_feat",
"key",
"shrinkage",
"selection"});
761 cv::Mat f8 = keyOutputs[1];
762 cv::Mat f4 = keyOutputs[2];
763 cv::Mat pixFeat = keyOutputs[3];
764 cv::Mat key = keyOutputs[4];
765 cv::Mat shrinkage = keyOutputs[5];
766 cv::Mat selection = keyOutputs[6];
769 if (!seedMask.empty()) {
770 foreground = MakeMaskBlob(seedMask, transform);
771 }
else if (HasMemory()) {
772 memoryReadout.setInput(key,
"query_key");
773 memoryReadout.setInput(selection,
"query_selection");
774 memoryReadout.setInput(MemoryKeyBlob(),
"memory_key");
775 memoryReadout.setInput(MemoryShrinkageBlob(),
"memory_shrinkage");
776 memoryReadout.setInput(MemoryValueBlob(),
"memory_value");
777 memoryReadout.setInput(MemoryValidBlob(),
"memory_valid");
778 memoryReadout.setInput(objectMemory,
"object_memory");
779 memoryReadout.setInput(pixFeat,
"pix_feat");
780 memoryReadout.setInput(sensory,
"sensory");
781 memoryReadout.setInput(lastMask,
"last_mask");
782 std::vector<cv::Mat> readoutOutputs;
783 memoryReadout.forward(readoutOutputs, std::vector<cv::String>{
"memory_readout"});
785 decode.setInput(f8,
"f8");
786 decode.setInput(f4,
"f4");
787 decode.setInput(readoutOutputs[0],
"memory_readout");
788 decode.setInput(sensory,
"sensory");
789 std::vector<cv::Mat> decodeOutputs;
790 decode.forward(decodeOutputs, std::vector<cv::String>{
"new_sensory",
"logits",
"prob"});
791 sensory = decodeOutputs[0].clone();
792 foreground = ForegroundFromProb(decodeOutputs[2]);
798 const bool isMemoryFrame = !seedMask.empty() || frameIndex - lastMemoryFrame >= memEvery;
800 encodeValue.setInput(image,
"image");
801 encodeValue.setInput(pixFeat,
"pix_feat");
802 encodeValue.setInput(sensory,
"sensory");
803 encodeValue.setInput(foreground,
"mask");
804 std::vector<cv::Mat> valueOutputs;
805 encodeValue.forward(valueOutputs, std::vector<cv::String>{
"mask_value",
"new_sensory",
"object_memory"});
806 sensory = valueOutputs[1].clone();
807 AddObjectMemory(valueOutputs[2]);
808 AddMemory(key, shrinkage, valueOutputs[0], validMask, !seedMask.empty());
809 lastMemoryFrame = frameIndex;
812 lastMask = foreground.clone();
813 cv::Mat outputMask = BinaryMaskFromForeground(foreground, transform);
822 : processingController(&controller)
829 return LoadONNXModel(modelPath,
nullptr);
834 if (!frame || efficientSamModelPath.empty() || promptKeyframes.empty())
835 return std::shared_ptr<Frame>();
837 std::string loadError = LoadONNXModel(efficientSamModelPath, &efficientSam);
838 if (!loadError.empty())
839 return std::shared_ptr<Frame>();
840 SetProcessingDevice();
843 cv::Mat frameImage = frame->GetImageCV();
844 cv::Mat seedMask = CreateEfficientSAMSeedMask(frameImage, prompts);
845 if (seedMask.empty())
846 return std::shared_ptr<Frame>();
848 auto maskImage = std::make_shared<QImage>(
849 seedMask.cols, seedMask.rows, QImage::Format_RGBA8888_Premultiplied);
850 maskImage->fill(Qt::transparent);
851 for (
int y = 0; y < seedMask.rows; ++y) {
852 const uint8_t* src = seedMask.ptr<uint8_t>(y);
853 QRgb* dst =
reinterpret_cast<QRgb*
>(maskImage->scanLine(y));
854 for (
int x = 0; x < seedMask.cols; ++x)
855 dst[x] = src[x] ? qRgba(255, 255, 255, 255) : qRgba(0, 0, 0, 0);
858 auto result = std::make_shared<Frame>(frame->number, seedMask.cols, seedMask.rows,
"#000000");
859 result->AddImage(maskImage);
863void CVObjectMask::SetProcessingDevice()
865 const std::string requestedDevice = processingDevice;
866 if (processingDevice ==
"CPU") {
867 efficientSam.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
868 efficientSam.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
869 Logger::Instance()->
Log(
"Object Mask EfficientSAM DNN device: requested CPU, selected CPU");
873 if (processingDevice ==
"GPU" || processingDevice ==
"GPU_AUTO" || processingDevice ==
"GPU_CUDA") {
875 const std::vector<cv::dnn::Target> targets = cv::dnn::getAvailableTargets(cv::dnn::DNN_BACKEND_CUDA);
876 if (std::find(targets.begin(), targets.end(), cv::dnn::DNN_TARGET_CUDA) != targets.end()) {
877 efficientSam.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
878 efficientSam.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
879 Logger::Instance()->
Log(
"Object Mask EfficientSAM DNN device: requested " + requestedDevice +
", selected CUDA");
882 }
catch (
const cv::Exception&) {
886 if (processingDevice ==
"GPU_OPENCL") {
888 const std::vector<cv::dnn::Target> targets = cv::dnn::getAvailableTargets(cv::dnn::DNN_BACKEND_OPENCV);
889 if (std::find(targets.begin(), targets.end(), cv::dnn::DNN_TARGET_OPENCL) != targets.end()) {
890 cv::ocl::setUseOpenCL(
true);
891 efficientSam.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
892 efficientSam.setPreferableTarget(cv::dnn::DNN_TARGET_OPENCL);
893 Logger::Instance()->
Log(
"Object Mask EfficientSAM DNN device: requested " + requestedDevice +
", selected OpenCL");
896 }
catch (
const cv::Exception&) {
900 processingDevice =
"CPU";
901 efficientSam.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
902 efficientSam.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
903 Logger::Instance()->
Log(
"Object Mask EfficientSAM DNN device: requested " + requestedDevice +
", selected CPU");
912 processingController->
SetError(
false,
"");
914 if (efficientSamModelPath.empty()) {
915 processingController->
SetError(
true,
"Missing path to EfficientSAM ONNX model file");
919 if (protobufDataPath.empty()) {
920 processingController->
SetError(
true,
"Missing path to object mask protobuf data file");
924 if (promptKeyframes.empty()) {
925 processingController->
SetError(
true,
"Missing positive prompt point for Object Mask preprocessing");
930 std::string loadError = LoadONNXModel(efficientSamModelPath, &efficientSam);
931 if (!loadError.empty()) {
932 processingController->
SetError(
true, loadError);
936 SetProcessingDevice();
938 CutiePropagator cutie;
939 if (cutieEncodeKeyModelPath.empty() && !cutieModelDir.empty())
940 cutieEncodeKeyModelPath = cutieModelDir +
"/cutie-encode-key-640x368.onnx";
941 if (cutieEncodeValueModelPath.empty() && !cutieModelDir.empty())
942 cutieEncodeValueModelPath = cutieModelDir +
"/cutie-encode-value-640x368.onnx";
943 if (cutieMemoryReadoutModelPath.empty() && !cutieModelDir.empty())
944 cutieMemoryReadoutModelPath = cutieModelDir +
"/cutie-memory-readout-floatmask-valid-640x368-m6-topk30-opencv.onnx";
945 if (cutieDecodeModelPath.empty() && !cutieModelDir.empty())
946 cutieDecodeModelPath = cutieModelDir +
"/cutie-decode-640x368.onnx";
947 if (cutieEncodeKeyModelPath.empty() || cutieEncodeValueModelPath.empty() ||
948 cutieMemoryReadoutModelPath.empty() || cutieDecodeModelPath.empty()) {
949 processingController->
SetError(
true,
"Missing path to Cutie ONNX model files");
954 cutie.Load(cutieEncodeKeyModelPath, cutieEncodeValueModelPath, cutieMemoryReadoutModelPath, cutieDecodeModelPath);
955 const std::string cutieDevice = cutie.SetDevice(processingDevice);
956 Logger::Instance()->
Log(
"Object Mask Cutie DNN device: requested " + processingDevice +
", selected " + cutieDevice);
957 }
catch (
const cv::Exception& e) {
958 processingController->
SetError(
true, std::string(
"Failed to load Cutie ONNX models: ") + e.what());
961 }
catch (
const std::exception& e) {
962 processingController->
SetError(
true, std::string(
"Failed to load Cutie ONNX models: ") + e.what());
967 if (!process_interval || end <= 1 || end - start == 0) {
968 start =
static_cast<size_t>(video.
Start() * video.
Reader()->info.fps.ToFloat());
969 end =
static_cast<size_t>(video.
End() * video.
Reader()->info.fps.ToFloat());
975 auto promptBeforeStart = promptKeyframes.upper_bound(start);
976 if (promptBeforeStart != promptKeyframes.begin()) {
978 activePrompts = promptBeforeStart->second;
980 auto firstPromptAtOrAfterStart = promptKeyframes.lower_bound(start);
982 for (
size_t frameNumber = start; frameNumber <= end; ++frameNumber) {
986 std::shared_ptr<openshot::Frame> frame = video.
GetFrame(frameNumber);
990 auto promptIt = promptKeyframes.find(frameNumber);
991 bool isPromptKeyframe = promptIt != promptKeyframes.end();
992 if (promptIt != promptKeyframes.end()) {
993 activePrompts = promptIt->second;
996 if (firstPromptAtOrAfterStart != promptKeyframes.end() && frameNumber >= firstPromptAtOrAfterStart->first) {
997 activePrompts = firstPromptAtOrAfterStart->second;
998 isPromptKeyframe =
true;
1002 emptyFrame.
frameId = frameNumber;
1008 const cv::Mat frameImage = frame->GetImageCV();
1010 if (isPromptKeyframe || !cutie.HasMemory()) {
1011 seedMask = CreateEfficientSAMSeedMask(frameImage, activePrompts);
1012 if (seedMask.empty()) {
1014 emptyFrame.
frameId = frameNumber;
1018 if (!isPromptKeyframe)
1022 cv::Mat propagatedMask;
1024 propagatedMask = cutie.Step(frameImage, seedMask);
1025 }
catch (
const cv::Exception& e) {
1026 processingController->
SetError(
true, std::string(
"Failed to propagate Object Mask with Cutie: ") + e.what());
1032 if (!seedMask.empty()) {
1033 outputMask = seedMask;
1034 }
else if (!propagatedMask.empty()) {
1035 cv::resize(propagatedMask, outputMask, frameImage.size(), 0, 0, cv::INTER_NEAREST);
1037 masksData[frameNumber] = FrameDataFromMask(outputMask, frameNumber, 1.0f);
1039 const size_t range = std::max<size_t>(1, end - start);
1040 processingController->
SetProgress(uint(100 * (frameNumber - start) / range));
1044cv::Mat CVObjectMask::CreateEfficientSAMSeedMask(
const cv::Mat& frame,
const CVObjectMaskPromptSet& prompts)
1046 EfficientSamPreprocessResult prep = MakeEfficientSamBlob(frame, modelSize);
1049 std::vector<cv::Point> backgroundPoints;
1050 std::vector<cv::Rect> backgroundRects;
1051 cv::Mat pointCoords = MakeEfficientSamPromptBlob(promptSet, prep, promptSlots, backgroundPoints, backgroundRects);
1052 cv::Mat pointLabels = MakeEfficientSamLabelBlob(promptSet, promptSlots);
1054 efficientSam.setInput(prep.blob,
"batched_images");
1055 efficientSam.setInput(pointCoords,
"batched_point_coords");
1056 efficientSam.setInput(pointLabels,
"batched_point_labels");
1058 std::vector<cv::Mat> outputs;
1059 efficientSam.forward(outputs, std::vector<cv::String>{
"output_masks",
"iou_predictions"});
1060 if (outputs.size() != 2)
1063 cv::Mat modelMask = SelectEfficientSamMask(outputs[0], outputs[1], backgroundPoints, backgroundRects, maskThreshold);
1064 if (modelMask.empty())
1066 return EfficientSamMaskToFrameMask(modelMask, frame.size(), maskThreshold);
1070 return runPromptSet(prompts);
1072 cv::Mat combinedMask(frame.rows, frame.cols, CV_8U, cv::Scalar(0));
1073 bool hasMask =
false;
1079 cv::Mat rectMask = runPromptSet(rectPrompt);
1080 if (rectMask.empty())
1082 cv::bitwise_or(combinedMask, rectMask, combinedMask);
1091 cv::Mat pointMask = runPromptSet(pointPrompt);
1092 if (!pointMask.empty()) {
1093 cv::bitwise_or(combinedMask, pointMask, combinedMask);
1098 return hasMask ? combinedMask : cv::Mat();
1103 if (protobufDataPath.empty()) {
1104 std::cerr <<
"Missing path to object mask protobuf data file." << std::endl;
1110 pb_objdetect::ObjDetect objMessage;
1111 objMessage.add_classnames()->assign(
"object mask");
1114 AddFrameDataToProto(objMessage.add_frame(), frameData.second);
1116 *objMessage.mutable_last_updated() = TimeUtil::SecondsToTimestamp(time(NULL));
1118 std::fstream output(protobufDataPath, std::ios::out | std::ios::trunc | std::ios::binary);
1119 if (!objMessage.SerializeToOstream(&output)) {
1120 std::cerr <<
"Failed to write object mask protobuf message." << std::endl;
1127void CVObjectMask::AddFrameDataToProto(pb_objdetect::Frame* pbFrameData,
const CVObjectMaskFrameData& frameData)
1129 pbFrameData->set_id(frameData.
frameId);
1133 pb_objdetect::Frame_Box* box = pbFrameData->add_bounding_box();
1134 box->set_x(frameData.
box.x);
1135 box->set_y(frameData.
box.y);
1136 box->set_w(frameData.
box.width);
1137 box->set_h(frameData.
box.height);
1138 box->set_classid(0);
1139 box->set_confidence(frameData.
score);
1140 box->set_objectid(frameData.
objectId);
1142 pb_objdetect::Frame_Box_Mask* mask = box->mutable_mask();
1143 mask->set_width(frameData.
width);
1144 mask->set_height(frameData.
height);
1145 for (uint32_t count : frameData.
rle)
1146 mask->add_rle(count);
1153 }
catch (
const std::exception&) {
1154 std::cout <<
"JSON is invalid (missing keys or invalid data types)" << std::endl;
1160 if (!root[
"protobuf_data_path"].isNull())
1161 protobufDataPath = root[
"protobuf_data_path"].asString();
1162 if (!root[
"efficient_sam_model"].isNull())
1163 efficientSamModelPath = root[
"efficient_sam_model"].asString();
1164 if (!root[
"efficient_sam_model_path"].isNull())
1165 efficientSamModelPath = root[
"efficient_sam_model_path"].asString();
1166 if (!root[
"sam_model"].isNull())
1167 efficientSamModelPath = root[
"sam_model"].asString();
1168 if (!root[
"sam_model_path"].isNull())
1169 efficientSamModelPath = root[
"sam_model_path"].asString();
1170 if (!root[
"encoder_model"].isNull())
1171 efficientSamModelPath = root[
"encoder_model"].asString();
1172 if (!root[
"encoder_model_path"].isNull())
1173 efficientSamModelPath = root[
"encoder_model_path"].asString();
1174 if (!root[
"cutie_model_dir"].isNull())
1175 cutieModelDir = root[
"cutie_model_dir"].asString();
1176 if (!root[
"cutie_encode_key_model"].isNull())
1177 cutieEncodeKeyModelPath = root[
"cutie_encode_key_model"].asString();
1178 if (!root[
"cutie_encode_key_model_path"].isNull())
1179 cutieEncodeKeyModelPath = root[
"cutie_encode_key_model_path"].asString();
1180 if (!root[
"cutie_encode_value_model"].isNull())
1181 cutieEncodeValueModelPath = root[
"cutie_encode_value_model"].asString();
1182 if (!root[
"cutie_encode_value_model_path"].isNull())
1183 cutieEncodeValueModelPath = root[
"cutie_encode_value_model_path"].asString();
1184 if (!root[
"cutie_memory_readout_model"].isNull())
1185 cutieMemoryReadoutModelPath = root[
"cutie_memory_readout_model"].asString();
1186 if (!root[
"cutie_memory_readout_model_path"].isNull())
1187 cutieMemoryReadoutModelPath = root[
"cutie_memory_readout_model_path"].asString();
1188 if (!root[
"cutie_decode_model"].isNull())
1189 cutieDecodeModelPath = root[
"cutie_decode_model"].asString();
1190 if (!root[
"cutie_decode_model_path"].isNull())
1191 cutieDecodeModelPath = root[
"cutie_decode_model_path"].asString();
1192 if (!root[
"processing-device"].isNull())
1193 processingDevice = root[
"processing-device"].asString();
1194 if (!root[
"processing_device"].isNull())
1195 processingDevice = root[
"processing_device"].asString();
1196 if (!root[
"prompt_slots"].isNull())
1197 promptSlots = std::max(1, std::min(6, root[
"prompt_slots"].asInt()));
1198 if (!root[
"mask_threshold"].isNull())
1199 maskThreshold = root[
"mask_threshold"].asFloat();
1200 if (!root[
"model_size"].isNull())
1201 modelSize = root[
"model_size"].asInt();
1202 promptKeyframes.clear();
1203 if (!root[
"object_mask_selection"].isNull()) {
1204 const Json::Value& selection = root[
"object_mask_selection"];
1205 const Json::Value& frames = selection[
"frames"];
1206 if (frames.isObject()) {
1207 for (
const auto& frameName : frames.getMemberNames()) {
1208 const size_t frameNumber = JsonFrameNumber(frameName);
1209 if (frameNumber == 0)
1213 promptKeyframes[frameNumber] = prompts;
1219 if (!root[
"positive_points"].isNull())
1220 AppendJsonPoints(root[
"positive_points"], legacyPrompts.
positivePoints);
1221 if (!root[
"negative_points"].isNull())
1222 AppendJsonPoints(root[
"negative_points"], legacyPrompts.
negativePoints);
1224 if (!root[
"positive_x"].isNull() && !root[
"positive_y"].isNull()) {
1225 cv::Point2f point(root[
"positive_x"].asFloat(), root[
"positive_y"].asFloat());
1229 if (!root[
"negative_x"].isNull() && !root[
"negative_y"].isNull()) {
1230 cv::Point2f point(root[
"negative_x"].asFloat(), root[
"negative_y"].asFloat());
1234 if (!root[
"rect_x1"].isNull() && !root[
"rect_y1"].isNull() &&
1235 !root[
"rect_x2"].isNull() && !root[
"rect_y2"].isNull()) {
1237 rect[
"x1"] = root[
"rect_x1"];
1238 rect[
"y1"] = root[
"rect_y1"];
1239 rect[
"x2"] = root[
"rect_x2"];
1240 rect[
"y2"] = root[
"rect_y2"];
1241 cv::Rect_<float> parsed;
1242 if (RectFromJson(rect, parsed))
1246 promptKeyframes[1] = legacyPrompts;
Header file for CVObjectMask class.
Header file for all Exception classes.
void SetError(bool err, std::string message)
std::map< size_t, CVObjectMaskFrameData > masksData
void maskClip(openshot::Clip &video, size_t start=0, size_t end=0, bool process_interval=false)
std::shared_ptr< Frame > PreviewSeedMask(std::shared_ptr< Frame > frame)
void SetJson(const std::string value)
static std::string ValidateONNXModel(std::string modelPath)
void SetJsonValue(const Json::Value root)
float Start() const
Get start position (in seconds) of clip (trim start of video)
This class represents a clip (used to arrange readers on the timeline)
void Open() override
Open the internal reader.
float End() const override
Get end position (in seconds) of clip (trim end of video), which can be affected by the time curve.
std::shared_ptr< openshot::Frame > GetFrame(int64_t clip_frame_number) override
Get an openshot::Frame object for a specific frame number of this clip. The image size and number of ...
void Reader(openshot::ReaderBase *new_reader)
Set the current reader.
void Log(std::string message, Level level=LevelDebug)
static Logger * Instance()
This namespace is the default namespace for all code in the openshot library.
const Json::Value stringToJson(const std::string value)
std::vector< uint32_t > rle
bool HasPositivePrompt() const
std::vector< cv::Rect_< float > > positiveRects
std::vector< cv::Point2f > positivePoints
std::vector< cv::Rect_< float > > negativeRects
std::vector< cv::Point2f > negativePoints