void OutputDataProcessingCPU(ImageManipulationModels imageManipulationModels, float[] outputData, Color32[] pixelBuffer, Color32[] srcBuffer = null)
{
switch (imageManipulationModels)
{
case ImageManipulationModels.Colorization:
// FIXME: stretch to original size
for (int i = 0; i < pixelBuffer.Length; i++)
{
var lab = AiliaColorConv.Color2Lab(srcBuffer[i]);
var nlab = new AiliaColorConv.LAB(
lab.l,
outputData[i + 0 * pixelBuffer.Length],
outputData[i + 1 * pixelBuffer.Length]
);
pixelBuffer[i] = AiliaColorConv.Lab2Color(nlab);
}
break;
default:
for (int i = 0; i < pixelBuffer.Length; i++)
{
pixelBuffer[i].r = (byte)Mathf.Clamp(outputData[i + 0 * pixelBuffer.Length] * 255, 0, 255);
pixelBuffer[i].g = (byte)Mathf.Clamp(outputData[i + 1 * pixelBuffer.Length] * 255, 0, 255);
pixelBuffer[i].b = (byte)Mathf.Clamp(outputData[i + 2 * pixelBuffer.Length] * 255, 0, 255);
pixelBuffer[i].a = 255;
}
break;
}
}
void InputDataPocessingCPU(ImageManipulationModels imageManipulationModels, Color32[] inputImage, float[] processedInputBuffer)
{
float weight = 1f / 255f;
float bias = 0;
bool rgbRepeats = false;
switch (imageManipulationModels)
{
case ImageManipulationModels.IlluminationCorrection:
weight = 1f / 127.5f;
bias = -1;
break;
case ImageManipulationModels.Colorization:
if (rgbRepeats)
{
for (int i = 0; i < inputImage.Length; i++)
{
var r = (inputImage[i].r) * weight + bias;
var g = (inputImage[i].g) * weight + bias;
var b = (inputImage[i].b) * weight + bias;
var lab = AiliaColorConv.Color2Lab(new Color(r, g, b));
processedInputBuffer[i * 3 + 0] = (float)lab.L;
processedInputBuffer[i * 3 + 1] = (float)lab.A;
processedInputBuffer[i * 3 + 2] = (float)lab.B;
}
}
else
{
for (int i = 0; i < inputImage.Length; i++)
{
var r = (inputImage[i].r) * weight + bias;
var g = (inputImage[i].g) * weight + bias;
var b = (inputImage[i].b) * weight + bias;
var lab = AiliaColorConv.Color2Lab(new Color(r, g, b));
processedInputBuffer[i + inputImage.Length * 0] = (float)lab.L;
}
}
return;
default:
break;
}
// flatten input data
if (rgbRepeats)
{
for (int i = 0; i < inputImage.Length; i++)
{
// rgbrgbrgb...
processedInputBuffer[i * 3 + 0] = (inputImage[i].r) * weight + bias;
processedInputBuffer[i * 3 + 1] = (inputImage[i].g) * weight + bias;
processedInputBuffer[i * 3 + 2] = (inputImage[i].b) * weight + bias;
}
}
else
{
for (int i = 0; i < inputImage.Length; i++)
{
processedInputBuffer[i + inputImage.Length * 0] = (inputImage[i].r) * weight + bias;
processedInputBuffer[i + inputImage.Length * 1] = (inputImage[i].g) * weight + bias;
processedInputBuffer[i + inputImage.Length * 2] = (inputImage[i].b) * weight + bias;
}
}
}
void Update()
{
if (!AiliaImageSource.IsPrepared || !modelPrepared)
{
return;
}
if (output == null)
{
float rawImageRatio = rawImageSize.x / rawImageSize.y;
float ratio = AiliaImageSource.Width / (float)AiliaImageSource.Height;
raw_image.rectTransform.sizeDelta = new Vector2(ratio / rawImageRatio * rawImageSize.x, rawImageSize.y);
SetShape(imageManipulationModels);
// texture & buffer allocate
if (!gpu_mode || outputDataToTextureShader == null)
{
resultTexture2D = new Texture2D(OutputWidth, OutputHeight, TextureFormat.RGBA32, false);
}
else
{
resultRenderTexture = new RenderTexture(OutputWidth, OutputHeight, 0);
resultRenderTexture.enableRandomWrite = true;
resultRenderTexture.Create();
}
baseTexture = AiliaImageSource.GetTexture(new Rect(0, 0, AiliaImageSource.Width, AiliaImageSource.Height));
AiliaImageSource.Resize(InputWidth, InputHeight);
input = new float[InputWidth * InputHeight * InputChannel];
output = new float[OutputWidth * OutputHeight * OutputChannel];
outputImage = new Color32[OutputWidth * OutputHeight];
}
if (oneshot)
{
oneshot = false;
// Make input data
Rect rect = new Rect(0, 0, InputWidth, InputHeight);
long start_time = DateTime.UtcNow.Ticks / TimeSpan.TicksPerMillisecond;
Color32[] inputImage = null;
if (!gpu_mode || inputDataProcessingShader == null)
{
inputImage = AiliaImageSource.GetPixels32(rect, true);
InputDataPocessingCPU(imageManipulationModels, inputImage, input);
}
else
{
originalTexture = AiliaImageSource.GetTexture(rect);
InputDataPocessing(imageManipulationModels, originalTexture, input, true);
}
long end_time = DateTime.UtcNow.Ticks / TimeSpan.TicksPerMillisecond;
// Predict
long start_time2 = DateTime.UtcNow.Ticks / TimeSpan.TicksPerMillisecond;
bool result = ailiaModel.Predict(output, input);
// convert result to image
if (!gpu_mode || outputDataToTextureShader == null)
{
OutputDataProcessingCPU(imageManipulationModels, output, outputImage, inputImage);
if (imageManipulationModels == ImageManipulationModels.Colorization)
{
resultTexture2D.SetPixels32(outputImage);
resultTexture2D.Apply();
{
var rTexture = AiliaImageUtil.ResizeTexture(resultTexture2D, baseTexture.width, baseTexture.height);
var base32 = baseTexture.GetPixels32();
var dst32 = rTexture.GetPixels32();
Color32[] dstColorBuffer = new Color32[dst32.Length];
for (int i = 0; i < dst32.Length; i++)
{
var px = i % baseTexture.width;
var py = (baseTexture.height - 1 - (i / baseTexture.width)) * baseTexture.width;
var baselab = AiliaColorConv.Color2Lab(base32[i]);
var dstlab = AiliaColorConv.Color2Lab(dst32[px + py]);
var resultlab = new AiliaColorConv.LAB(baselab.L, dstlab.A, dstlab.B);
dstColorBuffer[px + py] = AiliaColorConv.Lab2Color(resultlab);
}
outputImage = dstColorBuffer;
}
resultTexture2D = new Texture2D(baseTexture.width, baseTexture.height);
}
}
else
{
OutputDataProcessing(output, resultRenderTexture);
}
long end_time2 = DateTime.UtcNow.Ticks / TimeSpan.TicksPerMillisecond;
if (label_text != null)
{
label_text.text = ((end_time - start_time) + (end_time2 - start_time2)).ToString() + "ms\n" + ailiaModel.EnvironmentName();
}
// for viewer
if (!gpu_mode || inputDataProcessingShader == null)
{
originalTexture = new Texture2D(InputWidth, InputHeight, TextureFormat.RGBA32, false);
originalTexture.SetPixels32(inputImage);
originalTexture.Apply();
}
raw_image.texture = originalTexture;
blendMaterial.SetTexture(mainTexId, originalTexture);
if (!gpu_mode || outputDataToTextureShader == null)
{
resultTexture2D.SetPixels32(outputImage);
resultTexture2D.Apply();
blendMaterial.SetTexture(blendTexId, resultTexture2D);
}
else
{
blendMaterial.SetTexture(blendTexId, resultRenderTexture);
}
blendMaterial.SetFloat(mainVFlipId, 0);
blendMaterial.SetFloat(blendVFlipId, 1);
raw_image.gameObject.SetActive(true);
}
// When space key down, draw original image
if (Input.GetKey(KeyCode.Space))
{
blendMaterial.SetFloat(blendFlagId, 0);
}
else
{
blendMaterial.SetFloat(blendFlagId, 1);
}
}