Merge branch 'main' of https://github.com/dme-compunet/YoloV8

dme-compunet · Feb 4, 2024 · c2fd69e · c2fd69e
2 parents 8c92e75 + f7c6cb8
commit c2fd69e
Showing 1 changed file with 21 additions and 25 deletions.
diff --git a/Source/YoloV8/Utilities/PreprocessHelper.cs b/Source/YoloV8/Utilities/PreprocessHelper.cs
@@ -18,54 +18,50 @@ public static void ProcessToTensor(Image<Rgb24> image, Size modelSize, bool orig
         var width = image.Width;
         var height = image.Height;
 
+        // Pre-calculate strides for performance
+        var strideBatchR = target.Strides[0] * batch + target.Strides[1] * 0;
+        var strideBatchG = target.Strides[0] * batch + target.Strides[1] * 1;
+        var strideBatchB = target.Strides[0] * batch + target.Strides[1] * 2;
+        var strideY = target.Strides[2];
+        var strideX = target.Strides[3];
+
+        // Get a span of the whole tensor for fast access
+        var tensorSpan = target.Buffer;
+
         // Try get continuous memory block of the entire image data
         if (image.DangerousTryGetSinglePixelMemory(out var memory))
         {
             Parallel.For(0, width * height, index =>
             {
                 int x = index % width;
                 int y = index / width;
+                int tensorIndex = strideBatchR + strideY * (y + yPadding) + strideX * (x + xPadding);
 
                 var pixel = memory.Span[index];
-
-                WritePixel(batch, y + yPadding, x + xPadding, pixel, target);
+                WritePixel(tensorSpan.Span, tensorIndex, pixel, strideBatchR, strideBatchG, strideBatchB);
             });
         }
         else
         {
             Parallel.For(0, height, y =>
             {
-                var row = image.DangerousGetPixelRowMemory(y).Span;
+                var rowSpan = image.DangerousGetPixelRowMemory(y).Span;
+                int tensorYIndex = strideBatchR + strideY * (y + yPadding);
 
                 for (int x = 0; x < width; x++)
                 {
-                    var pixel = row[x];
-
-                    WritePixel(batch, y + yPadding, x + xPadding, pixel, target);
+                    int tensorIndex = tensorYIndex + strideX * (x + xPadding);
+                    var pixel = rowSpan[x];
+                    WritePixel(tensorSpan.Span, tensorIndex, pixel, strideBatchR, strideBatchG, strideBatchB);
                 }
             });
         }
     }
 
-    private static void WritePixel(int batch, int y, int x, Rgb24 pixel, DenseTensor<float> target)
+    private static void WritePixel(Span<float> tensorSpan, int tensorIndex, Rgb24 pixel, int strideBatchR, int strideBatchG, int strideBatchB)
     {
-        var offsetR = target.Strides[0] * batch
-                    + target.Strides[1] * 0
-                    + target.Strides[2] * y
-                    + target.Strides[3] * x;
-
-        var offsetG = target.Strides[0] * batch
-                    + target.Strides[1] * 1
-                    + target.Strides[2] * y
-                    + target.Strides[3] * x;
-
-        var offsetB = target.Strides[0] * batch
-                    + target.Strides[1] * 2
-                    + target.Strides[2] * y
-                    + target.Strides[3] * x;
-
-        target.Buffer.Span[offsetR] = pixel.R / 255f;
-        target.Buffer.Span[offsetG] = pixel.G / 255f;
-        target.Buffer.Span[offsetB] = pixel.B / 255f;
+        tensorSpan[tensorIndex] = pixel.R / 255f;
+        tensorSpan[tensorIndex + strideBatchG - strideBatchR] = pixel.G / 255f;
+        tensorSpan[tensorIndex + strideBatchB - strideBatchR] = pixel.B / 255f;
     }
 }