Added capacity for multiple rectangle labels; bug fixes

1 files changed, 22 insertions, 10 deletions

diff --git a/src/playback.cc b/src/playback.cc
index 700567a..3203f47 100644
--- a/src/playback.cc
+++ b/src/playback.cc
@@ -5,13 +5,15 @@
 #include <stdexcept>
 #include <iostream>
 #include <thread>
+#include <map>
 
 namespace frontend {
     struct playback_impl {
         std::vector<cv::Mat> images;
-        int fps;
+        double fps;
         double reportedDuration;
         std::size_t frameNum = 0;
+        std::map<std::string, cv::Scalar> names2colors;
     };
 
     playback::playback(const std::filesystem::path& video, std::size_t frameCap) {
@@ -29,14 +31,21 @@ namespace frontend {
         }
         data->reportedDuration = cap.get(cv::CAP_PROP_POS_MSEC) / 1000.0;
         std::cout << "Reported duration: " << data->reportedDuration << " seconds" << std::endl;
-        cap.release();
     }
 
-    void playback::display(const std::string& windowName, const backend::rect &r) const {
-        cv::Mat curr = data->images[data->frameNum].clone();
-        if(r.nonzero()) {
-            // Update with the rectangle
-            cv::rectangle(curr, cv::Point(r.x1, r.y1), cv::Point(r.x2, r.y2), cv::Scalar(255, 0, 0), 2, cv::LINE_8);
+    void playback::display(const std::string& windowName, const std::vector<backend::label>& labels) const {
+        auto curr = data->images[data->frameNum].clone();
+        std::vector<cv::Scalar> colors = {cv::Scalar(255, 0, 0), cv::Scalar(0, 255, 0), cv::Scalar(0, 0, 255), cv::Scalar(255, 255, 0), cv::Scalar(255, 0, 255), cv::Scalar(0, 255, 255), cv::Scalar(255, 255, 255)};
+        for(auto l : labels) {
+            if(! data->names2colors.contains(l.name)) {
+                data->names2colors.insert({l.name, colors[data->names2colors.size()%colors.size()]});
+            }
+            auto color = data->names2colors.at(l.name);
+            auto r = l.location;
+            if(r.nonzero()) {
+                // Update with the rectangle
+                cv::rectangle(curr, cv::Point(r.x1, r.y1), cv::Point(r.x2, r.y2), color, 2, cv::LINE_8);
+            }
         }
         cv::imshow(windowName, curr);
     }
@@ -58,11 +67,11 @@ namespace frontend {
     }
 
     double playback::getTime() const {
-        return getFrame() / (double) data->fps;
+        return getFrame() / data->fps;
     }
     
     void playback::interFrameSleep() const {
-        std::this_thread::sleep_for(std::chrono::milliseconds(1000/data->fps));
+        std::this_thread::sleep_for(std::chrono::milliseconds(int(1000/data->fps)));
     }
 
     std::size_t playback::getMaxFrame() const {
@@ -70,6 +79,9 @@ namespace frontend {
     }
 
     double playback::getMaxTime() const {
-        return getMaxFrame() / (double) data->fps;
+        return getMaxFrame() / data->fps;
+    }
+
+    void playback::release() const {
     }
 }