Refactor of ObjectDetection effect:

jonoomph · jonoomph · commit 98383fdddd26 · 2024-02-22T20:47:21.000-06:00
- 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
diff --git a/src/TrackedObjectBBox.cpp b/src/TrackedObjectBBox.cpp
@@ -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)
 {
diff --git a/src/effects/ObjectDetection.cpp b/src/effects/ObjectDetection.cpp
@@ -19,6 +19,7 @@
 #include "Exceptions.h"
 #include "Timeline.h"
 #include "objdetectdata.pb.h"
+#include "Tracker.h"
 
 #include <QImage>
 #include <QPainter>
@@ -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
@@ -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);
@@ -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();
diff --git a/src/effects/ObjectDetection.h b/src/effects/ObjectDetection.h
@@ -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
@@ -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
diff --git a/src/effects/Tracker.cpp b/src/effects/Tracker.cpp
@@ -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);
                 }
             }
diff --git a/src/effects/Tracker.h b/src/effects/Tracker.h

Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,7 @@ TrackedObjectBBox::TrackedObjectBBox(int Red, int Green, int Blue, int Alfa)`
`34`	`34`	`: delta_x(0.0), delta_y(0.0),`
`35`	`35`	`scale_x(1.0), scale_y(1.0), rotation(0.0),`
`36`	`36`	`background_alpha(0.0), background_corner(12),`
`37`		`- stroke_width(2) , stroke_alpha(0.5),`
	`37`	`+ stroke_width(2) , stroke_alpha(0.70),`
`38`	`38`	`stroke(Red, Green, Blue, Alfa),`
`39`	`39`	`background(0, 0, 255, Alfa)`
`40`	`40`	`{`
Original file line number	Diff line number	Diff line change
`@@ -142,9 +142,6 @@ std::shared_ptr<Frame> Tracker::GetFrame(std::shared_ptr<Frame> frame, int64_t f`
`142`	`142`
`143`	`143`	`// Scale the original bounding box to this image`
`144`	`144`	`QRectF scaledRect = scaleAndCenterRect(QRectF(childClipImage->rect()), boxRect);`
`145`		`- QImage scaledImage = childClipImage->scaled(scaledRect.size().toSize(), Qt::KeepAspectRatio, Qt::SmoothTransformation);`
`146`		`-`
`147`		`- // Draw the child clip image inside the bounding-box`
`148`	`145`	`painter.drawImage(scaledRect, *childClipImage);`
`149`	`146`	`}`
`150`	`147`	`}`