examples/extract_slice: translate input spec to 0.

IWYU driver/downsample Cleanups while I was attempting to use a simple invocation of the extract_slice example to repro #140 PiperOrigin-RevId: 613040988 Change-Id: I53a6653bb9a777d5dcf9a3101ab64218dd34ad94
google · Mar 6, 2024 · c9cd01c · c9cd01c
1 parent e5c13e1
commit c9cd01c
Show file tree

Hide file tree

Showing 8 changed files with 169 additions and 13 deletions.
diff --git a/examples/BUILD b/examples/BUILD
@@ -98,6 +98,7 @@ tensorstore_cc_binary(
         "//tensorstore:open",
         "//tensorstore:open_mode",
         "//tensorstore:spec",
+        "//tensorstore:strided_layout",
         "//tensorstore/index_space:dim_expression",
         "//tensorstore/index_space:index_transform",
         "//tensorstore/internal/image",
@@ -106,6 +107,7 @@ tensorstore_cc_binary(
         "//tensorstore/internal/image:png",
         "//tensorstore/internal/image:webp",
         "//tensorstore/util:json_absl_flag",
+        "//tensorstore/util:result",
         "//tensorstore/util:span",
         "//tensorstore/util:status",
         "//tensorstore/util:str_cat",
@@ -116,5 +118,6 @@ tensorstore_cc_binary(
         "@com_google_absl//absl/strings",
         "@com_google_riegeli//riegeli/bytes:fd_writer",
         "@com_google_riegeli//riegeli/bytes:std_io",
+        "@com_google_riegeli//riegeli/bytes:writer",
     ],
 )
diff --git a/examples/extract_slice.cc b/examples/extract_slice.cc
@@ -12,15 +12,17 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Extracts a slice of a volumetric dataset, outputtting it as a 2d jpeg image.
+// Extracts a slice of a volumetric dataset, outputting it as a 2d image.
 //
 // extract_slice --output_file=/tmp/foo.jpg --input_spec=...
 
-#include <cstdint>
-#include <fstream>
+#include <stdint.h>
+
 #include <iostream>
+#include <memory>
 #include <string>
 #include <type_traits>
+#include <utility>
 #include <vector>
 
 #include "absl/flags/flag.h"
@@ -30,6 +32,7 @@
 #include <nlohmann/json.hpp>
 #include "riegeli/bytes/fd_writer.h"
 #include "riegeli/bytes/std_io.h"
+#include "riegeli/bytes/writer.h"
 #include "tensorstore/array.h"
 #include "tensorstore/context.h"
 #include "tensorstore/data_type.h"
@@ -46,8 +49,10 @@
 #include "tensorstore/open.h"
 #include "tensorstore/open_mode.h"
 #include "tensorstore/spec.h"
+#include "tensorstore/strided_layout.h"
 #include "tensorstore/tensorstore.h"
 #include "tensorstore/util/json_absl_flag.h"
+#include "tensorstore/util/result.h"
 #include "tensorstore/util/span.h"
 #include "tensorstore/util/status.h"
 #include "tensorstore/util/str_cat.h"
@@ -56,7 +61,6 @@ namespace {
 
 using ::tensorstore::Context;
 using ::tensorstore::Index;
-using ::tensorstore::StrCat;
 using ::tensorstore::internal_image::AvifWriter;
 using ::tensorstore::internal_image::ImageInfo;
 using ::tensorstore::internal_image::ImageWriter;
@@ -115,6 +119,10 @@ absl::Status Run(tensorstore::Spec input_spec, std::string output_filename) {
 
   std::cerr << std::endl << "Before: " << *transform << std::endl;
 
+  // DimRange(...).IndexSlice(0) transform below assumes that extra dimensions
+  // are 0-based; so make all dimensions 0-based here.
+  transform = transform | tensorstore::AllDims().TranslateTo(0);
+
   // By convention, assume that the first dimension is Y, and the second is X,
   // and the third is C. The C++ api could use some help with labelling missing
   // dimensions, actually...
@@ -126,6 +134,7 @@ absl::Status Run(tensorstore::Spec input_spec, std::string output_filename) {
     has_y = has_y || l == "y";
     has_c = has_c || l == "c";
   }
+
   if (has_y) {
     transform = transform | tensorstore::Dims("y").MoveTo(0);
   }
@@ -139,6 +148,8 @@ absl::Status Run(tensorstore::Spec input_spec, std::string output_filename) {
     transform = transform | tensorstore::DimRange(2, -1).IndexSlice(0);
   }
 
+  std::cerr << std::endl << "After: " << *transform << std::endl;
+
   auto constrained_input = input | *transform;
   TENSORSTORE_RETURN_IF_ERROR(constrained_input);
 

diff --git a/tensorstore/driver/downsample/BUILD b/tensorstore/driver/downsample/BUILD
@@ -30,23 +30,52 @@ tensorstore_cc_library(
         ":downsample_nditerable",
         ":downsample_util",
         ":grid_occupancy_map",
+        "//tensorstore:array",
+        "//tensorstore:array_storage_statistics",
+        "//tensorstore:box",
+        "//tensorstore:chunk_layout",
+        "//tensorstore:codec_spec",
+        "//tensorstore:context",
+        "//tensorstore:contiguous_layout",
+        "//tensorstore:data_type",
         "//tensorstore:downsample_method",
+        "//tensorstore:index",
+        "//tensorstore:json_serialization_options",
+        "//tensorstore:open_mode",
+        "//tensorstore:open_options",
+        "//tensorstore:rank",
+        "//tensorstore:resize_options",
+        "//tensorstore:schema",
         "//tensorstore:spec",
+        "//tensorstore:transaction",
         "//tensorstore/driver",
+        "//tensorstore/driver:chunk",
         "//tensorstore/index_space:dim_expression",
+        "//tensorstore/index_space:dimension_units",
         "//tensorstore/index_space:index_transform",
+        "//tensorstore/index_space:transformed_array",
+        "//tensorstore/internal:arena",
         "//tensorstore/internal:intrusive_ptr",
+        "//tensorstore/internal:lock_collection",
         "//tensorstore/internal:nditerable_transformed_array",
         "//tensorstore/internal/json_binding",
+        "//tensorstore/kvstore",
         "//tensorstore/serialization",
         "//tensorstore/util:executor",
+        "//tensorstore/util:future",
+        "//tensorstore/util:iterate",
         "//tensorstore/util:result",
         "//tensorstore/util:span",
+        "//tensorstore/util:status",
+        "//tensorstore/util:str_cat",
+        "//tensorstore/util/execution",
         "//tensorstore/util/execution:any_receiver",
         "//tensorstore/util/execution:sender_util",
         "//tensorstore/util/garbage_collection",
         "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/synchronization",
     ],
     alwayslink = True,
 )
@@ -82,11 +111,12 @@ tensorstore_cc_test(
     deps = [
         ":downsample_array",
         "//tensorstore:array",
+        "//tensorstore:data_type",
+        "//tensorstore:downsample_method",
         "//tensorstore:index",
         "//tensorstore/index_space:dim_expression",
         "//tensorstore/index_space:transformed_array",
         "//tensorstore/util:span",
-        "//tensorstore/util:status_testutil",
         "@com_github_nlohmann_json//:nlohmann_json",
         "@com_google_googletest//:gtest_main",
     ],
@@ -101,11 +131,11 @@ tensorstore_cc_library(
         "//tensorstore:data_type",
         "//tensorstore:downsample_method",
         "//tensorstore:index",
+        "//tensorstore:rank",
         "//tensorstore/internal:arena",
         "//tensorstore/internal:elementwise_function",
         "//tensorstore/internal:nditerable",
         "//tensorstore/internal:nditerable_buffer_management",
-        "//tensorstore/internal:nditerable_util",
         "//tensorstore/internal:unique_with_intrusive_allocator",
         "//tensorstore/util:extents",
         "//tensorstore/util:iterate",
@@ -146,22 +176,38 @@ tensorstore_cc_test(
     srcs = ["downsample_test.cc"],
     deps = [
         ":downsample",
+        "//tensorstore",
+        "//tensorstore:array",
+        "//tensorstore:chunk_layout",
         "//tensorstore:context",
+        "//tensorstore:data_type",
         "//tensorstore:downsample",
+        "//tensorstore:downsample_method",
+        "//tensorstore:index",
         "//tensorstore:open",
+        "//tensorstore:open_mode",
+        "//tensorstore:schema",
         "//tensorstore:spec",
+        "//tensorstore:static_cast",
         "//tensorstore/driver",
         "//tensorstore/driver:driver_testutil",
         "//tensorstore/driver/array",
         "//tensorstore/driver/cast",
         "//tensorstore/driver/n5",
         "//tensorstore/driver/zarr",
         "//tensorstore/index_space:dim_expression",
+        "//tensorstore/index_space:index_transform",
+        "//tensorstore/index_space:transformed_array",
         "//tensorstore/internal:global_initializer",
         "//tensorstore/internal:json_gtest",
         "//tensorstore/kvstore/memory",
         "//tensorstore/util:status_testutil",
+        "//tensorstore/util:str_cat",
+        "//tensorstore/util:unit",
+        "//tensorstore/util/execution",
         "//tensorstore/util/execution:sender_util",
+        "@com_github_nlohmann_json//:nlohmann_json",
+        "@com_google_absl//absl/status",
         "@com_google_googletest//:gtest_main",
     ],
 )
@@ -176,12 +222,20 @@ tensorstore_cc_library(
         "//tensorstore:downsample_method",
         "//tensorstore:index",
         "//tensorstore:index_interval",
+        "//tensorstore:rank",
         "//tensorstore/index_space:index_transform",
+        "//tensorstore/index_space:output_index_method",
+        "//tensorstore/internal:integer_overflow",
+        "//tensorstore/util:byte_strided_pointer",
+        "//tensorstore/util:division",
         "//tensorstore/util:iterate",
         "//tensorstore/util:result",
         "//tensorstore/util:span",
+        "//tensorstore/util:status",
+        "//tensorstore/util:str_cat",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
     ],
 )

diff --git a/tensorstore/driver/downsample/downsample.cc b/tensorstore/driver/downsample/downsample.cc
@@ -14,31 +14,74 @@
 
 #include "tensorstore/driver/downsample/downsample.h"
 
+#include <stddef.h>
+
+#include <algorithm>
+#include <cassert>
 #include <mutex>
+#include <utility>
+#include <vector>
 
 #include "absl/base/thread_annotations.h"
+#include "absl/container/inlined_vector.h"
 #include "absl/status/status.h"
+#include "absl/synchronization/mutex.h"
+#include "tensorstore/array.h"
+#include "tensorstore/array_storage_statistics.h"
+#include "tensorstore/box.h"
+#include "tensorstore/chunk_layout.h"
+#include "tensorstore/codec_spec.h"
+#include "tensorstore/context.h"
+#include "tensorstore/contiguous_layout.h"
+#include "tensorstore/data_type.h"
+#include "tensorstore/downsample_method.h"
+#include "tensorstore/driver/chunk.h"
 #include "tensorstore/driver/downsample/downsample_array.h"
 #include "tensorstore/driver/downsample/downsample_method_json_binder.h"
 #include "tensorstore/driver/downsample/downsample_nditerable.h"
 #include "tensorstore/driver/downsample/downsample_util.h"
 #include "tensorstore/driver/downsample/grid_occupancy_map.h"
 #include "tensorstore/driver/driver.h"
 #include "tensorstore/driver/driver_handle.h"
+#include "tensorstore/driver/driver_spec.h"
 #include "tensorstore/driver/read.h"
 #include "tensorstore/driver/registry.h"
+#include "tensorstore/index.h"
 #include "tensorstore/index_space/dim_expression.h"
+#include "tensorstore/index_space/dimension_units.h"
+#include "tensorstore/index_space/index_domain.h"
 #include "tensorstore/index_space/index_domain_builder.h"
+#include "tensorstore/index_space/index_transform.h"
 #include "tensorstore/index_space/index_transform_builder.h"
+#include "tensorstore/index_space/transformed_array.h"
+#include "tensorstore/internal/arena.h"
 #include "tensorstore/internal/intrusive_ptr.h"
+#include "tensorstore/internal/json_binding/json_binding.h"
 #include "tensorstore/internal/json_binding/std_array.h"
+#include "tensorstore/internal/lock_collection.h"
 #include "tensorstore/internal/nditerable_transformed_array.h"
+#include "tensorstore/json_serialization_options.h"
+#include "tensorstore/kvstore/kvstore.h"
+#include "tensorstore/kvstore/spec.h"
+#include "tensorstore/open_mode.h"
+#include "tensorstore/open_options.h"
+#include "tensorstore/rank.h"
+#include "tensorstore/resize_options.h"
+#include "tensorstore/schema.h"
 #include "tensorstore/serialization/std_vector.h"  // IWYU pragma: keep
 #include "tensorstore/spec.h"
+#include "tensorstore/transaction.h"
 #include "tensorstore/util/execution/any_receiver.h"
+#include "tensorstore/util/execution/execution.h"
 #include "tensorstore/util/execution/sender_util.h"
 #include "tensorstore/util/executor.h"
+#include "tensorstore/util/future.h"
 #include "tensorstore/util/garbage_collection/std_vector.h"  // IWYU pragma: keep
+#include "tensorstore/util/iterate.h"
+#include "tensorstore/util/result.h"
+#include "tensorstore/util/span.h"
+#include "tensorstore/util/status.h"
+#include "tensorstore/util/str_cat.h"
 
 namespace tensorstore {
 namespace internal_downsample {

diff --git a/tensorstore/driver/downsample/downsample_array_test.cc b/tensorstore/driver/downsample/downsample_array_test.cc
@@ -14,15 +14,18 @@
 
 #include "tensorstore/driver/downsample/downsample_array.h"
 
+#include <stdint.h>
+
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <nlohmann/json.hpp>
 #include "tensorstore/array.h"
+#include "tensorstore/data_type.h"
+#include "tensorstore/downsample_method.h"
 #include "tensorstore/index.h"
 #include "tensorstore/index_space/dim_expression.h"
 #include "tensorstore/index_space/transformed_array.h"
 #include "tensorstore/util/span.h"
-#include "tensorstore/util/status_testutil.h"
 
 namespace {
 

diff --git a/tensorstore/driver/downsample/downsample_nditerable.cc b/tensorstore/driver/downsample/downsample_nditerable.cc
@@ -14,11 +14,18 @@
 
 #include "tensorstore/driver/downsample/downsample_nditerable.h"
 
+#include <stddef.h>
+#include <stdint.h>
+
 #include <algorithm>
 #include <array>
 #include <cassert>
+#include <complex>
+#include <functional>
+#include <limits>
 #include <numeric>
 #include <type_traits>
+#include <utility>
 #include <vector>
 
 #include "absl/log/absl_log.h"
@@ -33,8 +40,8 @@
 #include "tensorstore/internal/elementwise_function.h"
 #include "tensorstore/internal/nditerable.h"
 #include "tensorstore/internal/nditerable_buffer_management.h"
-#include "tensorstore/internal/nditerable_util.h"
 #include "tensorstore/internal/unique_with_intrusive_allocator.h"
+#include "tensorstore/rank.h"
 #include "tensorstore/util/extents.h"
 #include "tensorstore/util/iterate.h"
 #include "tensorstore/util/span.h"
@@ -49,7 +56,6 @@
 
 namespace tensorstore {
 namespace internal_downsample {
-
 namespace {
 
 using ::tensorstore::internal::ArenaAllocator;
@@ -1005,7 +1011,7 @@ class DownsampledNDIterator : public NDIterator::Base<DownsampledNDIterator> {
       Index base_inclusive_min =
           initial_base_index * downsample_factor - downsample_dim_origin[i];
       if (base_dim >= base_iteration_rank - 2) {
-        // This downsampling factor applies to one of the 2 inner-most
+        // This downsampling factor applies to one of the 2 innermost
         // dimensions.
         const DimensionIndex inner_dim_i = base_dim - (base_iteration_rank - 2);
         --num_outer_downsample_dims;
@@ -1238,7 +1244,7 @@ class DownsampledNDIterable : public NDIterable::Base<DownsampledNDIterable> {
     DimensionIndex base_iteration_dimensions_[kMaxRank];
   };
 
-  std::ptrdiff_t GetWorkingMemoryBytesPerElement(
+  ptrdiff_t GetWorkingMemoryBytesPerElement(
       NDIterable::IterationLayoutView layout,
       IterationBufferKind buffer_kind) const override {
     NDIterable::IterationLayoutView base_layout;