diff --git a/include/CVStabilization.h b/include/CVStabilization.h index 45e15c605..249f518f1 100644 --- a/include/CVStabilization.h +++ b/include/CVStabilization.h @@ -96,7 +96,7 @@ class CVStabilization { ProcessingController *processingController; // Track current frame features and find the relative transformation - void TrackFrameFeatures(cv::Mat frame, size_t frameNum); + bool TrackFrameFeatures(cv::Mat frame, size_t frameNum); std::vector ComputeFramesTrajectory(); std::map SmoothTrajectory(std::vector &trajectory); diff --git a/include/effects/ObjectDetection.h b/include/effects/ObjectDetection.h index 81afa6854..5c1ef6f4b 100644 --- a/include/effects/ObjectDetection.h +++ b/include/effects/ObjectDetection.h @@ -36,6 +36,7 @@ #include #include #include +#include #include "../Color.h" #include "../Json.h" #include "../KeyFrame.h" @@ -70,6 +71,8 @@ namespace openshot std::string protobuf_data_path; std::map detectionsData; std::vector classNames; + + std::vector classesColor; /// Init effect settings void init_effect_details(); diff --git a/src/CVObjectDetection.cpp b/src/CVObjectDetection.cpp index 651d880ef..07d63e8d9 100644 --- a/src/CVObjectDetection.cpp +++ b/src/CVObjectDetection.cpp @@ -166,17 +166,21 @@ void CVObjectDetection::postprocess(const cv::Size &frameDims, const std::vector std::vector indices; cv::dnn::NMSBoxes(boxes, confidences, confThreshold, nmsThreshold, indices); - // std::vector sortBoxes; - // for(auto box : boxes) - // sortBoxes.push_back(box); - // sort.update(sortBoxes, frameId, sqrt(pow(frameDims.width,2) + pow(frameDims.height, 2))); + std::vector sortBoxes; + for(auto box : boxes) + sortBoxes.push_back(box); + sort.update(sortBoxes, frameId, sqrt(pow(frameDims.width,2) + pow(frameDims.height, 2)), confidences, classIds); - // sortBoxes.clear(); - // for(auto TBox : sort.frameTrackingResult) - // if(TBox.frame == frameId){ - // sortBoxes.push_back(TBox.box); - // } + sortBoxes.clear(); boxes.clear(); confidences.clear(); classIds.clear(); + + for(auto TBox : sort.frameTrackingResult){ + if(TBox.frame == frameId){ + boxes.push_back(TBox.box); + confidences.push_back(TBox.confidence); + classIds.push_back(TBox.classId); + } + } // for(int i = 0; iGetImageCV(); cv::cvtColor(cvimage, cvimage, cv::COLOR_RGB2GRAY); - TrackFrameFeatures(cvimage, frame_number); + if(! TrackFrameFeatures(cvimage, frame_number)) + prev_to_cur_transform.push_back(TransformParam(0, 0, 0)); // Update progress processingController->SetProgress(uint(100*(frame_number-start)/(end-start))); @@ -81,17 +84,17 @@ void CVStabilization::stabilizeClip(openshot::Clip& video, size_t _start, size_t } // Track current frame features and find the relative transformation -void CVStabilization::TrackFrameFeatures(cv::Mat frame, size_t frameNum){ +bool CVStabilization::TrackFrameFeatures(cv::Mat frame, size_t frameNum){ std::cout<<"frame "< status; std::vector err; // Extract new image features - cv::goodFeaturesToTrack(prev_grey, prev_corner, 200, 0.01, 30); + cv::goodFeaturesToTrack(prev_grey, prev_corner, 200, 0.01, 15); // Track features cv::calcOpticalFlowPyrLK(prev_grey, frame, prev_corner, cur_corner, status, err); // Remove untracked features + mediastatus+=status.size(); + if(status.size() > maiorstatus) + maiorstatus = status.size(); + for(size_t i=0; i < status.size(); i++) { if(status[i]) { prev_corner2.push_back(prev_corner[i]); @@ -114,18 +121,15 @@ void CVStabilization::TrackFrameFeatures(cv::Mat frame, size_t frameNum){ if(prev_corner2.empty() || cur_corner2.empty()){ last_T = cv::Mat(); prev_grey = cv::Mat(); - return; + return false; } // Translation + rotation only - cv::Mat T = estimateAffinePartial2D(prev_corner2, cur_corner2); // false = rigid transform, no scaling/shearing + cv::Mat T = cv::estimateAffinePartial2D(prev_corner2, cur_corner2); // false = rigid transform, no scaling/shearing double da, dx, dy; if(T.size().width == 0 || T.size().height == 0){ - - dx = 0; - dy = 0; - da = 0; + return false; } else{ // If no transformation is found, just use the last known good transform. @@ -137,10 +141,26 @@ void CVStabilization::TrackFrameFeatures(cv::Mat frame, size_t frameNum){ da = atan2(T.at(1,0), T.at(0,0)); } + if(dx > 100 || dy > 100 || da > 0.1){ + return false; + } + + mediax+=fabs(dx); + mediay+=fabs(dy); + mediaa+=fabs(da); + + if(fabs(dx) > maiorx) + maiorx = dx; + if(fabs(dy) > maiory) + maiory = dy; + if(fabs(da) > maiora) + maiora = da; + + std::cout< bboxes = {bbox}; + std::vector confidence = {1.0}; + std::vector classId = {1}; - sort.update(bboxes, frameId, sqrt(pow(frame.rows, 2) + pow(frame.cols, 2))); + sort.update(bboxes, frameId, sqrt(pow(frame.rows, 2) + pow(frame.cols, 2)), confidence, classId); for(auto TBox : sort.frameTrackingResult) bbox = TBox.box; diff --git a/src/effects/ObjectDetection.cpp b/src/effects/ObjectDetection.cpp index 20cfa12d7..545db9450 100644 --- a/src/effects/ObjectDetection.cpp +++ b/src/effects/ObjectDetection.cpp @@ -133,8 +133,9 @@ std::shared_ptr ObjectDetection::GetFrame(std::shared_ptr frame, i void ObjectDetection::drawPred(int classId, float conf, cv::Rect2d box, cv::Mat& frame) { + //Draw a rectangle displaying the bounding box - cv::rectangle(frame, box, cv::Scalar(255, 178, 50), 3); + cv::rectangle(frame, box, classesColor[classId], 2); //Get the label for the class name and its confidence std::string label = cv::format("%.2f", conf); @@ -144,12 +145,15 @@ void ObjectDetection::drawPred(int classId, float conf, cv::Rect2d box, cv::Mat& label = classNames[classId] + ":" + label; } - //Display the label at the box.y of the bounding box + //Display the label at the top of the bounding box int baseLine; - cv::Size labelSize = getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine); - box.y = std::max((int)box.y, labelSize.height); - cv::rectangle(frame, cv::Point(box.x, box.y - round(1.5*labelSize.height)), cv::Point(box.x + round(1.5*labelSize.width), box.y + baseLine), cv::Scalar(255, 255, 255), cv::FILLED); - putText(frame, label, cv::Point(box.x, box.y), cv::FONT_HERSHEY_SIMPLEX, 0.75, cv::Scalar(0,0,0),1); + 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(frame, cv::Point(left, top - round(1.025*labelSize.height)), cv::Point(left + round(1.025*labelSize.width), top + baseLine), classesColor[classId], cv::FILLED); + putText(frame, label, cv::Point(left+1, top), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0,0,0),1); } // Load protobuf data file @@ -172,6 +176,7 @@ bool ObjectDetection::LoadObjDetectdData(std::string inputFilePath){ for(int i = 0; i < objMessage.classnames_size(); i++){ classNames.push_back(objMessage.classnames(i)); + classesColor.push_back(cv::Scalar(std::rand()%205 + 50, std::rand()%205 + 50, std::rand()%205 + 50)); } // Iterate over all frames of the saved message diff --git a/src/sort_filter/KalmanTracker.h b/src/sort_filter/KalmanTracker.h index ac9b70296..03360f0da 100644 --- a/src/sort_filter/KalmanTracker.h +++ b/src/sort_filter/KalmanTracker.h @@ -24,7 +24,7 @@ class KalmanTracker m_id = kf_count; //kf_count++; } - KalmanTracker(StateType initRect) + KalmanTracker(StateType initRect, float confidence, int classId) : confidence(confidence), classId(classId) { init_kf(initRect); m_time_since_update = 0; @@ -54,6 +54,8 @@ class KalmanTracker int m_hit_streak; int m_age; int m_id; + float confidence; + int classId; private: void init_kf(StateType stateMat); diff --git a/src/sort_filter/sort.cpp b/src/sort_filter/sort.cpp index 44daf6761..299bfb1a0 100644 --- a/src/sort_filter/sort.cpp +++ b/src/sort_filter/sort.cpp @@ -37,7 +37,7 @@ double SortTracker::GetCentroidsDistance( return distance; } -void SortTracker::update(vector detections_cv, int frame_count, double image_diagonal) +void SortTracker::update(vector detections_cv, int frame_count, double image_diagonal, std::vector confidences, std::vector classIds) { vector detections; if (trackers.size() == 0) // the first frame met @@ -49,9 +49,11 @@ void SortTracker::update(vector detections_cv, int frame_count, double TrackingBox tb; tb.box = cv::Rect_(detections_cv[i]); + tb.classId = classIds[i]; + tb.confidence = confidences[i]; detections.push_back(tb); - KalmanTracker trk = KalmanTracker(detections[i].box); + KalmanTracker trk = KalmanTracker(detections[i].box, detections[i].confidence, detections[i].classId); trackers.push_back(trk); } return; @@ -62,6 +64,8 @@ void SortTracker::update(vector detections_cv, int frame_count, double { TrackingBox tb; tb.box = cv::Rect_(detections_cv[i]); + tb.classId = classIds[i]; + tb.confidence = confidences[i]; detections.push_back(tb); } for (auto it = frameTrackingResult.begin(); it != frameTrackingResult.end(); it++) @@ -156,12 +160,14 @@ void SortTracker::update(vector detections_cv, int frame_count, double trkIdx = matchedPairs[i].x; detIdx = matchedPairs[i].y; trackers[trkIdx].update(detections[detIdx].box); + trackers[trkIdx].classId = detections[detIdx].classId; + trackers[trkIdx].confidence = detections[detIdx].confidence; } // create and initialise new trackers for unmatched detections for (auto umd : unmatchedDetections) { - KalmanTracker tracker = KalmanTracker(detections[umd].box); + KalmanTracker tracker = KalmanTracker(detections[umd].box, detections[umd].confidence, detections[umd].classId); trackers.push_back(tracker); } @@ -189,6 +195,8 @@ void SortTracker::update(vector detections_cv, int frame_count, double res.box = trackers[i].get_state(); res.id = trackers[i].m_id; res.frame = frame_count; + res.classId = trackers[i].classId; + res.confidence = trackers[i].confidence; frameTrackingResult.push_back(res); } diff --git a/src/sort_filter/sort.hpp b/src/sort_filter/sort.hpp index eaa313375..4eab6dcda 100644 --- a/src/sort_filter/sort.hpp +++ b/src/sort_filter/sort.hpp @@ -32,7 +32,7 @@ class SortTracker // Initialize tracker // Update position based on the new frame - void update(std::vector detection, int frame_count, double image_diagonal); + void update(std::vector detection, int frame_count, double image_diagonal, std::vector confidences, std::vector classIds); double GetIOU(cv::Rect_ bb_test, cv::Rect_ bb_gt); double GetCentroidsDistance(cv::Rect_ bb_test, cv::Rect_ bb_gt); std::vector trackers;