add json, none and pretty download prints

Author: lnotspotl

bindings/python/src/modelzoo/ZooBindings.cpp

@@ -32,13 +32,15 @@ void ZooBindings::bind(pybind11::module& m, void* pCallstack) {
           py::arg("useCached") = true,
           py::arg("cacheDirectory") = "",
           py::arg("apiKey") = "",
+          py::arg("progressFormat") = "none",
           DOC(dai, getModelFromZoo));
 
     m.def("downloadModelsFromZoo",
           downloadModelsFromZoo,
           py::arg("path"),
           py::arg("cacheDirectory") = "",
           py::arg("apiKey") = "",
+          py::arg("progressFormat") = "none",
           DOC(dai, downloadModelsFromZoo));
 
     // Bind NNModelDescription

include/depthai/modelzoo/Zoo.hpp

@@ -86,12 +86,14 @@ struct NNModelDescription {
  * DEPTHAI_ZOO_CACHE_PATH environment variable and uses that if set, otherwise the default value is used (see getDefaultCachePath).
  * @param apiKey: API key for the model zoo, default is "". If apiKey is set to "", this function checks the DEPTHAI_ZOO_API_KEY environment variable and uses
  * that if set. Otherwise, no API key is used.
+ * @param progressFormat: Format to use for progress output (possible values: pretty, json, none), default is "none"
  * @return std::string: Path to the model in cache
  */
 std::string getModelFromZoo(const NNModelDescription& modelDescription,
                             bool useCached = true,
                             const std::string& cacheDirectory = "",
-                            const std::string& apiKey = "");
+                            const std::string& apiKey = "",
+                            const std::string& progressFormat = "none");
 
 /**
  * @brief Helper function allowing one to download all models specified in yaml files in the given path and store them in the cache directory
@@ -101,10 +103,10 @@ std::string getModelFromZoo(const NNModelDescription& modelDescription,
  * DEPTHAI_ZOO_CACHE_PATH environment variable and uses that if set, otherwise the default is used (see getDefaultCachePath).
  * @param apiKey: API key for the model zoo, default is "". If apiKey is set to "", this function checks the DEPTHAI_ZOO_API_KEY environment variable and uses
  * that if set. Otherwise, no API key is used.
- * @param format: Format to use for output (possible values: pretty, json), default is "pretty"
+ * @param progressFormat: Format to use for progress output (possible values: pretty, json, none), default is "none"
  * @return bool: True if all models were downloaded successfully, false otherwise
  */
-bool downloadModelsFromZoo(const std::string& path, const std::string& cacheDirectory = "", const std::string& apiKey = "", const std::string& format = "pretty");
+bool downloadModelsFromZoo(const std::string& path, const std::string& cacheDirectory = "", const std::string& apiKey = "", const std::string& progressFormat = "none");
 
 std::ostream& operator<<(std::ostream& os, const NNModelDescription& modelDescription);
 

src/modelzoo/Zoo.cpp

@@ -4,8 +4,6 @@
 #include <filesystem>
 #include <iostream>
 #include <memory>
-#include <thread>
-#include <atomic>
 #include <mutex>
 #include <nlohmann/json.hpp>
 #include <nlohmann/json_fwd.hpp>
@@ -332,8 +330,6 @@ nlohmann::json ZooManager::fetchModelDownloadLinks() {
         throw std::runtime_error(generateErrorMessageHub(response));
     }
 
-    std::cout << "Response: " << response.text << std::endl;
-
     // Extract download links from response
     nlohmann::json responseJson = nlohmann::json::parse(response.text);
     return responseJson;
@@ -406,11 +402,145 @@ std::string ZooManager::loadModelFromCache() const {
     return std::filesystem::absolute(folderFiles[0]).string();
 }
 
+class CprCallback {
+   public:
+    virtual ~CprCallback() = default;
+    CprCallback(const std::string& modelName) : modelName(modelName) {}
+
+    virtual void cprCallback(
+        cpr::cpr_off_t downloadTotal, cpr::cpr_off_t downloadNow, cpr::cpr_off_t uploadTotal, cpr::cpr_off_t uploadNow, intptr_t userdata) = 0;
+
+    virtual std::unique_ptr<cpr::ProgressCallback> getCprProgressCallback() {
+        return std::make_unique<cpr::ProgressCallback>(
+            [this](cpr::cpr_off_t downloadTotal, cpr::cpr_off_t downloadNow, cpr::cpr_off_t uploadTotal, cpr::cpr_off_t uploadNow, intptr_t userdata) {
+                this->cprCallback(downloadTotal, downloadNow, uploadTotal, uploadNow, userdata);
+                return true;
+            });
+    }
+
+   protected:
+    std::string modelName;
+};
+
+class JsonCprCallback : public CprCallback {
+    constexpr static long long PRINT_INTERVAL_MS = 100;
+
+   public:
+    JsonCprCallback(const std::string& modelName) : CprCallback(modelName) {
+        startTime = std::chrono::steady_clock::time_point::min();
+    }
+
+    void print(long downloadTotal, long downloadNow, const std::string& modelName) {
+        nlohmann::json json = {
+            {"download_total", downloadTotal},
+            {"download_now", downloadNow},
+            {"model_name", modelName},
+        };
+        std::cout << json.dump() << std::endl;
+    }
+
+    void cprCallback(
+        cpr::cpr_off_t downloadTotal, cpr::cpr_off_t downloadNow, cpr::cpr_off_t uploadTotal, cpr::cpr_off_t uploadNow, intptr_t userdata) override {
+        (void)uploadTotal;
+        (void)uploadNow;
+        (void)userdata;
+
+        bool firstCall = startTime == std::chrono::steady_clock::time_point::min();
+        if(firstCall || downloadTotal == 0) {
+            startTime = std::chrono::steady_clock::now();
+        }
+
+        bool shouldPrint = std::chrono::steady_clock::now() - startTime > std::chrono::milliseconds(PRINT_INTERVAL_MS) || this->downloadTotal != downloadTotal;
+
+        if(shouldPrint) {
+            print(downloadTotal, downloadNow, modelName);
+            startTime = std::chrono::steady_clock::now();
+        }
+
+        this->downloadTotal = downloadTotal;
+        this->downloadNow = downloadNow;
+    }
+
+    ~JsonCprCallback() override {
+        if(downloadTotal != 0) {
+            print(downloadTotal, downloadNow, modelName);
+        }
+    }
+
+   private:
+    long downloadTotal = 0;
+    long downloadNow = 0;
+    std::chrono::steady_clock::time_point startTime;
+};
+
+class PrettyCprCallback : public CprCallback {
+   public:
+    PrettyCprCallback(const std::string& modelName) : CprCallback(modelName), finalProgressPrinted(false) {}
+
+    void cprCallback(
+        cpr::cpr_off_t downloadTotal, cpr::cpr_off_t downloadNow, cpr::cpr_off_t uploadTotal, cpr::cpr_off_t uploadNow, intptr_t userdata) override {
+        (void)uploadTotal;
+        (void)uploadNow;
+        (void)userdata;
+
+        if(finalProgressPrinted) return;
+
+        if(downloadTotal > 0) {
+            float progress = static_cast<float>(downloadNow) / downloadTotal;
+            int barWidth = 50;
+            int pos = static_cast<int>(barWidth * progress);
+
+            std::cout << "\rDownloading " << modelName << " [";
+            for(int i = 0; i < barWidth; ++i) {
+                if(i < pos)
+                    std::cout << "=";
+                else if(i == pos)
+                    std::cout << ">";
+                else
+                    std::cout << " ";
+            }
+            std::cout << "] " << std::fixed << std::setprecision(3) << progress * 100.0f << "% " << downloadNow / 1024.0f / 1024.0f << "/"
+                      << downloadTotal / 1024.0f / 1024.0f << " MB";
+
+            if(downloadNow == downloadTotal) {
+                std::cout << std::endl;
+                finalProgressPrinted = true;
+            } else {
+                std::cout << "\r";
+                std::cout.flush();
+            }
+        }
+    }
+
+   private:
+    bool finalProgressPrinted;
+};
+
+class NoneCprCallback : public CprCallback {
+   public:
+    NoneCprCallback(const std::string& modelName) : CprCallback(modelName) {}
+
+    void cprCallback(cpr::cpr_off_t, cpr::cpr_off_t, cpr::cpr_off_t, cpr::cpr_off_t, intptr_t) override {
+        // Do nothing
+    }
+};
+
+std::unique_ptr<CprCallback> getCprCallback(const std::string& format, const std::string& name) {
+    if(format == "json") {
+        return std::make_unique<JsonCprCallback>(name);
+    } else if(format == "pretty") {
+        return std::make_unique<PrettyCprCallback>(name);
+    } else if(format == "none") {
+        return std::make_unique<NoneCprCallback>(name);
+    }
+    throw std::runtime_error("Invalid format: " + format);
+}
+
 std::string getModelFromZoo(const NNModelDescription& modelDescription,
                             bool useCached,
                             const std::string& cacheDirectory,
                             const std::string& apiKey,
-                            const std::function<void(cpr::cpr_off_t, cpr::cpr_off_t, cpr::cpr_off_t, cpr::cpr_off_t, intptr_t)>& progressCallback) {
+                            const std::string& progressFormat) {
     // Check if model description is valid
     if(!modelDescription.check()) throw std::runtime_error("Invalid model description:\n" + modelDescription.toString());
 
@@ -482,20 +612,11 @@ std::string getModelFromZoo(const NNModelDescription& modelDescription,
     zooManager.createCacheFolder();
 
     // Create download progress callback
-    std::unique_ptr<cpr::ProgressCallback> cprCallback;
-    if(progressCallback) {
-        cprCallback = std::make_unique<cpr::ProgressCallback>(
-            [&](cpr::cpr_off_t downloadTotal, cpr::cpr_off_t downloadNow, cpr::cpr_off_t uploadTotal, cpr::cpr_off_t uploadNow, intptr_t userdata) {
-                progressCallback(downloadTotal, downloadNow, uploadTotal, uploadNow, userdata);
-                return true;
-            });
-    } else {
-        cprCallback = std::make_unique<cpr::ProgressCallback>([](cpr::cpr_off_t, cpr::cpr_off_t, cpr::cpr_off_t, cpr::cpr_off_t, intptr_t) { return true; });
-    }
+    std::unique_ptr<CprCallback> cprCallback = getCprCallback(progressFormat, modelDescription.globalMetadataEntryName.size() > 0 ? modelDescription.globalMetadataEntryName : modelDescription.model);
 
     // Download model
     logger::info("Downloading model from model zoo");
-    zooManager.downloadModel(responseJson, std::move(cprCallback));
+    zooManager.downloadModel(responseJson, cprCallback->getCprProgressCallback());
 
     // Store model as yaml in the cache folder
     std::string yamlPath = combinePaths(zooManager.getModelCacheFolderPath(cacheDirectory), "model.yaml");
@@ -506,85 +627,7 @@ std::string getModelFromZoo(const NNModelDescription& modelDescription,
     return modelPath;
 }
 
-std::string getModelFromZoo(const NNModelDescription& modelDescription, bool useCached, const std::string& cacheDirectory, const std::string& apiKey) {
-    return getModelFromZoo(modelDescription, useCached, cacheDirectory, apiKey, nullptr);
-}
-
-struct JsonDownloadProgressManager {
-    JsonDownloadProgressManager(size_t updateIntervalMs = 1000) : updateIntervalMs(updateIntervalMs) {
-        pause();
-    }
-
-    size_t updateIntervalMs;
-
-    std::string model;
-    size_t bytesDownloaded = 0;
-    size_t bytesTotal = 0;
-    std::mutex mutex;
-
-    std::thread thread;
-    std::atomic<bool> running;
-    std::atomic<bool> started;
-    std::atomic<bool> firstUpdate;
-
-    bool pause() {
-        std::lock_guard<std::mutex> lock(mutex);
-        started = false;
-        firstUpdate = true;
-        return true;
-    }
-
-    bool resume() {
-        std::lock_guard<std::mutex> lock(mutex);
-        started = true;
-        firstUpdate = true;
-        return true;
-    }
-
-    void update(cpr::cpr_off_t downloadTotal, cpr::cpr_off_t downloadNow, cpr::cpr_off_t uploadTotal, cpr::cpr_off_t uploadNow, intptr_t userdata) {
-        (void)uploadTotal;
-        (void)uploadNow;
-        (void)userdata;
-        std::lock_guard<std::mutex> lock(mutex);
-        this->bytesDownloaded = downloadNow;
-        this->bytesTotal = downloadTotal;
-        firstUpdate = false;
-    }
-
-    nlohmann::json getJson() {
-        std::lock_guard<std::mutex> lock(mutex);
-        nlohmann::json json;
-        json["model"] = model;
-        json["bytes_downloaded"] = bytesDownloaded;
-        json["bytes_total"] = bytesTotal;
-        return json;
-    }
-
-    void setModel(const std::string& model) {
-        std::lock_guard<std::mutex> lock(mutex);
-        this->model = model;
-    }
-
-    void startThread() {
-        running = true;
-        thread = std::thread([this]() {
-            while(running) {
-                std::this_thread::sleep_for(std::chrono::milliseconds(updateIntervalMs));
-                if(!started || firstUpdate) continue;
-                auto json = getJson();
-                std::string jsonStr = json.dump();
-                std::cout << jsonStr << std::endl;
-            }
-        });
-    }
-
-    void stopThread() {
-        running = false;
-        thread.join();
-    }
-};
-
-bool downloadModelsFromZoo(const std::string& path, const std::string& cacheDirectory, const std::string& apiKey, const std::string& format) {
+bool downloadModelsFromZoo(const std::string& path, const std::string& cacheDirectory, const std::string& apiKey, const std::string& progressFormat) {
     logger::info("Downloading models from zoo");
     // Make sure 'path' exists
     if(!std::filesystem::exists(path)) throw std::runtime_error("Path does not exist: " + path);
@@ -602,9 +645,6 @@ bool downloadModelsFromZoo(const std::string& path, const std::string& cacheDire
         }
     }
 
-    auto jsonDownloadProgressManager = std::make_unique<JsonDownloadProgressManager>(100);
-    jsonDownloadProgressManager->startThread();
-    auto callback = std::bind(&JsonDownloadProgressManager::update, jsonDownloadProgressManager.get(), std::placeholders::_1, std::placeholders::_2, std::placeholders::_3, std::placeholders::_4, std::placeholders::_5);
     // Download models from yaml files
     int numSuccess = 0, numFail = 0;
     for(size_t i = 0; i < models.size(); ++i) {
@@ -613,21 +653,17 @@ bool downloadModelsFromZoo(const std::string& path, const std::string& cacheDire
 
         // Download model - ignore the returned model path here == we are only interested in downloading the model
         try {
-            jsonDownloadProgressManager->setModel(modelName);
-            jsonDownloadProgressManager->resume();
             logger::info("Downloading model [{} / {}]: {}", i + 1, models.size(), modelName);
             auto modelDescription = NNModelDescription::fromYamlFile(modelName, path);
-            getModelFromZoo(modelDescription, true, cacheDirectory, apiKey, callback);
+            getModelFromZoo(modelDescription, true, cacheDirectory, apiKey, progressFormat);
             logger::info("Downloaded model [{} / {}]: {}", i + 1, models.size(), modelName);
             numSuccess++;
         } catch(const std::exception& e) {
             logger::error("Failed to download model [{} / {}]: {} in folder {}\n{}", i + 1, models.size(), modelName, path, e.what());
             numFail++;
         }
-        jsonDownloadProgressManager->pause();
     }
 
-    jsonDownloadProgressManager->stopThread();
     logger::info("Downloaded {} models from folder {} | {} failed.", numSuccess, path, numFail);
     return numFail == 0;
 }
@@ -658,18 +694,24 @@ std::string ZooManager::getGlobalMetadataFilePath() const {
 
 #else
 
-std::string getModelFromZoo(const NNModelDescription& modelDescription, bool useCached, const std::string& cacheDirectory, const std::string& apiKey) {
+std::string getModelFromZoo(const NNModelDescription& modelDescription,
+                            bool useCached,
+                            const std::string& cacheDirectory,
+                            const std::string& apiKey,
+                            const std::string& progressFormat) {
     (void)modelDescription;
     (void)useCached;
     (void)cacheDirectory;
     (void)apiKey;
+    (void)progressFormat;
     throw std::runtime_error("getModelFromZoo requires libcurl to be enabled. Please recompile DepthAI with libcurl enabled.");
 }
 
-bool downloadModelsFromZoo(const std::string& path, const std::string& cacheDirectory, const std::string& apiKey) {
+bool downloadModelsFromZoo(const std::string& path, const std::string& cacheDirectory, const std::string& apiKey, const std::string& progressFormat) {
     (void)path;
     (void)cacheDirectory;
     (void)apiKey;
+    (void)progressFormat;
     throw std::runtime_error("downloadModelsFromZoo requires libcurl to be enabled. Please recompile DepthAI with libcurl enabled.");
 }
 

src/modelzoo/zoo_helper.cpp

@@ -74,7 +74,7 @@ int main(int argc, char* argv[]) {
     }
 
     // Download models
-    bool success = dai::downloadModelsFromZoo(yamlFolder, cacheFolder, apiKey);
+    bool success = dai::downloadModelsFromZoo(yamlFolder, cacheFolder, apiKey, format);
     if(!success) {
         std::cerr << "Failed to download all models from " << yamlFolder << std::endl;
         return EXIT_FAILURE;