Methods to split and join a tensor into tiles

Author: saddam213

OnnxStack.Core/Extensions/TensorExtension.cs

@@ -4,10 +4,11 @@
 using System.Linq;
 using System.Numerics.Tensors;
 using System.Numerics;
+using OnnxStack.Core.Model;
 
 namespace OnnxStack.Core
 {
-    public static class TensorExtension
+    public static partial class TensorExtension
     {
         /// <summary>
         /// Divides the tensor by float.
@@ -423,5 +424,78 @@ private static DenseTensor<float> ConcatenateAxis2(DenseTensor<float> tensor1, D
 
             return concatenatedTensor;
         }
+
+
+        /// <summary>
+        /// Splits the Tensor into 4 equal tiles.
+        /// </summary>
+        /// <param name="sourceTensor">The source tensor.</param>
+        /// <returns>TODO: Optimize</returns>
+        public static ImageTiles SplitTiles(this DenseTensor<float> sourceTensor)
+        {
+            int tileWidth = sourceTensor.Dimensions[3] / 2;
+            int tileHeight = sourceTensor.Dimensions[2] / 2;
+
+            return new ImageTiles(
+                SplitTile(sourceTensor, 0, 0, tileHeight, tileWidth),
+                SplitTile(sourceTensor, 0, tileWidth, tileHeight, tileWidth * 2),
+                SplitTile(sourceTensor, tileHeight, 0, tileHeight * 2, tileWidth),
+                SplitTile(sourceTensor, tileHeight, tileWidth, tileHeight * 2, tileWidth * 2));
+        }
+
+        private static DenseTensor<float> SplitTile(DenseTensor<float> tensor, int startRow, int startCol, int endRow, int endCol)
+        {
+            int height = endRow - startRow;
+            int width = endCol - startCol;
+            int channels = tensor.Dimensions[1];
+            var slicedData = new DenseTensor<float>(new[] { 1, channels, height, width });
+            for (int c = 0; c < channels; c++)
+            {
+                for (int i = 0; i < height; i++)
+                {
+                    for (int j = 0; j < width; j++)
+                    {
+                        slicedData[0, c, i, j] = tensor[0, c, startRow + i, startCol + j];
+                    }
+                }
+            }
+            return slicedData;
+        }
+
+
+        /// <summary>
+        /// Rejoins the tiles into a single Tensor.
+        /// </summary>
+        /// <param name="tiles">The tiles.</param>
+        /// <returns>TODO: Optimize</returns>
+        public static DenseTensor<float> RejoinTiles(this ImageTiles tiles)
+        {
+            int totalHeight = tiles.Tile1.Dimensions[2] + tiles.Tile3.Dimensions[2];
+            int totalWidth = tiles.Tile1.Dimensions[3] + tiles.Tile2.Dimensions[3];
+            int channels = tiles.Tile1.Dimensions[1];
+            var destination = new DenseTensor<float>(new[] { 1, channels, totalHeight, totalWidth });
+            RejoinTile(destination, tiles.Tile1, 0, 0);
+            RejoinTile(destination, tiles.Tile2, 0, tiles.Tile1.Dimensions[3]);
+            RejoinTile(destination, tiles.Tile3, tiles.Tile1.Dimensions[2], 0);
+            RejoinTile(destination, tiles.Tile4, tiles.Tile1.Dimensions[2], tiles.Tile1.Dimensions[3]);
+            return destination;
+        }
+
+        private static void RejoinTile(DenseTensor<float> destination, DenseTensor<float> tile, int startRow, int startCol)
+        {
+            int channels = tile.Dimensions[1];
+            int height = tile.Dimensions[2];
+            int width = tile.Dimensions[3];
+            for (int c = 0; c < channels; c++)
+            {
+                for (int i = 0; i < height; i++)
+                {
+                    for (int j = 0; j < width; j++)
+                    {
+                        destination[0, c, startRow + i, startCol + j] = tile[0, c, i, j];
+                    }
+                }
+            }
+        }
     }
 }

OnnxStack.Core/Model/ImageTiles.cs

@@ -0,0 +1,6 @@
+using Microsoft.ML.OnnxRuntime.Tensors;
+
+namespace OnnxStack.Core.Model
+{
+    public record ImageTiles(DenseTensor<float> Tile1, DenseTensor<float> Tile2, DenseTensor<float> Tile3, DenseTensor<float> Tile4);
+}

OnnxStack.Core/OnnxStack.Core.csproj

@@ -44,7 +44,7 @@
     <PackageReference Include="Microsoft.ML" Version="3.0.1" />
     <PackageReference Include="Microsoft.ML.OnnxRuntime.Extensions" Version="0.10.0" />
     <PackageReference Include="Microsoft.ML.OnnxRuntime.Managed" Version="1.17.1" />
-    <PackageReference Include="SixLabors.ImageSharp" Version="3.1.2" />
+    <PackageReference Include="SixLabors.ImageSharp" Version="3.1.3" />
     <PackageReference Include="System.Linq.Async" Version="6.0.1" />
   </ItemGroup>