Minor refactors for PR#8

PolusAI · Jun 26, 2024 · 407e898 · 407e898
1 parent 99ed75f
commit 407e898
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 30 deletions.
diff --git a/src/cpp/utilities/utilities.cpp b/src/cpp/utilities/utilities.cpp
@@ -1,11 +1,15 @@
 #include <iomanip>
 #include <ctime>
-#include <chrono>
 #include "utilities.h"
 #include <cassert>
 #include <tiffio.h>
 #include <thread>
 
+#include "tensorstore/driver/zarr/dtype.h"
+
+
+using ::tensorstore::internal_zarr::ChooseBaseDType;
+
 namespace bfiocpp {
 tensorstore::Spec GetOmeTiffSpecToRead(const std::string& filename){
     return tensorstore::Spec::FromJson({{"driver", "ometiff"},
@@ -46,6 +50,45 @@ uint16_t GetDataTypeCode (std::string_view type_name){
   else {return 2;}
 }
 
+std::string GetEncodedType(uint16_t data_type_code){
+    switch (data_type_code)
+    {
+    case 1:
+        return ChooseBaseDType(tensorstore::dtype_v<std::uint8_t>).value().encoded_dtype;
+        break;
+    case 2:
+        return ChooseBaseDType(tensorstore::dtype_v<std::uint16_t>).value().encoded_dtype;
+        break;
+    case 4:
+        return ChooseBaseDType(tensorstore::dtype_v<std::uint32_t>).value().encoded_dtype;
+        break;
+    case 8:
+        return ChooseBaseDType(tensorstore::dtype_v<std::uint16_t>).value().encoded_dtype;
+        break;
+    case 16:
+        return ChooseBaseDType(tensorstore::dtype_v<std::int8_t>).value().encoded_dtype;
+        break;
+    case 32:
+        return ChooseBaseDType(tensorstore::dtype_v<std::int16_t>).value().encoded_dtype;
+        break;
+    case 64:
+        return ChooseBaseDType(tensorstore::dtype_v<std::int32_t>).value().encoded_dtype;
+        break;
+    case 128:
+        return ChooseBaseDType(tensorstore::dtype_v<std::int64_t>).value().encoded_dtype;
+        break;
+    case 256:
+        return ChooseBaseDType(tensorstore::dtype_v<float>).value().encoded_dtype;
+        break;
+    case 512:
+        return ChooseBaseDType(tensorstore::dtype_v<double>).value().encoded_dtype;
+        break;
+    default:
+        return ChooseBaseDType(tensorstore::dtype_v<std::uint16_t>).value().encoded_dtype;
+        break;
+    }
+}
+
 std::string GetUTCString() {
     // Get the current UTC time
     auto now = std::chrono::system_clock::now();

diff --git a/src/cpp/utilities/utilities.h b/src/cpp/utilities/utilities.h
@@ -15,6 +15,7 @@ tensorstore::Spec GetOmeTiffSpecToRead(const std::string& filename);
 tensorstore::Spec GetZarrSpecToRead(const std::string& filename);
 
 uint16_t GetDataTypeCode (std::string_view type_name);
+std::string GetEncodedType(uint16_t data_type_code);
 std::string GetUTCString();
 std::string GetOmeXml(const std::string& file_path);
 std::tuple<std::optional<int>, std::optional<int>, std::optional<int>>ParseMultiscaleMetadata(const std::string& axes_list, int len);

diff --git a/src/cpp/writer/tswriter.cpp b/src/cpp/writer/tswriter.cpp
@@ -1,11 +1,11 @@
-#include "tswriter.h"
+#include <string>
 
-#include "../utilities/utilities.h"
+#include "tensorstore/array.h"
+#include "tensorstore/open.h"
 
-#include <variant>
-#include <string>
+#include "tswriter.h"
+#include "../utilities/utilities.h"
 
-using ::tensorstore::internal_zarr::ChooseBaseDType;
 
 namespace bfiocpp {
 
@@ -14,14 +14,13 @@ TsWriterCPP::TsWriterCPP(
     const std::vector<std::int64_t>& image_shape, 
     const std::vector<std::int64_t>& chunk_shape,
     const std::string& dtype_str
-  ): _filename(fname), _image_shape(image_shape), _chunk_shape(chunk_shape) {
-
-  _dtype_code = GetDataTypeCode(dtype_str); 
-
-  auto spec = GetZarrSpecToWrite(_filename, image_shape, chunk_shape, dtype_str);
-
+  ): _filename(fname), 
+     _image_shape(image_shape), 
+     _chunk_shape(chunk_shape),  
+     _dtype_code(GetDataTypeCode(dtype_str)) {
+
   TENSORSTORE_CHECK_OK_AND_ASSIGN(_source, tensorstore::Open(
-                            spec,
+                            GetZarrSpecToWrite(_filename, _image_shape, _chunk_shape, GetEncodedType(_dtype_code)),
                             tensorstore::OpenMode::create |
                             tensorstore::OpenMode::delete_existing,
                             tensorstore::ReadWriteMode::write).result());

diff --git a/src/cpp/writer/tswriter.h b/src/cpp/writer/tswriter.h
@@ -1,25 +1,9 @@
 #pragma once
 
 #include <string>
-#include <memory>
 #include <vector>
-#include <variant>
-#include <iostream>
-#include <tuple>
-#include <optional>
-#include <unordered_map>
-#include "../reader/sequence.h"
-
 #include "tensorstore/tensorstore.h"
-#include "tensorstore/context.h"
-#include "tensorstore/array.h"
-#include "tensorstore/driver/zarr/dtype.h"
-#include "tensorstore/index_space/dim_expression.h"
-#include "tensorstore/kvstore/kvstore.h"
-#include "tensorstore/open.h"
-
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
+#include "../reader/sequence.h"
 #include <pybind11/numpy.h>
 
 namespace py = pybind11;