Skip to content

Commit

Permalink
Refactor of ObjectDetection effect:
Browse files Browse the repository at this point in the history 
- Draw bounding boxes using QPainter and support corner radius (and faster drawing performance by drawing directly on the frame)
- Draw child clips with correct aspect ratio
- Adding "Yes/No" option for "Draw Text"
- default to fully opaque stroke
  • Loading branch information
@jonoomph
jonoomph committed Feb 23, 2024
1 parent 3af6e1f commit 98383fd
Show file tree
Hide file tree
Showing 5 changed files with 123 additions and 224 deletions.
2 changes: 1 addition & 1 deletion src/TrackedObjectBBox.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ TrackedObjectBBox::TrackedObjectBBox(int Red, int Green, int Blue, int Alfa)
: delta_x(0.0), delta_y(0.0),
scale_x(1.0), scale_y(1.0), rotation(0.0),
background_alpha(0.0), background_corner(12),
stroke_width(2) , stroke_alpha(0.5),
stroke_width(2) , stroke_alpha(0.70),
stroke(Red, Green, Blue, Alfa),
background(0, 0, 255, Alfa)
{
Expand Down
329 changes: 119 additions & 210 deletions src/effects/ObjectDetection.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
#include "Exceptions.h"
#include "Timeline.h"
#include "objdetectdata.pb.h"
#include "Tracker.h"

#include <QImage>
#include <QPainter>
Expand Down Expand Up @@ -67,213 +68,120 @@ void ObjectDetection::init_effect_details()

// This method is required for all derived classes of EffectBase, and returns a
// modified openshot::Frame object
std::shared_ptr<Frame> ObjectDetection::GetFrame(std::shared_ptr<Frame> frame, int64_t frame_number)
{
// Get the frame's image
cv::Mat cv_image = frame->GetImageCV();

// Check if frame isn't NULL
if(cv_image.empty()){
return frame;
}

// Initialize the Qt rectangle that will hold the positions of the bounding-box
std::vector<QRectF> boxRects;
// Initialize the image of the TrackedObject child clip
std::vector<std::shared_ptr<QImage>> childClipImages;

// Check if track data exists for the requested frame
if (detectionsData.find(frame_number) != detectionsData.end()) {
float fw = cv_image.size().width;
float fh = cv_image.size().height;

DetectionData detections = detectionsData[frame_number];
for(int i = 0; i<detections.boxes.size(); i++){

// Does not show boxes with confidence below the threshold
if(detections.confidences.at(i) < confidence_threshold){
continue;
}
// Just display selected classes
if( display_classes.size() > 0 &&
std::find(display_classes.begin(), display_classes.end(), classNames[detections.classIds.at(i)]) == display_classes.end()){
continue;
}

// Get the object id
int objectId = detections.objectIds.at(i);

// Search for the object in the trackedObjects map
auto trackedObject_it = trackedObjects.find(objectId);

// Cast the object as TrackedObjectBBox
std::shared_ptr<TrackedObjectBBox> trackedObject = std::static_pointer_cast<TrackedObjectBBox>(trackedObject_it->second);

// Check if the tracked object has data for this frame
if (trackedObject->Contains(frame_number) &&
trackedObject->visible.GetValue(frame_number) == 1)
{
// Get the bounding-box of given frame
BBox trackedBox = trackedObject->GetBox(frame_number);
bool draw_text = !display_box_text.GetValue(frame_number);
std::vector<int> stroke_rgba = trackedObject->stroke.GetColorRGBA(frame_number);
int stroke_width = trackedObject->stroke_width.GetValue(frame_number);
float stroke_alpha = trackedObject->stroke_alpha.GetValue(frame_number);
std::vector<int> bg_rgba = trackedObject->background.GetColorRGBA(frame_number);
float bg_alpha = trackedObject->background_alpha.GetValue(frame_number);

cv::Rect2d box(
(int)( (trackedBox.cx-trackedBox.width/2)*fw),
(int)( (trackedBox.cy-trackedBox.height/2)*fh),
(int)( trackedBox.width*fw),
(int)( trackedBox.height*fh)
);

// If the Draw Box property is off, then make the box invisible
if (trackedObject->draw_box.GetValue(frame_number) == 0)
{
bg_alpha = 1.0;
stroke_alpha = 1.0;
}

drawPred(detections.classIds.at(i), detections.confidences.at(i),
box, cv_image, detections.objectIds.at(i), bg_rgba, bg_alpha, 1, true, draw_text);
drawPred(detections.classIds.at(i), detections.confidences.at(i),
box, cv_image, detections.objectIds.at(i), stroke_rgba, stroke_alpha, stroke_width, false, draw_text);


// Get the Detected Object's child clip
if (trackedObject->ChildClipId() != ""){
// Cast the parent timeline of this effect
Timeline* parentTimeline = static_cast<Timeline *>(ParentTimeline());
if (parentTimeline){
// Get the Tracked Object's child clip
Clip* childClip = parentTimeline->GetClip(trackedObject->ChildClipId());

if (childClip){
// Get the image of the child clip for this frame
std::shared_ptr<Frame> childClipFrame = childClip->GetFrame(frame_number);
childClipImages.push_back(childClipFrame->GetImage());

// Set the Qt rectangle with the bounding-box properties
QRectF boxRect;
boxRect.setRect((int)((trackedBox.cx-trackedBox.width/2)*fw),
(int)((trackedBox.cy - trackedBox.height/2)*fh),
(int)(trackedBox.width*fw),
(int)(trackedBox.height*fh));
boxRects.push_back(boxRect);
}
}
}
}
}
}

// Update Qt image with new Opencv frame
frame->SetImageCV(cv_image);

// Set the bounding-box image with the Tracked Object's child clip image
if(boxRects.size() > 0){
// Get the frame image
QImage frameImage = *(frame->GetImage());
for(int i; i < boxRects.size();i++){
// Set a Qt painter to the frame image
QPainter painter(&frameImage);
// Draw the child clip image inside the bounding-box
painter.drawImage(boxRects[i], *childClipImages[i]);
}
// Set the frame image as the composed image
frame->AddImage(std::make_shared<QImage>(frameImage));
}

return frame;
}

void ObjectDetection::DrawRectangleRGBA(cv::Mat &frame_image, cv::RotatedRect box, std::vector<int> color, float alpha,
int thickness, bool is_background){
// Get the bouding box vertices
cv::Point2f vertices2f[4];
box.points(vertices2f);

// TODO: take a rectangle of frame_image by refencence and draw on top of that to improve speed
// select min enclosing rectangle to draw on a small portion of the image
// cv::Rect rect = box.boundingRect();
// cv::Mat image = frame_image(rect)

if(is_background){
cv::Mat overlayFrame;
frame_image.copyTo(overlayFrame);

// draw bounding box background
cv::Point vertices[4];
for(int i = 0; i < 4; ++i){
vertices[i] = vertices2f[i];}

cv::Rect rect = box.boundingRect();
cv::fillConvexPoly(overlayFrame, vertices, 4, cv::Scalar(color[2],color[1],color[0]), cv::LINE_AA);
// add opacity
cv::addWeighted(overlayFrame, 1-alpha, frame_image, alpha, 0, frame_image);
}
else{
cv::Mat overlayFrame;
frame_image.copyTo(overlayFrame);

// Draw bounding box
for (int i = 0; i < 4; i++)
{
cv::line(overlayFrame, vertices2f[i], vertices2f[(i+1)%4], cv::Scalar(color[2],color[1],color[0]),
thickness, cv::LINE_AA);
}

// add opacity
cv::addWeighted(overlayFrame, 1-alpha, frame_image, alpha, 0, frame_image);
}
}

void ObjectDetection::drawPred(int classId, float conf, cv::Rect2d box, cv::Mat& frame, int objectNumber, std::vector<int> color,
float alpha, int thickness, bool is_background, bool display_text)
{

if(is_background){
cv::Mat overlayFrame;
frame.copyTo(overlayFrame);

//Draw a rectangle displaying the bounding box
cv::rectangle(overlayFrame, box, cv::Scalar(color[2],color[1],color[0]), cv::FILLED);

// add opacity
cv::addWeighted(overlayFrame, 1-alpha, frame, alpha, 0, frame);
}
else{
cv::Mat overlayFrame;
frame.copyTo(overlayFrame);

//Draw a rectangle displaying the bounding box
cv::rectangle(overlayFrame, box, cv::Scalar(color[2],color[1],color[0]), thickness);

if(display_text){
//Get the label for the class name and its confidence
std::string label = cv::format("%.2f", conf);
if (!classNames.empty())
{
CV_Assert(classId < (int)classNames.size());
label = classNames[classId] + ":" + label;
}

//Display the label at the top of the bounding box
int baseLine;
cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);

double left = box.x;
double top = std::max((int)box.y, labelSize.height);

cv::rectangle(overlayFrame, cv::Point(left, top - round(1.025*labelSize.height)), cv::Point(left + round(1.025*labelSize.width), top + baseLine),
cv::Scalar(color[2],color[1],color[0]), cv::FILLED);
putText(overlayFrame, label, cv::Point(left+1, top), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0,0,0),1);
}
// add opacity
cv::addWeighted(overlayFrame, 1-alpha, frame, alpha, 0, frame);
}
std::shared_ptr<Frame> ObjectDetection::GetFrame(std::shared_ptr<Frame> frame, int64_t frame_number) {
// Get the frame's QImage
std::shared_ptr<QImage> frame_image = frame->GetImage();

// Check if frame isn't NULL
if(!frame_image || frame_image->isNull()) {
return frame;
}

QPainter painter(frame_image.get());
painter.setRenderHints(QPainter::Antialiasing | QPainter::SmoothPixmapTransform);

if (detectionsData.find(frame_number) != detectionsData.end()) {
float fw = frame_image->width();
float fh = frame_image->height();

DetectionData detections = detectionsData[frame_number];
for (int i = 0; i < detections.boxes.size(); i++) {
if (detections.confidences.at(i) < confidence_threshold ||
(!display_classes.empty() &&
std::find(display_classes.begin(), display_classes.end(), classNames[detections.classIds.at(i)]) == display_classes.end())) {
continue;
}

int objectId = detections.objectIds.at(i);
auto trackedObject_it = trackedObjects.find(objectId);

if (trackedObject_it != trackedObjects.end()) {
std::shared_ptr<TrackedObjectBBox> trackedObject = std::static_pointer_cast<TrackedObjectBBox>(trackedObject_it->second);

if (trackedObject->Contains(frame_number) && trackedObject->visible.GetValue(frame_number) == 1) {
BBox trackedBox = trackedObject->GetBox(frame_number);

QRectF boxRect((trackedBox.cx - trackedBox.width / 2) * fw,
(trackedBox.cy - trackedBox.height / 2) * fh,
trackedBox.width * fw,
trackedBox.height * fh);

if (trackedObject->draw_box.GetValue(frame_number) == 1) {
// Draw bounding box
bool display_text = !display_box_text.GetValue(frame_number);
std::vector<int> stroke_rgba = trackedObject->stroke.GetColorRGBA(frame_number);
std::vector<int> bg_rgba = trackedObject->background.GetColorRGBA(frame_number);
int stroke_width = trackedObject->stroke_width.GetValue(frame_number);
float stroke_alpha = trackedObject->stroke_alpha.GetValue(frame_number);
float bg_alpha = trackedObject->background_alpha.GetValue(frame_number);
float bg_corner = trackedObject->background_corner.GetValue(frame_number);

// Set the pen for the border
QPen pen(QColor(stroke_rgba[0], stroke_rgba[1], stroke_rgba[2], 255 * stroke_alpha));
pen.setWidth(stroke_width);
painter.setPen(pen);

// Set the brush for the background
QBrush brush(QColor(bg_rgba[0], bg_rgba[1], bg_rgba[2], 255 * bg_alpha));
painter.setBrush(brush);

// Draw the rounded rectangle
painter.drawRoundedRect(boxRect, bg_corner, bg_corner);

if(display_text) {
// Draw text label above bounding box
// Get the confidence and classId for the current detection
float conf = detections.confidences.at(i);
int classId = detections.classIds.at(i);

// Get the label for the class name and its confidence
QString label = QString::number(conf, 'f', 2); // Format confidence to two decimal places
if (!classNames.empty()) {
label = QString::fromStdString(classNames[classId]) + ":" + label;
}

// Set up the painter, font, and pen
QFont font;
font.setPixelSize(14);
painter.setFont(font);

// Calculate the size of the text
QFontMetrics fontMetrics(font);
QSize labelSize = fontMetrics.size(Qt::TextSingleLine, label);

// Define the top left point of the rectangle
double left = boxRect.center().x() - (labelSize.width() / 2.0);
double top = std::max(static_cast<int>(boxRect.top()), labelSize.height()) - 4.0;

// Draw the text
painter.drawText(QPointF(left, top), label);
}
}

// Get the image of the Tracked Object' child clip
if (trackedObject->ChildClipId() != "") {
Timeline* parentTimeline = static_cast<Timeline *>(ParentTimeline());
if (parentTimeline) {
Clip* childClip = parentTimeline->GetClip(trackedObject->ChildClipId());
if (childClip) {
std::shared_ptr<Frame> childClipFrame = childClip->GetFrame(frame_number);
std::shared_ptr<QImage> childClipImage = childClipFrame->GetImage();

// Scale the original bounding box to this image
QRectF scaledRect = Tracker::scaleAndCenterRect(QRectF(childClipImage->rect()), boxRect);
painter.drawImage(scaledRect, *childClipImage);
}
}
}
}
}
}
}

painter.end();

// The frame's QImage has been modified in place, so we just return the original frame
return frame;
}

// Load protobuf data file
Expand Down Expand Up @@ -348,6 +256,7 @@ bool ObjectDetection::LoadObjDetectdData(std::string inputFilePath){
{
// There is no tracked object with that id, so insert a new one
TrackedObjectBBox trackedObj((int)classesColor[classId](0), (int)classesColor[classId](1), (int)classesColor[classId](2), (int)0);
trackedObj.stroke_alpha = Keyframe(1.0);
trackedObj.AddBox(id, x+(w/2), y+(h/2), w, h, 0.0);

std::shared_ptr<TrackedObjectBBox> trackedObjPtr = std::make_shared<TrackedObjectBBox>(trackedObj);
Expand Down Expand Up @@ -551,9 +460,9 @@ std::string ObjectDetection::PropertiesJSON(int64_t requested_frame) const {
root["confidence_threshold"] = add_property_json("Confidence Theshold", confidence_threshold, "float", "", NULL, 0, 1, false, requested_frame);
root["class_filter"] = add_property_json("Class Filter", 0.0, "string", class_filter, NULL, -1, -1, false, requested_frame);

root["display_box_text"] = add_property_json("Draw Box Text", display_box_text.GetValue(requested_frame), "int", "", &display_box_text, 0, 1.0, false, requested_frame);
root["display_box_text"]["choices"].append(add_property_choice_json("Off", 1, display_box_text.GetValue(requested_frame)));
root["display_box_text"]["choices"].append(add_property_choice_json("On", 0, display_box_text.GetValue(requested_frame)));
root["display_box_text"] = add_property_json("Draw Box Text", display_box_text.GetValue(requested_frame), "int", "", &display_box_text, 0, 1, false, requested_frame);
root["display_box_text"]["choices"].append(add_property_choice_json("Yes", true, display_box_text.GetValue(requested_frame)));
root["display_box_text"]["choices"].append(add_property_choice_json("No", false, display_box_text.GetValue(requested_frame)));

// Return formatted string
return root.toStyledString();
Expand Down
7 changes: 0 additions & 7 deletions src/effects/ObjectDetection.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,6 @@ namespace openshot
std::string protobuf_data_path;
std::map<size_t, DetectionData> detectionsData;
std::vector<std::string> classNames;

std::vector<cv::Scalar> classesColor;

/// Draw class name and confidence score on top of the bounding box
Expand All @@ -73,12 +72,6 @@ namespace openshot

/// Init effect settings
void init_effect_details();
/// Draw bounding box with class and score text
void drawPred(int classId, float conf, cv::Rect2d box, cv::Mat& frame, int objectNumber, std::vector<int> color, float alpha,
int thickness, bool is_background, bool draw_text);
/// Draw rotated rectangle with alpha channel
void DrawRectangleRGBA(cv::Mat &frame_image, cv::RotatedRect box, std::vector<int> color, float alpha, int thickness, bool is_background);


public:
/// Index of the Tracked Object that was selected to modify it's properties
Expand Down
3 changes: 0 additions & 3 deletions src/effects/Tracker.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -142,9 +142,6 @@ std::shared_ptr<Frame> Tracker::GetFrame(std::shared_ptr<Frame> frame, int64_t f

// Scale the original bounding box to this image
QRectF scaledRect = scaleAndCenterRect(QRectF(childClipImage->rect()), boxRect);
QImage scaledImage = childClipImage->scaled(scaledRect.size().toSize(), Qt::KeepAspectRatio, Qt::SmoothTransformation);

// Draw the child clip image inside the bounding-box
painter.drawImage(scaledRect, *childClipImage);
}
}
Expand Down
Loading

0 comments on commit 98383fd

Please sign in to comment.