11using Microsoft . ML . OnnxRuntime . Tensors ;
2+ using OnnxStack . Core . Model ;
23using SixLabors . ImageSharp ;
34using SixLabors . ImageSharp . PixelFormats ;
45using SixLabors . ImageSharp . Processing ;
6+ using System . Threading . Tasks ;
57
68namespace OnnxStack . Core . Image
79{
@@ -14,6 +16,17 @@ public static class Extensions
1416 /// <param name="imageTensor">The image tensor.</param>
1517 /// <returns></returns>
1618 public static OnnxImage ToImageMask ( this DenseTensor < float > imageTensor )
19+ {
20+ return new OnnxImage ( imageTensor . FromMaskTensor ( ) ) ;
21+ }
22+
23+
24+ /// <summary>
25+ /// Convert from single channle mask tensor to Rgba32 (Greyscale)
26+ /// </summary>
27+ /// <param name="imageTensor">The image tensor.</param>
28+ /// <returns></returns>
29+ public static Image < Rgba32 > FromMaskTensor ( this DenseTensor < float > imageTensor )
1730 {
1831 var width = imageTensor . Dimensions [ 3 ] ;
1932 var height = imageTensor . Dimensions [ 2 ] ;
@@ -26,7 +39,7 @@ public static OnnxImage ToImageMask(this DenseTensor<float> imageTensor)
2639 result [ x , y ] = new L8 ( ( byte ) ( imageTensor [ 0 , 0 , y , x ] * 255.0f ) ) ;
2740 }
2841 }
29- return new OnnxImage ( result . CloneAs < Rgba32 > ( ) ) ;
42+ return result . CloneAs < Rgba32 > ( ) ;
3043 }
3144 }
3245
@@ -40,6 +53,104 @@ public static ResizeMode ToResizeMode(this ImageResizeMode resizeMode)
4053 } ;
4154 }
4255
56+
57+ /// <summary>
58+ /// Splits the Tensor into 4 equal tiles.
59+ /// </summary>
60+ /// <param name="sourceTensor">The source tensor.</param>
61+ /// <returns>TODO: Optimize</returns>
62+ public static ImageTiles SplitImageTiles ( this DenseTensor < float > sourceTensor , int overlap = 20 )
63+ {
64+ var tileWidth = sourceTensor . Dimensions [ 3 ] / 2 ;
65+ var tileHeight = sourceTensor . Dimensions [ 2 ] / 2 ;
66+ return new ImageTiles ( tileWidth , tileHeight , overlap ,
67+ SplitImageTile ( sourceTensor , 0 , 0 , tileHeight + overlap , tileWidth + overlap ) ,
68+ SplitImageTile ( sourceTensor , 0 , tileWidth - overlap , tileHeight + overlap , tileWidth * 2 ) ,
69+ SplitImageTile ( sourceTensor , tileHeight - overlap , 0 , tileHeight * 2 , tileWidth + overlap ) ,
70+ SplitImageTile ( sourceTensor , tileHeight - overlap , tileWidth - overlap , tileHeight * 2 , tileWidth * 2 ) ) ;
71+ }
72+
73+
74+ /// <summary>
75+ /// Splits a tile from the source.
76+ /// </summary>
77+ /// <param name="source">The tensor.</param>
78+ /// <param name="startRow">The start row.</param>
79+ /// <param name="startCol">The start col.</param>
80+ /// <param name="endRow">The end row.</param>
81+ /// <param name="endCol">The end col.</param>
82+ /// <returns></returns>
83+ private static DenseTensor < float > SplitImageTile ( DenseTensor < float > source , int startRow , int startCol , int endRow , int endCol )
84+ {
85+ int height = endRow - startRow ;
86+ int width = endCol - startCol ;
87+ int channels = source . Dimensions [ 1 ] ;
88+ var splitTensor = new DenseTensor < float > ( new [ ] { 1 , channels , height , width } ) ;
89+ Parallel . For ( 0 , channels , ( c ) =>
90+ {
91+ Parallel . For ( 0 , height , ( i ) =>
92+ {
93+ Parallel . For ( 0 , width , ( j ) =>
94+ {
95+ splitTensor [ 0 , c , i , j ] = source [ 0 , c , startRow + i , startCol + j ] ;
96+ } ) ;
97+ } ) ;
98+ } ) ;
99+ return splitTensor ;
100+ }
101+
102+
103+ /// <summary>
104+ /// Joins the tiles into a single Tensor.
105+ /// </summary>
106+ /// <param name="tiles">The tiles.</param>
107+ /// <returns>TODO: Optimize</returns>
108+ public static DenseTensor < float > JoinImageTiles ( this ImageTiles tiles )
109+ {
110+ var totalWidth = tiles . Width * 2 ;
111+ var totalHeight = tiles . Height * 2 ;
112+ var channels = tiles . Tile1 . Dimensions [ 1 ] ;
113+ var destination = new DenseTensor < float > ( new [ ] { 1 , channels , totalHeight , totalWidth } ) ;
114+ JoinImageTile ( destination , tiles . Tile1 , 0 , 0 , tiles . Height + tiles . Overlap , tiles . Width + tiles . Overlap ) ;
115+ JoinImageTile ( destination , tiles . Tile2 , 0 , tiles . Width - tiles . Overlap , tiles . Height + tiles . Overlap , totalWidth ) ;
116+ JoinImageTile ( destination , tiles . Tile3 , tiles . Height - tiles . Overlap , 0 , totalHeight , tiles . Width + tiles . Overlap ) ;
117+ JoinImageTile ( destination , tiles . Tile4 , tiles . Height - tiles . Overlap , tiles . Width - tiles . Overlap , totalHeight , totalWidth ) ;
118+ return destination ;
119+ }
120+
121+
122+ /// <summary>
123+ /// Joins the tile to the destination tensor.
124+ /// </summary>
125+ /// <param name="destination">The destination.</param>
126+ /// <param name="tile">The tile.</param>
127+ /// <param name="startRow">The start row.</param>
128+ /// <param name="startCol">The start col.</param>
129+ /// <param name="endRow">The end row.</param>
130+ /// <param name="endCol">The end col.</param>
131+ private static void JoinImageTile ( DenseTensor < float > destination , DenseTensor < float > tile , int startRow , int startCol , int endRow , int endCol )
132+ {
133+ int height = endRow - startRow ;
134+ int width = endCol - startCol ;
135+ int channels = tile . Dimensions [ 1 ] ;
136+ Parallel . For ( 0 , channels , ( c ) =>
137+ {
138+ Parallel . For ( 0 , height , ( i ) =>
139+ {
140+ Parallel . For ( 0 , width , ( j ) =>
141+ {
142+ var value = tile [ 0 , c , i , j ] ;
143+ var existing = destination [ 0 , c , startRow + i , startCol + j ] ;
144+ if ( existing > 0 )
145+ {
146+ // Blend ovelap
147+ value = ( existing + value ) / 2f ;
148+ }
149+ destination [ 0 , c , startRow + i , startCol + j ] = value ;
150+ } ) ;
151+ } ) ;
152+ } ) ;
153+ }
43154 }
44155
45156 public enum ImageNormalizeType
0 commit comments