You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Zero/ZeroLevel.NN/Services/ImagePreprocessor.cs

378 lines
19 KiB

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

using Microsoft.ML.OnnxRuntime.Tensors;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;
using SixLabors.ImageSharp.Processing;
using ZeroLevel.NN.Models;
namespace ZeroLevel.NN
{
public static class ImagePreprocessor
{
private static Action<Tensor<float>, float, int, int, int, int> _precompiledChannelFirstAction = new Action<Tensor<float>, float, int, int, int, int>((t, v, ind, c, i, j) => { t[ind, c, i, j] = v; });
private static Action<Tensor<float>, float, int, int, int, int> _precompiledChannelLastAction = new Action<Tensor<float>, float, int, int, int, int>((t, v, ind, c, i, j) => { t[ind, i, j, c] = v; });
private static Func<byte, int, float> PixelToTensorMethod(ImagePreprocessorOptions options)
{
if (options.Normalize)
{
if (options.Correction)
{
if (options.CorrectionFunc == null)
{
return new Func<byte, int, float>((b, i) => ((options.NormalizationMultiplier * (float)b) - options.Mean[i]) / options.Std[i]);
}
else
{
return new Func<byte, int, float>((b, i) => options.CorrectionFunc.Invoke(i, options.NormalizationMultiplier * (float)b));
}
}
else
{
return new Func<byte, int, float>((b, i) => options.NormalizationMultiplier * (float)b);
}
}
else if (options.Correction)
{
if (options.CorrectionFunc == null)
{
return new Func<byte, int, float>((b, i) => (((float)b) - options.Mean[i]) / options.Std[i]);
}
else
{
return new Func<byte, int, float>((b, i) => options.CorrectionFunc.Invoke(i, (float)b));
}
}
return new Func<byte, int, float>((b, _) => (float)b);
}
//private static int CalculateFragmentsCount(Image image, ImagePreprocessorOptions options)
//{
// int count = (options.Crop.SaveOriginal ? 1 : 0);
// var Sw = image.Width; // ширина оригинала
// var Sh = image.Height; // высота оригинала
// var CRw = options.Crop.Width; // ширина кропа
// var CRh = options.Crop.Height; // высота кропа
// var Dx = options.Crop.Overlap ? (int)(options.Crop.OverlapKoefWidth * CRw) : CRw; // сдвиг по OX к следующему кропу
// var Dy = options.Crop.Overlap ? (int)(options.Crop.OverlapKoefHeight * CRh) : CRh; // сдвиг по OY к следующему кропу
// for (int x = 0; x < Sw; x += Dx)
// {
// for (int y = 0; y < Sh; y += Dy)
// {
// count++;
// }
// }
// return count;
//}
private static int CalculateFragmentsCount(Image image, ImagePreprocessorOptions options)
{
int count = (options.Crop.SaveOriginal ? 1 : 0);
var Sw = image.Width; // ширина оригинала
var Sh = image.Height; // высота оригинала
var CRw = options.InputWidth; // ширина кропа (равна ширине входа, т.к. изображение отресайзено подобающим образом)
var CRh = options.InputHeight; // высота кропа (равна высоте входа, т.к. изображение отресайзено подобающим образом)
var Dx = options.Crop.Overlap ? (int)(options.Crop.OverlapKoefWidth * CRw) : CRw; // сдвиг по OX к следующему кропу
var Dy = options.Crop.Overlap ? (int)(options.Crop.OverlapKoefHeight * CRh) : CRh; // сдвиг по OY к следующему кропу
for (int x = 0; x < Sw; x += Dx)
{
for (int y = 0; y < Sh; y += Dy)
{
count++;
}
}
return count;
}
private static void FillTensor(Tensor<float> tensor, Image image, int index, ImagePreprocessorOptions options, Func<byte, int, float> pixToTensor)
{
var append = options.ChannelType == PredictorChannelType.ChannelFirst ? _precompiledChannelFirstAction : _precompiledChannelLastAction;
((Image<Rgb24>)image).ProcessPixelRows(pixels =>
{
if (options.InvertXY)
{
for (int y = 0; y < pixels.Height; y++)
{
Span<Rgb24> pixelSpan = pixels.GetRowSpan(y);
for (int x = 0; x < pixels.Width; x++)
{
if (options.BGR)
{
append(tensor, pixToTensor(pixelSpan[x].B, 0), index, 0, y, x);
append(tensor, pixToTensor(pixelSpan[x].G, 1), index, 1, y, x);
append(tensor, pixToTensor(pixelSpan[x].R, 2), index, 2, y, x);
}
else
{
append(tensor, pixToTensor(pixelSpan[x].R, 0), index, 0, y, x);
append(tensor, pixToTensor(pixelSpan[x].G, 1), index, 1, y, x);
append(tensor, pixToTensor(pixelSpan[x].B, 2), index, 2, y, x);
}
}
}
}
else
{
for (int y = 0; y < pixels.Height; y++)
{
Span<Rgb24> pixelSpan = pixels.GetRowSpan(y);
for (int x = 0; x < pixels.Width; x++)
{
if (options.BGR)
{
append(tensor, pixToTensor(pixelSpan[x].B, 0), index, 0, x, y);
append(tensor, pixToTensor(pixelSpan[x].G, 1), index, 1, x, y);
append(tensor, pixToTensor(pixelSpan[x].R, 2), index, 2, x, y);
}
else
{
append(tensor, pixToTensor(pixelSpan[x].R, 0), index, 0, x, y);
append(tensor, pixToTensor(pixelSpan[x].G, 1), index, 1, x, y);
append(tensor, pixToTensor(pixelSpan[x].B, 2), index, 2, x, y);
}
}
}
}
});
}
private static void FillTensor(Tensor<float> tensor, Image image, int startX, int startY, int w, int h, int index, ImagePreprocessorOptions options, Func<byte, int, float> pixToTensor)
{
var append = options.ChannelType == PredictorChannelType.ChannelFirst ? _precompiledChannelFirstAction : _precompiledChannelLastAction;
if (image.PixelType.BitsPerPixel != 24)
{
var i = image;
image = i.CloneAs<Rgb24>();
i.Dispose();
}
((Image<Rgb24>)image).ProcessPixelRows(pixels =>
{
if (options.InvertXY)
{
for (int y = startY; y < h; y++)
{
Span<Rgb24> pixelSpan = pixels.GetRowSpan(y);
for (int x = startX; x < w; x++)
{
if (options.BGR)
{
append(tensor, pixToTensor(pixelSpan[x].B, 0), index, 0, y, x);
append(tensor, pixToTensor(pixelSpan[x].G, 1), index, 1, y, x);
append(tensor, pixToTensor(pixelSpan[x].R, 2), index, 2, y, x);
}
else
{
append(tensor, pixToTensor(pixelSpan[x].R, 0), index, 0, y, x);
append(tensor, pixToTensor(pixelSpan[x].G, 1), index, 1, y, x);
append(tensor, pixToTensor(pixelSpan[x].B, 2), index, 2, y, x);
}
}
}
}
else
{
for (int y = startY; y < h; y++)
{
Span<Rgb24> pixelSpan = pixels.GetRowSpan(y);
for (int x = startX; x < w; x++)
{
if (options.BGR)
{
append(tensor, pixToTensor(pixelSpan[x].B, 0), index, 0, x, y);
append(tensor, pixToTensor(pixelSpan[x].G, 1), index, 1, x, y);
append(tensor, pixToTensor(pixelSpan[x].R, 2), index, 2, x, y);
}
else
{
append(tensor, pixToTensor(pixelSpan[x].R, 0), index, 0, x, y);
append(tensor, pixToTensor(pixelSpan[x].G, 1), index, 1, x, y);
append(tensor, pixToTensor(pixelSpan[x].B, 2), index, 2, x, y);
}
}
}
}
});
}
private static Tensor<float> InitInputTensor(ImagePreprocessorOptions options, int batchSize = 1)
{
switch (options.ChannelType)
{
case PredictorChannelType.ChannelFirst:
return options.InvertXY
? new DenseTensor<float>(new[] { batchSize, options.Channels, options.InputHeight, options.InputWidth })
: new DenseTensor<float>(new[] { batchSize, options.Channels, options.InputWidth, options.InputHeight });
default:
return options.InvertXY
? new DenseTensor<float>(new[] { batchSize, options.InputHeight, options.InputWidth, options.Channels })
: new DenseTensor<float>(new[] { batchSize, options.InputWidth, options.InputHeight, options.Channels });
}
}
public static ImagePredictionInput[] ToTensors(this Image image, ImagePreprocessorOptions options)
{
ImagePredictionInput[] result = null;
var pixToTensor = PixelToTensorMethod(options);
options.Channels = image.PixelType.BitsPerPixel >> 3;
if (options.Crop.Enabled)
{
// Размеры оригинального изображения
var Sw = image.Width;
var Sh = image.Height;
// Создание ресайза для целочисленного прохода кропами шириной CRw и высотой CRh
var resizedForCropWidthKoef = options.InputWidth / (double)options.Crop.Width;
var resizedForCropHeightKoef = options.InputHeight / (double)options.Crop.Height;
// Размеры для ресайза изображения к размеру по которому удобно идти кропами
var resizedForCropWidth = (int)Math.Round(Sw * resizedForCropWidthKoef, MidpointRounding.ToEven);
var resizedForCropHeight = (int)Math.Round(Sh * resizedForCropHeightKoef, MidpointRounding.ToEven);
// Размеры кропа, равны входу сети, а не (options.Crop.Width, options.Crop.Height), т.к. для оптимизации изображение будет предварительно отресайзено
var CRw = options.InputWidth;
var CRh = options.InputHeight;
// Расчет сдвигов между кропами
var Dx = options.Crop.Overlap ? (int)(options.Crop.OverlapKoefWidth * CRw) : CRw;
var Dy = options.Crop.Overlap ? (int)(options.Crop.OverlapKoefHeight * CRh) : CRh;
using (var source = image.Clone(img => img.Resize(resizedForCropWidth, resizedForCropHeight, KnownResamplers.Bicubic)))
{
// Количество тензоров всего, во всех батчах суммарно
var count = CalculateFragmentsCount(source, options);
// Проверка, укладывается ли количество тензоров поровну в батчи
int offset = count % options.MaxBatchSize;
// Количество батчей
int count_tensor_batches = count / options.MaxBatchSize + (offset == 0 ? 0 : 1);
// Батчи
var tensors = new ImagePredictionInput[count_tensor_batches];
// Инициализация батчей
Parallel.For(0, count_tensor_batches, batch_index =>
{
if (batch_index < count_tensor_batches - 1)
{
tensors[batch_index] = new ImagePredictionInput
{
Tensor = InitInputTensor(options, options.MaxBatchSize),
Offsets = new OffsetBox[options.MaxBatchSize],
Count = options.MaxBatchSize
};
}
else
{
tensors[batch_index] = new ImagePredictionInput
{
Tensor = InitInputTensor(options, offset == 0 ? options.MaxBatchSize : offset),
Offsets = new OffsetBox[offset == 0 ? options.MaxBatchSize : offset],
Count = offset == 0 ? options.MaxBatchSize : offset
};
}
});
// Заполнение батчей
int tensor_index = 0;
// Если используется ресайз оригинала кроме кропов, пишется в первый батч в первый тензор
if (options.Crop.SaveOriginal)
{
using (var copy = source.Clone(img => img.Resize(options.InputWidth, options.InputHeight, KnownResamplers.Bicubic)))
{
FillTensor(tensors[0].Tensor, copy, 0, options, pixToTensor);
tensors[tensor_index].Offsets[0] = new OffsetBox(0, 0, image.Width, image.Height);
}
tensor_index++;
}
tensor_index--;
Parallel.ForEach(SteppedIterator(0, source.Width, Dx), x =>
{
// Можно запараллелить и тут, но выигрыш дает малоощутимый
for (int y = 0; y < source.Height; y += Dy)
{
var current_index = Interlocked.Increment(ref tensor_index);
// Индекс тензора внутри батча
var b_index = current_index % options.MaxBatchSize;
// Индекс батча
var p_index = (int)Math.Round((double)current_index / (double)options.MaxBatchSize, MidpointRounding.ToNegativeInfinity);
int w = CRw;
if ((x + CRw) > source.Width)
{
w = source.Width - x;
}
int h = CRh;
if ((y + CRh) > source.Height)
{
h = source.Height - y;
}
// Заполнение b_index тензора в p_index батче
FillTensor(tensors[p_index].Tensor, source, x, y, w, h, b_index, options, pixToTensor);
// Указание смещений для данного тензора
tensors[p_index].Offsets[b_index] = new OffsetBox(x, y, options.Crop.Width, options.Crop.Height);
}
});
return tensors;
}
}
// if resize only
result = new ImagePredictionInput[1];
using (var copy = image.Clone(img => img.Resize(options.InputWidth, options.InputHeight, KnownResamplers.Bicubic)))
{
Tensor<float> tensor = InitInputTensor(options);
FillTensor(tensor, copy, 0, options, pixToTensor);
result[0] = new ImagePredictionInput
{
Count = 1,
Offsets = null,
Tensor = tensor
};
}
return result;
}
private static IEnumerable<int> SteppedIterator(int startIndex, int endIndex, int stepSize)
{
for (int i = startIndex; i < endIndex; i += stepSize)
{
yield return i;
}
}
public static Image Crop(Image source, float x1, float y1, float x2, float y2)
{
int left = 0;
int right = 0;
int top = 0;
int bottom = 0;
int width = (int)(x2 - x1);
int height = (int)(y2 - y1);
if (x1 < 0) { left = (int)-x1; x1 = 0; }
if (x2 > source.Width) { right = (int)(x2 - source.Width); x2 = source.Width - 1; }
if (y1 < 0) { top = (int)-y1; y1 = 0; }
if (y2 > source.Height) { bottom = (int)(y2 - source.Height); y2 = source.Height - 1; }
if (left + right + top + bottom > 0)
{
var backgroundImage = new Image<Rgb24>(SixLabors.ImageSharp.Configuration.Default, width, height, new Rgb24(0, 0, 0));
using (var crop = source.Clone(img => img.Crop(new Rectangle((int)x1, (int)y1, (int)(x2 - x1), (int)(y2 - y1)))))
{
backgroundImage.Mutate(bg => bg.DrawImage(crop, new Point(left, top), 1f));
}
return backgroundImage;
}
return source.Clone(img => img.Crop(new Rectangle((int)x1, (int)y1, (int)(x2 - x1), (int)(y2 - y1))));
}
}
}

Powered by TurnKey Linux.