ImageProcessor.cs

// Copyright (c) Zhongkai Fu. All rights reserved.
// https://github.com/zhongkaifu/TensorSharp
//
// This file is part of TensorSharp.
//
// TensorSharp is licensed under the BSD-3-Clause license found in the LICENSE file in the root directory of this source tree.
//
// TensorSharp is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the BSD-3-Clause License for more details.
using System;
using System.IO;
using System.Threading.Tasks;
using ImageMagick;
using StbImageSharp;

namespace TensorSharp.Models
{
    public class Gemma3ImageProcessor
    {
        public int TokensPerImage { get; }
        public int ImageSize { get; }
        public int PatchSize { get; }

        public const int StartOfImageToken = 255999;
        public const int EndOfImageToken = 256000;
        public const int NewlineNewlineToken = 108;
        public const int PadToken = 0;

        private const float ImageMean = 0.5f;
        private const float ImageStd = 0.5f;

        public Gemma3ImageProcessor(int tokensPerImage = 256, int imageSize = 896, int patchSize = 14)
        {
            TokensPerImage = tokensPerImage;
            ImageSize = imageSize;
            PatchSize = patchSize;
        }

        /// <summary>
        /// Process an image file into normalized pixel values in channel-first format [C, H, W].
        /// Matches Ollama's processing: composite over white, bilinear resize, normalize with mean=0.5, std=0.5.
        /// </summary>
        public float[] ProcessImage(string imagePath)
        {
            byte[] fileBytes = File.ReadAllBytes(imagePath);
            int origWidth, origHeight;
            byte[] rgba = DecodeImageToRGBA(fileBytes, out origWidth, out origHeight);
            return ResizeRgbaToChannelFirstNormalized(rgba, origWidth, origHeight, ImageSize, ImageSize);
        }

        internal static byte[] DecodeImageToRGBA(byte[] fileBytes, out int width, out int height)
        {
            if (IsPng(fileBytes))
                return DecodePNG(fileBytes, out width, out height);

            if (IsJpeg(fileBytes))
                return DecodeJPEG(fileBytes, out width, out height);

            if (IsHeic(fileBytes))
                return DecodeHEIC(fileBytes, out width, out height);

            throw new NotSupportedException("Only PNG, JPEG, and HEIC/HEIF image formats are supported");
        }

        internal static (int width, int height) ReadImageDimensions(string imagePath)
        {
            byte[] fileBytes = File.ReadAllBytes(imagePath);

            if (IsPng(fileBytes))
                return ReadPngDimensions(fileBytes);

            if (IsJpeg(fileBytes))
            {
                DecodeJPEG(fileBytes, out int width, out int height);
                return (width, height);
            }

            if (IsHeic(fileBytes))
                return ReadHeicDimensions(fileBytes);

            throw new NotSupportedException("Only PNG, JPEG, and HEIC/HEIF image formats are supported");
        }

        private static bool IsPng(byte[] fileBytes) =>
            fileBytes.Length >= 8 &&
            fileBytes[0] == 0x89 &&
            fileBytes[1] == 0x50 &&
            fileBytes[2] == 0x4E &&
            fileBytes[3] == 0x47;

        private static bool IsJpeg(byte[] fileBytes) =>
            fileBytes.Length >= 2 &&
            fileBytes[0] == 0xFF &&
            fileBytes[1] == 0xD8;

        // HEIC/HEIF files use the ISOBMFF container: a leading "ftyp" box whose
        // major_brand (offset 8..11) or compatible_brands (offsets 16, 20, 24, ...
        // up to the box size) carry the HEIF/HEIC marker. Detect the major brand
        // first and then scan the compatible_brands list so files authored with a
        // generic major_brand (e.g. "mif1") but a HEIC/HEIF compatible brand are
        // still recognised.
        private static bool IsHeic(byte[] fileBytes)
        {
            if (fileBytes.Length < 12)
                return false;

            if (fileBytes[4] != (byte)'f' || fileBytes[5] != (byte)'t' ||
                fileBytes[6] != (byte)'y' || fileBytes[7] != (byte)'p')
            {
                return false;
            }

            int boxSize = (fileBytes[0] << 24) | (fileBytes[1] << 16) | (fileBytes[2] << 8) | fileBytes[3];
            if (boxSize <= 0 || boxSize > fileBytes.Length)
                boxSize = fileBytes.Length;

            if (IsHeifBrand(fileBytes, 8))
                return true;

            for (int offset = 16; offset + 4 <= boxSize; offset += 4)
            {
                if (IsHeifBrand(fileBytes, offset))
                    return true;
            }

            return false;
        }

        private static bool IsHeifBrand(byte[] data, int offset)
        {
            if (offset < 0 || offset + 4 > data.Length)
                return false;

            string brand = System.Text.Encoding.ASCII.GetString(data, offset, 4);
            return brand switch
            {
                "heic" or "heix" or "heim" or "heis" or
                "hevc" or "hevx" or "hevm" or "hevs" or
                "mif1" or "msf1" or "heif" => true,
                _ => false,
            };
        }

        private static byte[] DecodePNG(byte[] data, out int width, out int height)
        {
            using var ms = new MemoryStream(data);
            using var reader = new BinaryReader(ms);

            byte[] sig = reader.ReadBytes(8);
            if (sig[0] != 0x89 || sig[1] != 0x50 || sig[2] != 0x4E || sig[3] != 0x47)
                throw new InvalidDataException("Not a PNG file");

            width = 0; height = 0;
            int bitDepth = 0, colorType = 0;
            using var idatStream = new MemoryStream();

            while (ms.Position < ms.Length)
            {
                int length = ReadBigEndianInt32(reader);
                byte[] chunkType = reader.ReadBytes(4);
                string type = System.Text.Encoding.ASCII.GetString(chunkType);

                if (type == "IHDR")
                {
                    width = ReadBigEndianInt32(reader);
                    height = ReadBigEndianInt32(reader);
                    bitDepth = reader.ReadByte();
                    colorType = reader.ReadByte();
                    reader.ReadBytes(3 + 4); // compression, filter, interlace + CRC
                }
                else if (type == "IDAT")
                {
                    byte[] idatData = reader.ReadBytes(length);
                    idatStream.Write(idatData, 0, idatData.Length);
                    reader.ReadBytes(4); // CRC
                }
                else if (type == "IEND")
                {
                    break;
                }
                else
                {
                    reader.ReadBytes(length + 4);
                }
            }

            idatStream.Position = 0;
            using var deflateStream = new System.IO.Compression.DeflateStream(
                new MemoryStream(idatStream.ToArray(), 2, (int)idatStream.Length - 2),
                System.IO.Compression.CompressionMode.Decompress);

            int channels = colorType switch { 0 => 1, 2 => 3, 4 => 2, 6 => 4, _ => 3 };
            int bytesPerPixel = channels * (bitDepth / 8);
            int stride = width * bytesPerPixel;
            byte[] rawPixels = new byte[height * stride];
            byte[] prevRow = new byte[stride];

            for (int y = 0; y < height; y++)
            {
                int filterByte = deflateStream.ReadByte();
                byte[] row = new byte[stride];
                int read = 0;
                while (read < stride)
                {
                    int n = deflateStream.Read(row, read, stride - read);
                    if (n <= 0) break;
                    read += n;
                }

                for (int x = 0; x < stride; x++)
                {
                    byte a = x >= bytesPerPixel ? row[x - bytesPerPixel] : (byte)0;
                    byte b = prevRow[x];
                    byte c = x >= bytesPerPixel ? prevRow[x - bytesPerPixel] : (byte)0;

                    row[x] = filterByte switch
                    {
                        0 => row[x],
                        1 => (byte)(row[x] + a),
                        2 => (byte)(row[x] + b),
                        3 => (byte)(row[x] + (a + b) / 2),
                        4 => (byte)(row[x] + PaethPredictor(a, b, c)),
                        _ => row[x],
                    };
                }

                Buffer.BlockCopy(row, 0, rawPixels, y * stride, stride);
                Buffer.BlockCopy(row, 0, prevRow, 0, stride);
            }

            byte[] rgba = new byte[width * height * 4];
            for (int i = 0; i < width * height; i++)
            {
                switch (colorType)
                {
                    case 2: // RGB
                        rgba[i * 4] = rawPixels[i * 3];
                        rgba[i * 4 + 1] = rawPixels[i * 3 + 1];
                        rgba[i * 4 + 2] = rawPixels[i * 3 + 2];
                        rgba[i * 4 + 3] = 255;
                        break;
                    case 6: // RGBA
                        rgba[i * 4] = rawPixels[i * 4];
                        rgba[i * 4 + 1] = rawPixels[i * 4 + 1];
                        rgba[i * 4 + 2] = rawPixels[i * 4 + 2];
                        rgba[i * 4 + 3] = rawPixels[i * 4 + 3];
                        break;
                    case 0: // Grayscale
                        rgba[i * 4] = rgba[i * 4 + 1] = rgba[i * 4 + 2] = rawPixels[i];
                        rgba[i * 4 + 3] = 255;
                        break;
                    case 4: // Grayscale + Alpha
                        rgba[i * 4] = rgba[i * 4 + 1] = rgba[i * 4 + 2] = rawPixels[i * 2];
                        rgba[i * 4 + 3] = rawPixels[i * 2 + 1];
                        break;
                }
            }

            return rgba;
        }

        private static (int width, int height) ReadPngDimensions(byte[] data)
        {
            if (data.Length < 24 || !IsPng(data))
                throw new InvalidDataException("Not a PNG file");

            int width = (data[16] << 24) | (data[17] << 16) | (data[18] << 8) | data[19];
            int height = (data[20] << 24) | (data[21] << 16) | (data[22] << 8) | data[23];
            return (width, height);
        }

        private static byte PaethPredictor(byte a, byte b, byte c)
        {
            int p = a + b - c;
            int pa = Math.Abs(p - a);
            int pb = Math.Abs(p - b);
            int pc = Math.Abs(p - c);
            return pa <= pb && pa <= pc ? a : pb <= pc ? b : c;
        }

        private static int ReadBigEndianInt32(BinaryReader reader)
        {
            byte[] bytes = reader.ReadBytes(4);
            return (bytes[0] << 24) | (bytes[1] << 16) | (bytes[2] << 8) | bytes[3];
        }

        private static byte[] DecodeJPEG(byte[] data, out int width, out int height)
        {
            try
            {
                ImageResult decoded = ImageResult.FromMemory(data, ColorComponents.RedGreenBlueAlpha);
                width = decoded.Width;
                height = decoded.Height;
                return decoded.Data;
            }
            catch (Exception ex)
            {
                throw new InvalidDataException("Failed to decode JPEG image.", ex);
            }
        }

        // HEIC/HEIF decoding is delegated to ImageMagick, which bundles libheif in the
        // Magick.NET-Q8-AnyCPU package. The Q8 build uses byte-sized quanta which
        // matches the rest of this pipeline (8-bit RGBA), so Magick.NET will tonemap
        // 10/12-bit HEIC sources down to 8-bit automatically.
        private static byte[] DecodeHEIC(byte[] data, out int width, out int height)
        {
            try
            {
                using var image = new MagickImage(data);
                if (image.ColorSpace != ColorSpace.sRGB)
                    image.ColorSpace = ColorSpace.sRGB;
                if (!image.HasAlpha)
                    image.Alpha(AlphaOption.Set);

                width = (int)image.Width;
                height = (int)image.Height;

                using var pixels = image.GetPixels();
                byte[] rgba = pixels.ToByteArray(0, 0, image.Width, image.Height, "RGBA");
                if (rgba == null || rgba.Length != width * height * 4)
                    throw new InvalidDataException("HEIC decoder returned an unexpected pixel buffer.");
                return rgba;
            }
            catch (Exception ex) when (ex is not InvalidDataException)
            {
                throw new InvalidDataException(
                    "Failed to decode HEIC/HEIF image. Ensure the Magick.NET native binaries with libheif support are available.",
                    ex);
            }
        }

        private static (int width, int height) ReadHeicDimensions(byte[] data)
        {
            try
            {
                var info = new MagickImageInfo(data);
                return ((int)info.Width, (int)info.Height);
            }
            catch (Exception ex)
            {
                throw new InvalidDataException("Failed to read HEIC/HEIF image dimensions.", ex);
            }
        }

        internal static byte[] CompositeOverWhite(byte[] rgba, int width, int height)
        {
            byte[] result = new byte[width * height * 4];
            Parallel.For(0, height, y =>
            {
                int srcRow = y * width * 4;
                for (int x = 0; x < width; x++)
                {
                    int pixBase = srcRow + x * 4;
                    int a = rgba[pixBase + 3];
                    if (a == 255)
                    {
                        result[pixBase] = rgba[pixBase];
                        result[pixBase + 1] = rgba[pixBase + 1];
                        result[pixBase + 2] = rgba[pixBase + 2];
                    }
                    else
                    {
                        float alpha = a / 255f;
                        result[pixBase] = (byte)(rgba[pixBase] * alpha + 255 * (1 - alpha));
                        result[pixBase + 1] = (byte)(rgba[pixBase + 1] * alpha + 255 * (1 - alpha));
                        result[pixBase + 2] = (byte)(rgba[pixBase + 2] * alpha + 255 * (1 - alpha));
                    }

                    result[pixBase + 3] = 255;
                }
            });
            return result;
        }

        internal static byte[] BilinearResize(byte[] rgba, int srcW, int srcH, int dstW, int dstH)
        {
            byte[] result = new byte[dstW * dstH * 4];
            double xRatio = (double)srcW / dstW;
            double yRatio = (double)srcH / dstH;

            Parallel.For(0, dstH, dy =>
            {
                double srcY = (dy + 0.5) * yRatio - 0.5;
                int y0 = Math.Max(0, (int)srcY);
                int y1 = Math.Min(srcH - 1, y0 + 1);
                double fy = srcY - y0;

                for (int dx = 0; dx < dstW; dx++)
                {
                    double srcX = (dx + 0.5) * xRatio - 0.5;
                    int x0 = Math.Max(0, (int)srcX);
                    int x1 = Math.Min(srcW - 1, x0 + 1);
                    double fx = srcX - x0;

                    for (int c = 0; c < 3; c++)
                    {
                        double v00 = rgba[(y0 * srcW + x0) * 4 + c];
                        double v01 = rgba[(y0 * srcW + x1) * 4 + c];
                        double v10 = rgba[(y1 * srcW + x0) * 4 + c];
                        double v11 = rgba[(y1 * srcW + x1) * 4 + c];

                        double v = v00 * (1 - fx) * (1 - fy) + v01 * fx * (1 - fy) +
                                   v10 * (1 - fx) * fy + v11 * fx * fy;
                        result[(dy * dstW + dx) * 4 + c] = (byte)Math.Clamp(v + 0.5, 0, 255);
                    }
                    result[(dy * dstW + dx) * 4 + 3] = 255;
                }
            });

            return result;
        }

        internal static float[] ResizeRgbaToChannelFirstNormalized(byte[] rgba, int srcW, int srcH, int dstW, int dstH)
        {
            int pixels = dstW * dstH;
            float[] result = new float[3 * pixels];
            double xRatio = (double)srcW / dstW;
            double yRatio = (double)srcH / dstH;

            Parallel.For(0, dstH, dy =>
            {
                double srcY = (dy + 0.5) * yRatio - 0.5;
                int y0 = Math.Max(0, (int)srcY);
                int y1 = Math.Min(srcH - 1, y0 + 1);
                double fy = srcY - y0;

                for (int dx = 0; dx < dstW; dx++)
                {
                    double srcX = (dx + 0.5) * xRatio - 0.5;
                    int x0 = Math.Max(0, (int)srcX);
                    int x1 = Math.Min(srcW - 1, x0 + 1);
                    double fx = srcX - x0;

                    int dstIdx = dy * dstW + dx;
                    result[dstIdx] = BilinearSampleNormalized(rgba, srcW, x0, y0, x1, y1, fx, fy, 0);
                    result[pixels + dstIdx] = BilinearSampleNormalized(rgba, srcW, x0, y0, x1, y1, fx, fy, 1);
                    result[2 * pixels + dstIdx] = BilinearSampleNormalized(rgba, srcW, x0, y0, x1, y1, fx, fy, 2);
                }
            });

            return result;
        }

        private static float BilinearSampleNormalized(byte[] rgba, int srcW, int x0, int y0, int x1, int y1,
            double fx, double fy, int channel)
        {
            float v00 = CompositeChannelToNormalized(rgba, (y0 * srcW + x0) * 4, channel);
            float v01 = CompositeChannelToNormalized(rgba, (y0 * srcW + x1) * 4, channel);
            float v10 = CompositeChannelToNormalized(rgba, (y1 * srcW + x0) * 4, channel);
            float v11 = CompositeChannelToNormalized(rgba, (y1 * srcW + x1) * 4, channel);

            double v = v00 * (1 - fx) * (1 - fy) + v01 * fx * (1 - fy) +
                       v10 * (1 - fx) * fy + v11 * fx * fy;
            return Math.Clamp((float)v, -1f, 1f);
        }

        private static float CompositeChannelToNormalized(byte[] rgba, int pixelBase, int channel)
        {
            float alpha = rgba[pixelBase + 3] / 255f;
            float composited = rgba[pixelBase + channel] * alpha + 255f * (1f - alpha);
            return composited / 255f * 2f - 1f;
        }

        /// <summary>
        /// Pack RGBA pixels into channel-first float format [R..., G..., B...] normalized with mean/std.
        /// Matches Ollama's pack(): channel-first with (pixel/255 - mean) / std.
        /// </summary>
        private float[] PackChannelFirst(byte[] rgba, int width, int height)
        {
            int pixels = width * height;
            float[] result = new float[3 * pixels];

            for (int y = 0; y < height; y++)
            {
                for (int x = 0; x < width; x++)
                {
                    int pixIdx = y * width + x;
                    float r = (rgba[pixIdx * 4] / 255f - ImageMean) / ImageStd;
                    float g = (rgba[pixIdx * 4 + 1] / 255f - ImageMean) / ImageStd;
                    float b = (rgba[pixIdx * 4 + 2] / 255f - ImageMean) / ImageStd;

                    result[pixIdx] = r;
                    result[pixels + pixIdx] = g;
                    result[2 * pixels + pixIdx] = b;
                }
            }

            return result;
        }
    }
}