From 63937172fd1920433b238e38c5d5dac9ae8b6b21 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20=C4=8Cuhel?=
 <79118988+HonzaCuhel@users.noreply.github.com>
Date: Thu, 27 Jun 2024 10:39:06 +0200
Subject: [PATCH] Feat/upload nn archive (#79)

* Update the uplod code

* Add upload logging
---
 pyproject.toml                | 2 +-
 tools/main.py                 | 3 ++-
 tools/modules/exporter.py     | 2 ++
 tools/yolo/yolov8_exporter.py | 2 ++
 4 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index e401ae3..1b5c758 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "tools"
-version = "0.0.1"
+version = "0.0.2"
 description = "Converter for YOLO models into .ONNX format."
 readme = "README.md"
 requires-python = ">=3.8"
diff --git a/tools/main.py b/tools/main.py
index eb83c74..95f8dc3 100644
--- a/tools/main.py
+++ b/tools/main.py
@@ -137,7 +137,8 @@ def convert(
     
     # Upload to remote
     if config.output_remote_url:
-        upload_file_to_remote(exporter.f_onnx, config.output_remote_url, config.put_file_plugin)
+        upload_file_to_remote(exporter.f_nn_archive, config.output_remote_url, config.put_file_plugin)
+        logger.info(f"Uploaded NN archive to {config.output_remote_url}")
 
 
 if __name__ == "__main__":
diff --git a/tools/modules/exporter.py b/tools/modules/exporter.py
index e162d0c..0eecb9e 100644
--- a/tools/modules/exporter.py
+++ b/tools/modules/exporter.py
@@ -42,6 +42,7 @@ def __init__(
         self.model_name = os.path.basename(self.model_path).split(".")[0]
         # Set up file paths
         self.f_onnx = None
+        self.f_nn_archive = None
         self.use_rvc2 = use_rvc2
         self.number_of_channels = None
         self.subtype = subtype
@@ -103,6 +104,7 @@ def make_nn_archive(
             conf_threshold (float): Confidence threshold
             max_det (int): Maximum number of detections
         """
+        self.f_nn_archive = (self.output_folder / f"{self.model_name}.tar.xz").resolve()
         archive = ArchiveGenerator(
             archive_name=self.model_name,
             save_path=str(self.output_folder),
diff --git a/tools/yolo/yolov8_exporter.py b/tools/yolo/yolov8_exporter.py
index 5e3892c..c7e398d 100644
--- a/tools/yolo/yolov8_exporter.py
+++ b/tools/yolo/yolov8_exporter.py
@@ -148,6 +148,8 @@ def export_nn_archive(self, class_names: Optional[List[str]] = None):
             assert len(class_names) == len(names), f"Number of the given class names {len(class_names)} does not match number of classes {len(names)} provided in the model!"
             names = class_names
         
+        self.f_nn_archive = (self.output_folder / f"{self.model_name}.tar.xz").resolve()
+
         if self.mode == DETECT_MODE:
             self.make_nn_archive(names, self.model.model[-1].nc)
         elif self.mode == SEGMENT_MODE: