Initial commit

9 files changed, 454 insertions, 0 deletions

diff --git a/src/annotator.cc b/src/annotator.cc
new file mode 100644
index 0000000..c74f04f
--- /dev/null
+++ b/src/annotator.cc
@@ -0,0 +1,45 @@
+#include "playback.h"
+#include "ui.h"
+#include "labeller.h"
+#include <vector>
+#include <iostream>
+#include <memory>
+
+// Removes flags from args (in-place) and returns vector of flags
+std::vector<std::string> extractFlags(std::vector<std::string>& args) {
+    std::vector<std::string> ret;
+    auto it = args.begin();
+    while(it != args.end()) {
+        if((*it)[0] == '-') {
+            while((*it)[0] == '-') {
+                (*it).erase((*it).begin());
+            }
+            ret.push_back(*it);
+            args.erase(it);
+        } else {
+            it++;
+        }
+    }
+    return ret;
+}
+
+int main(int argc, char *argv[]) {
+    std::string exename = argv[0];
+    std::vector<std::string> args(&argv[1], &argv[argc]);
+    std::vector<std::string> flags = extractFlags(args);
+    if(args.empty()) {
+        std::cout << "Must provide a path to a video to process!" << 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));
+    // Format save path
+    auto savedir = std::filesystem::path("save") / p.parent_path().filename();
+    std::filesystem::create_directories(savedir);
+    auto savepath = savedir / (p.stem().string() + ".txt");
+    std::cout << "Saves are written to: " << savepath << std::endl;
+    std::shared_ptr<backend::labeller> labeller(new backend::labeller(savepath));
+    frontend::ui ui(playback, labeller);
+    ui.begin();
+    return 0;
+}
diff --git a/src/labeller.cc b/src/labeller.cc
new file mode 100644
index 0000000..64da2d7
--- /dev/null
+++ b/src/labeller.cc
@@ -0,0 +1,116 @@
+#include "labeller.h"
+#include "settings.h"
+#include <fstream>
+
+
+namespace backend {
+    struct labeller_impl {
+        std::vector<std::string> labels;
+        std::vector<label> annotations;
+        // Bool in undo/redo buffers is whether was applied or deleted
+        std::vector<std::pair<bool, label>> undoBuffer;
+        std::vector<std::pair<bool, label>> redoBuffer;
+        std::filesystem::path savepath;
+    };
+
+    labeller::labeller(const std::filesystem::path& savepath) {
+        data = std::shared_ptr<labeller_impl>(new labeller_impl);
+        data->labels = settings::getLabels();
+        data->savepath = savepath;
+    }
+
+    std::pair<label, label> labeller::getSurrounding(double time) const {
+        std::pair<label, label> labs;
+        if(!data->annotations.empty()) {
+            auto above = data->annotations.rbegin();
+            auto below = above;
+            auto rit = above;
+            while(++rit != data->annotations.rend()) {
+                if(rit->time > time && (above->time <= time || rit->time < above->time)) above = rit;
+                if(rit->time <= time && (below->time > time || rit->time > below->time)) below = rit;
+            }
+            if(below->time <= time) labs.first = *below;
+            if(above->time > time) labs.second = *above;
+        }
+        return labs;
+    }
+
+    std::vector<std::string> labeller::getLabels() const {
+        return data->labels;
+    }
+
+    void appLab(std::vector<label>& annotations, label lab) {
+        annotations.push_back(lab);
+    }
+
+    void labeller::applyLabel(std::string name, double time) {
+        label lab(name, time);
+        appLab(data->annotations, lab);
+        data->undoBuffer.push_back({true, lab});
+        data->redoBuffer.clear();
+    }
+
+    double abs(double x) {
+        if(x < 0) x *= -1;
+        return x;
+    }
+
+    label delLab(std::vector<label>& annotations, double time) {
+        if(annotations.empty()) {
+            return label();
+        }
+        auto closest = annotations.rbegin();
+        auto rit = closest;
+        while(++rit != annotations.rend()) {
+            if(abs(rit->time - time) < abs(closest->time - time)) {
+                closest = rit;
+            }
+        }
+        label c(*closest);
+        annotations.erase(closest.base());
+        return c;
+    }
+
+    void labeller::deleteLabel(double time) {
+        label deleted = delLab(data->annotations, time);
+        data->undoBuffer.push_back({false, deleted});
+        data->redoBuffer.clear();
+    }
+
+    // Pops action from 'from', applies, and appends to 'to'.
+    void handleUndoRedo(std::vector<std::pair<bool, label>>& from, std::vector<std::pair<bool, label>>& to, std::vector<label>& annotations) {
+        if(! from.empty()) {
+            auto elem = from.back();
+            from.pop_back();
+            if(elem.first) { // It was applied, so we must unapply
+                delLab(annotations, elem.second.time);
+            } else {
+                appLab(annotations, elem.second);
+            }
+            to.push_back({! elem.first, elem.second});
+        }
+    }
+
+    void labeller::undo() {
+        handleUndoRedo(data->undoBuffer, data->redoBuffer, data->annotations);
+    }
+
+    void labeller::redo() {
+        handleUndoRedo(data->redoBuffer, data->undoBuffer, data->annotations);
+    }
+
+    bool compareLabels(label l1, label l2) {
+        return (l1.time < l2.time);
+    }
+
+    void labeller::save() const {
+        std::ofstream out;
+        out.open(data->savepath);
+        std::vector<label> a(data->annotations);
+        std::sort(a.begin(), a.end(), compareLabels);
+        for(label l : a) {
+            out << l.name << "," << l.time << std::endl;
+        }
+        out.close();
+    }
+}
diff --git a/src/labeller.h b/src/labeller.h
new file mode 100644
index 0000000..4248e2a
--- /dev/null
+++ b/src/labeller.h
@@ -0,0 +1,32 @@
+#pragma once
+#include <memory>
+#include <vector>
+#include <string>
+#include <utility>
+#include <filesystem>
+
+namespace backend {
+    struct labeller_impl;
+
+    struct label {
+        label() {}
+        label(const std::string& name, double time) : name(name), time(time) {}
+        std::string name;
+        double time;
+    };
+
+    class labeller {
+        public:
+            labeller(const std::filesystem::path& savepath);
+            std::pair<label, label> getSurrounding(double time) const;
+            std::vector<std::string> getLabels() const;
+            void applyLabel(std::string name, double time);
+            void deleteLabel(double time); // Deletes closest to time, last added if ties
+            void undo();
+            void redo();
+            void save() const;
+
+        private:
+            std::shared_ptr<labeller_impl> data;
+    };
+}
diff --git a/src/playback.cc b/src/playback.cc
new file mode 100644
index 0000000..045abb6
--- /dev/null
+++ b/src/playback.cc
@@ -0,0 +1,69 @@
+#include "playback.h"
+#include <opencv2/opencv.hpp>
+#include <vector>
+#include <stdexcept>
+#include <iostream>
+#include <thread>
+
+namespace frontend {
+    struct playback_impl {
+        std::vector<cv::Mat> images;
+        int fps;
+        double reportedDuration;
+        std::size_t frameNum = 0;
+    };
+
+    playback::playback(const std::filesystem::path& video) {
+        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) {
+            cv::Mat frame;
+            cap >> frame;
+            if(frame.empty()) break;
+            data->images.push_back(frame);
+        }
+        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 {
+        cv::imshow(windowName, data->images[data->frameNum]);
+    }
+
+    bool playback::seekFrame(std::size_t frameNum) {
+        if(frameNum >= data->images.size()) {
+            return false;
+        }
+        data->frameNum = frameNum;
+        return true;
+    }
+
+    std::size_t playback::getFrame() const {
+        return data->frameNum;
+    }
+
+    bool playback::seekTime(double time) {
+        return seekFrame(std::size_t(time * data->fps));
+    }
+
+    double playback::getTime() const {
+        return getFrame() / (double) data->fps;
+    }
+    
+    void playback::interFrameSleep() const {
+        std::this_thread::sleep_for(std::chrono::milliseconds(1000/data->fps));
+    }
+
+    std::size_t playback::getMaxFrame() const {
+        return data->images.size() - 1;
+    }
+
+    double playback::getMaxTime() const {
+        return getMaxFrame() / (double) data->fps;
+    }
+}
diff --git a/src/playback.h b/src/playback.h
new file mode 100644
index 0000000..f0f1f09
--- /dev/null
+++ b/src/playback.h
@@ -0,0 +1,23 @@
+#pragma once
+#include <filesystem>
+#include <memory>
+#include <string>
+
+namespace frontend {
+    struct playback_impl;
+    class playback {
+        public:
+            playback(const std::filesystem::path& video);
+            void display(const std::string& windowName) const;
+            bool seekFrame(std::size_t frameNum);
+            std::size_t getFrame() const;
+            bool seekTime(double time);
+            double getTime() const;
+            void interFrameSleep() const;
+            std::size_t getMaxFrame() const;
+            double getMaxTime() const;
+
+        private:
+            std::shared_ptr<playback_impl> data;
+    };
+}
diff --git a/src/settings.cc b/src/settings.cc
new file mode 100644
index 0000000..d19acde
--- /dev/null
+++ b/src/settings.cc
@@ -0,0 +1,28 @@
+#include "settings.h"
+#include <confuse.h>
+#include <errno.h>
+#include <cstring>
+#include <stdexcept>
+
+namespace settings {
+    std::vector<std::string> getLabels() {
+        cfg_opt_t opts[] = {
+            CFG_STR_LIST("labels", NULL, CFGF_NONE),
+            CFG_END()
+        };
+        cfg_t *cfg = cfg_init(opts, CFGF_IGNORE_UNKNOWN);
+        if(cfg_parse(cfg, "annotator.conf") == CFG_PARSE_ERROR) {
+            throw std::runtime_error("Configuration file annotator.conf could not be read: " + std::string(strerror(errno)));
+        }
+        try {
+            std::vector<std::string> ret;
+            std::size_t i;
+            for(i = 0; i < cfg_size(cfg, "labels"); i++) {
+                ret.push_back(cfg_getnstr(cfg, "labels", i));
+            }
+            return ret;
+        } catch(std::exception& e) {
+            throw std::runtime_error("Cannot find 'labels' in configuration file");
+        }
+    }
+}
diff --git a/src/settings.h b/src/settings.h
new file mode 100644
index 0000000..3b3ad71
--- /dev/null
+++ b/src/settings.h
@@ -0,0 +1,7 @@
+#pragma once
+#include <string>
+#include <vector>
+
+namespace settings {
+    std::vector<std::string> getLabels(void);
+}
diff --git a/src/ui.cc b/src/ui.cc
new file mode 100644
index 0000000..3d0e193
--- /dev/null
+++ b/src/ui.cc
@@ -0,0 +1,115 @@
+#include "ui.h"
+#include "playback.h"
+#include "labeller.h"
+#include <opencv2/opencv.hpp>
+#include <iostream>
+#include <sys/ioctl.h>
+#include <unistd.h>
+
+namespace frontend {
+    struct ui_impl {
+        std::shared_ptr<playback> pb;
+        std::shared_ptr<backend::labeller> llr;
+        bool play = true;
+        bool stalePrintout = true;
+    };
+
+    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;
+    }
+
+    bool handleKey(int keycode, const std::shared_ptr<ui_impl>& data) {
+        if(keycode != -1) {
+            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());
+                }
+            } else {
+                switch(keycode) {
+                    case 27: // Escape quits
+                        std::cout << std::endl;
+                        return true;
+                    case 32: // Space pauses
+                        data->play = ! data->play;
+                        break;
+                    case 65361: // Left seeks backward 1 frame
+                        data->pb->seekFrame(data->pb->getFrame() - 1);
+                        break;
+                    case 65362: // Up seeks backward 1 second
+                        if(data->pb->getTime() < 1) {
+                            data->pb->seekTime(data->pb->getTime() * -1);
+                        } else {
+                            data->pb->seekTime(data->pb->getTime() - 1);
+                        }
+                        break;
+                    case 65363: // Right seeks forward 1 frame
+                        data->pb->seekFrame(data->pb->getFrame() + 1);
+                        break;
+                    case 65364: // Down seeks forward 1 second
+                        data->pb->seekTime(data->pb->getTime() + 1);
+                        break;
+                    case 117: // u undoes
+                        data->llr->undo();
+                        break;
+                    case 114: // r redoes
+                        data->llr->redo();
+                        break;
+                    case 65535: // DEL deletes
+                        data->llr->deleteLabel(data->pb->getTime());
+                        break;
+                    case 115: // s saves
+                        data->llr->save();
+                        break;
+                    default:
+                        std::cout << "Pressed the " << keycode << " key" << std::endl;
+                        break;
+                }
+            }
+        }
+        return false;
+    }
+
+    void ui::begin() {
+        std::cout << "Playing a video that's " << data->pb->getMaxFrame() << " frames (" << data->pb->getMaxTime() << " seconds) long." << std::endl;
+        std::cout << "Annotations:" << std::endl;
+        int num = 0;
+        for(auto ann : data->llr->getLabels()) {
+            std::cout << num++ << ": " << ann << std::endl;
+        }
+        // Get window size
+        struct winsize size;
+        ioctl(STDOUT_FILENO, TIOCGWINSZ, &size);
+        int cols = size.ws_col;
+        while(true) {
+            data->pb->display("Video");
+            data->pb->interFrameSleep();
+            if(handleKey(cv::pollKey(), data)) {
+                break;
+            }
+            if(data->play) {
+                data->pb->seekFrame(data->pb->getFrame() + 1);
+            }
+            if(data->play || data->stalePrintout) {
+                std::stringstream toPrint;
+                toPrint << "Frame: " << data->pb->getFrame() << " (" << std::setprecision(4) << std::setw(6) << std::left << data->pb->getTime() << " s)";
+                std::string timeString = toPrint.str();
+                toPrint.str(""); // clear contents
+                // Get surrounding labels
+                auto labs = data->llr->getSurrounding(data->pb->getTime());
+                if(! labs.first.name.empty()) toPrint << labs.first.name << " (t=" << std::setprecision(4) << labs.first.time << ")";
+                else toPrint << "START (t=0)";
+                toPrint << " CURRENT ";
+                if(! labs.second.name.empty()) toPrint << labs.second.name << " (t=" << std::setprecision(4) << labs.second.time << ")";
+                else toPrint << "END (t=" << data->pb->getMaxTime() << ")";
+                toPrint << std::string(cols - toPrint.str().size() - timeString.size() - 1, ' ') << timeString;
+                std::cout << "\r\b\r" << toPrint.str();
+                std::cout.flush();
+                data->stalePrintout = false;
+            }
+        }
+    }
+}
diff --git a/src/ui.h b/src/ui.h
new file mode 100644
index 0000000..b4ce859
--- /dev/null
+++ b/src/ui.h
@@ -0,0 +1,19 @@
+#pragma once
+#include <memory>
+
+namespace backend {
+    class labeller;
+}
+
+namespace frontend {
+    struct ui_impl;
+    class playback;
+    class ui {
+        public:
+            ui(std::shared_ptr<playback> playback, std::shared_ptr<backend::labeller> labeller);
+            void begin();
+
+        private:
+            std::shared_ptr<ui_impl> data;
+    };
+}