5 files changed, 35 insertions, 12 deletions

diff --git a/src/annotator.cc b/src/annotator.cc
index f1fccbc..1a6d557 100644
--- a/src/annotator.cc
+++ b/src/annotator.cc
@@ -29,14 +29,18 @@ int main(int argc, char *argv[]) {
     std::vector<std::string> flags = extractFlags(args);
     if(args.empty()) {
         std::cout << "Must provide a path to a video to process!" << std::endl;
+        return 1;
+    }
+    int frameCap = -1;
+    if(args.size() >= 2) {
+        frameCap = std::stoi(args[1]);
+        std::cout << "Only using first " << frameCap << " frames." << std::endl;
     }
     std::cout << "Loading video " << args[0] << std::endl;
     std::filesystem::path p(args[0]);
-    std::shared_ptr<frontend::playback> playback(new frontend::playback(p));
+    std::shared_ptr<frontend::playback> playback(new frontend::playback(p, frameCap));
     // Format save path
-    auto savedir = std::filesystem::path("save") / p.parent_path().filename();
-    std::filesystem::create_directories(savedir);
-    auto savepath = savedir / (p.stem().string() + ".csv");
+    auto savepath = std::filesystem::path("saves") / p.parent_path().filename() / (p.stem().string() + ".csv");
     std::cout << "Saves are written to: " << savepath << std::endl;
     std::shared_ptr<backend::labeller> labeller(new backend::labeller(savepath));
     frontend::ui ui(playback, labeller);
diff --git a/src/labeller.cc b/src/labeller.cc
index 024b720..f7fa69e 100644
--- a/src/labeller.cc
+++ b/src/labeller.cc
@@ -1,7 +1,7 @@
 #include "labeller.h"
 #include "settings.h"
 #include <fstream>
-
+#include <iostream>
 
 namespace backend {
     struct labeller_impl {
@@ -14,11 +14,28 @@ namespace backend {
         std::filesystem::path savepath;
     };
 
+    void load(const std::filesystem::path& savepath, labeller_impl& data) {
+        std::ifstream in(savepath);
+        std::string line;
+        const char delim = ',';
+        while(std::getline(in, line)) {
+            std::size_t split = line.find(delim);
+            if(split == std::string::npos) {
+                std::cerr << "Error reading " << savepath << std::endl;
+                throw std::runtime_error(savepath.string());
+            }
+            data.annotations.push_back(label(line.substr(0, split), std::stod(line.substr(split+1))));
+        }
+    }
+
     labeller::labeller(const std::filesystem::path& savepath) {
         data = std::shared_ptr<labeller_impl>(new labeller_impl);
         data->labels = settings::getLabels();
         data->rectangleLabels = settings::getRectangleLabels();
         data->savepath = savepath;
+        if(std::filesystem::directory_entry(savepath).exists()) {
+            load(savepath, *data);
+        }
     }
 
     std::pair<label, label> labeller::getSurrounding(double time) const {
@@ -111,6 +128,7 @@ namespace backend {
 
     void labeller::save() const {
         std::ofstream out;
+        std::filesystem::create_directories(data->savepath.parent_path());
         out.open(data->savepath);
         std::vector<label> a(data->annotations);
         std::sort(a.begin(), a.end(), compareLabels);
@@ -119,4 +137,5 @@ namespace backend {
         }
         out.close();
     }
+
 }
diff --git a/src/playback.cc b/src/playback.cc
index 045abb6..8839528 100644
--- a/src/playback.cc
+++ b/src/playback.cc
@@ -13,14 +13,14 @@ namespace frontend {
         std::size_t frameNum = 0;
     };
 
-    playback::playback(const std::filesystem::path& video) {
+    playback::playback(const std::filesystem::path& video, std::size_t frameCap) {
         data = std::shared_ptr<playback_impl>(new playback_impl);
         cv::VideoCapture cap(video);
         if(!cap.isOpened()) {
             throw std::runtime_error("Error loading video " + video.string());
         }
         data->fps = cap.get(cv::CAP_PROP_FPS);
-        while(true) {
+        while(frameCap == 0 or data->images.size() < frameCap) {
             cv::Mat frame;
             cap >> frame;
             if(frame.empty()) break;
diff --git a/src/playback.h b/src/playback.h
index f0f1f09..2fd5fa9 100644
--- a/src/playback.h
+++ b/src/playback.h
@@ -7,7 +7,7 @@ namespace frontend {
     struct playback_impl;
     class playback {
         public:
-            playback(const std::filesystem::path& video);
+            playback(const std::filesystem::path& video, std::size_t frameCap=0);
             void display(const std::string& windowName) const;
             bool seekFrame(std::size_t frameNum);
             std::size_t getFrame() const;
diff --git a/src/ui.cc b/src/ui.cc
index 858d74c..f852bde 100644
--- a/src/ui.cc
+++ b/src/ui.cc
@@ -10,7 +10,7 @@ namespace frontend {
     struct ui_impl {
         std::shared_ptr<playback> pb;
         std::shared_ptr<backend::labeller> llr;
-        bool play = true;
+        bool play = false;
         bool stalePrintout = true;
     };
 
@@ -25,8 +25,8 @@ namespace frontend {
             data->stalePrintout = true;
             if(keycode >= 48 && keycode <= 57) { // Number -> label
                 std::size_t num = keycode - 48;
-                if(num < data->llr->getLabels().size()) {
-                    data->llr->applyLabel(data->llr->getLabels()[num], data->pb->getTime());
+                if(num <= data->llr->getLabels().size()) {
+                    data->llr->applyLabel(data->llr->getLabels()[num-1], data->pb->getTime());
                 }
             } else {
                 switch(keycode) {
@@ -104,7 +104,7 @@ namespace frontend {
         std::cout << "Annotations:" << std::endl;
         int num = 0;
         for(auto ann : data->llr->getLabels()) {
-            std::cout << num++ << ": " << ann << std::endl;
+            std::cout << ++num << ": " << ann << std::endl;
         }
         // Get window size
         struct winsize size;