Added capacity for multiple rectangle labels; bug fixes

1 files changed, 69 insertions, 24 deletions

diff --git a/src/ui.cc b/src/ui.cc
index 3940955..541444a 100644
--- a/src/ui.cc
+++ b/src/ui.cc
@@ -12,12 +12,64 @@ namespace frontend {
         std::shared_ptr<backend::labeller> llr;
         bool play = false;
         bool stalePrintout = true;
+        int rectIndex = -1;
     };
 
     ui::ui(std::shared_ptr<playback> playback, std::shared_ptr<backend::labeller> labeller) {
         data = std::shared_ptr<ui_impl>(new ui_impl);
         data->pb = playback;
         data->llr = labeller;
+        if(!data->llr->getRectangleLabels().empty()) {
+            data->rectIndex = 0;
+        }
+    }
+
+    std::string getCurrentRectName(ui_impl* data) {
+        std::string name = "Rectangle";
+        if(data->rectIndex >= 0) {
+            name = data->llr->getRectangleLabels()[data->rectIndex];
+        }
+        return name;
+    }
+
+    std::pair<backend::label*, backend::label*> getSurroundingLabels(ui_impl* data, std::string name) {
+        std::pair<backend::label*, backend::label*> ret = {nullptr, nullptr};
+        auto &labs = data->llr->getEditableLabels();
+        for(auto it = labs.begin(); it != labs.end(); it++) {
+            if(it->name != name) {
+                continue;
+            }
+            if(it->time <= data->pb->getTime() && (!ret.first || ret.first->time < it->time)) {
+                ret.first = &(*it);
+            }
+            if(it->time > data->pb->getTime() && (!ret.second || ret.second->time > it->time)) {
+                ret.second = &(*it);
+            }
+        }
+        return ret;
+    }
+
+    std::vector<backend::label> getInterpolatedRectLabels(ui_impl* data) {
+        std::vector<backend::label> ret;
+        for(auto name : data->llr->getRectangleLabels()) {
+            auto p = getSurroundingLabels(data, name);
+            if(p.first && ! p.second) {
+                ret.push_back(*p.first);
+            } else if(p.first && p.second) {
+                backend::label l(p.first->name, data->pb->getTime());
+                double weight = (p.second->time-l.time)/(p.second->time-p.first->time);
+                l.location.x1 = weight*p.first->location.x1 + (1-weight)*p.second->location.x1;
+                l.location.x2 = weight*p.first->location.x2 + (1-weight)*p.second->location.x2;
+                l.location.y1 = weight*p.first->location.y1 + (1-weight)*p.second->location.y1;
+                l.location.y2 = weight*p.first->location.y2 + (1-weight)*p.second->location.y2;
+                ret.push_back(l);
+            }
+        }
+        return ret;
+    }
+
+    backend::label* getCurrentRectLabel(ui_impl* data) {
+        return getSurroundingLabels(data, getCurrentRectName(data)).first;
     }
 
     bool handleKey(int keycode, const std::shared_ptr<ui_impl>& data) {
@@ -31,6 +83,7 @@ namespace frontend {
             } else {
                 switch(keycode) {
                     case 27: // Escape quits
+                        data->pb->release();
                         std::cout << std::endl;
                         return true;
                     case 32: // Space pauses
@@ -68,6 +121,12 @@ namespace frontend {
                     case 115: // s saves
                         data->llr->save();
                         break;
+                    case 99: // c cycles rect labels
+                        if(!data->llr->getRectangleLabels().empty()) {
+                            data->rectIndex = (data->rectIndex + 1) % data->llr->getRectangleLabels().size();
+                        }
+                        std::cout << "Set rect name to " << getCurrentRectName(data.get()) << std::endl;
+                        break;
                     default:
                         std::cout << "Pressed the " << keycode << " key" << std::endl;
                         break;
@@ -77,16 +136,6 @@ namespace frontend {
         return false;
     }
 
-    backend::label* getCurrentRectLabel(ui_impl *data) {
-        auto &labs = data->llr->getEditableLabels();
-        for(auto rit = labs.rbegin(); rit != labs.rend(); rit++) {
-            if(rit->location.nonzero() && rit->time <= data->pb->getTime()) {
-                return &(*rit);
-            }
-        }
-        return nullptr;
-    }
-
     void mouseCallback(int event, int x, int y, int flags, void* userdata) {
         // userdata is actually a &shared_ptr<ui_impl>
         auto data = static_cast<ui_impl*>(userdata);
@@ -100,14 +149,7 @@ namespace frontend {
             if(lab && lab->time == time) {
                 // Use this one
             } else { // Make one
-                // Try to use the first rectangle annotation, fall back on normal, then "Rectangle"
-                std::string name = "Rectangle";
-                if(!data->llr->getRectangleLabels().empty()) {
-                    name = data->llr->getRectangleLabels()[0];
-                } else if(!data->llr->getLabels().empty()) {
-                    name = data->llr->getLabels()[0];
-                }
-                data->llr->applyLabel(name, time);
+                data->llr->applyLabel(getCurrentRectName(data), time);
                 lab = &data->llr->getEditableLabels().back();
             }
             lab->location.x1 = lab->location.x2 = x;
@@ -118,10 +160,13 @@ namespace frontend {
         } else if(event == cv::EVENT_MOUSEMOVE) {
             // If the mouse is down, update x2 and y2 of current window
             if(mouseDown) {
-                auto lab = getCurrentRectLabel(data); // It will be valid! TODO: Will it?
-                //std::cout << "Updating location from " << lab->location.x2 << ", " << lab->location.y2 << " to " << x << ", " << y << "\n";
-                lab->location.x2 = x;
-                lab->location.y2 = y;
+                auto lab = getCurrentRectLabel(data);
+                if(lab) {
+                    lab->location.x2 = x;
+                    lab->location.y2 = y;
+                } else {
+                    std::cout << "ERROR: Invalid label!" << std::endl;
+                }
             }
         }
     }
@@ -138,10 +183,10 @@ namespace frontend {
         ioctl(STDOUT_FILENO, TIOCGWINSZ, &size);
         int cols = size.ws_col;
         // Register mouse callback
-        data->pb->display("Video", (getCurrentRectLabel(data.get()))? getCurrentRectLabel(data.get())->location : backend::rect());
+        data->pb->display("Video", getInterpolatedRectLabels(data.get()));
         cv::setMouseCallback("Video", mouseCallback, data.get());
         while(true) {
-            data->pb->display("Video", (getCurrentRectLabel(data.get()))? getCurrentRectLabel(data.get())->location : backend::rect());
+            data->pb->display("Video", getInterpolatedRectLabels(data.get()));
             data->pb->interFrameSleep();
             if(handleKey(cv::pollKey(), data)) {
                 break;