Iceshrimp.NET/Iceshrimp.Backend/Core/Services/ImageProcessing/VipsProcessor.cs

using System.Runtime.InteropServices;
using System.Security;
using CommunityToolkit.HighPerformance;
using Iceshrimp.Backend.Core.Helpers;
using NetVips;
using SixLabors.ImageSharp.PixelFormats;

namespace Iceshrimp.Backend.Core.Services.ImageProcessing;

public class VipsProcessor : ImageProcessorBase, IImageProcessor
{
	private readonly ILogger<VipsProcessor> _logger;

	// Set to false until https://github.com/libvips/libvips/issues/2537 is implemented
	public bool CanIdentify         => false;
	public bool CanGenerateBlurhash => true;

	public VipsProcessor(ILogger<VipsProcessor> logger) : base("LibVips", 0)
	{
		_logger = logger;

		//TODO: Implement something similar to https://github.com/lovell/sharp/blob/da655a1859744deec9f558effa5c9981ef5fd6d3/lib/utility.js#L153C5-L158
		NetVips.NetVips.Concurrency = 1;

		// We want to know when we have a memory leak
		NetVips.NetVips.Leak = true;

		// We don't need the VIPS operation or file cache
		Cache.Max      = 0;
		Cache.MaxFiles = 0;
		Cache.MaxMem   = 0;

		Log.SetLogHandler("VIPS", Enums.LogLevelFlags.Warning | Enums.LogLevelFlags.Error,
		                  VipsLogDelegate);
	}

	public bool CanEncode(ImageFormat format)
	{
		return format switch
		{
			ImageFormat.Webp => true,
			ImageFormat.Jxl  => true,
			ImageFormat.Avif => true,
			_                => throw new ArgumentOutOfRangeException(nameof(format), format, null)
		};
	}

	public Stream Encode(byte[] input, IImageInfo _, ImageFormat format)
	{
		return format switch
		{
			ImageFormat.Webp opts => EncodeWebp(input, opts),
			ImageFormat.Jxl opts  => EncodeJxl(input, opts),
			ImageFormat.Avif opts => EncodeAvif(input, opts),
			_                     => throw new ArgumentOutOfRangeException(nameof(format))
		};
	}

	public string Blurhash(byte[] buf, IImageInfo ident)
	{
		using var blurhashImageSource =
			Image.ThumbnailBuffer(buf, 100, height: 100, size: Enums.Size.Down);
		using var blurhashImage = blurhashImageSource.Interpretation == Enums.Interpretation.Srgb
			? blurhashImageSource
			: blurhashImageSource.Colourspace(Enums.Interpretation.Srgb);
		using var blurhashImageFlattened = blurhashImage.HasAlpha() ? blurhashImage.Flatten() : blurhashImage;
		using var blurhashImageActual    = blurhashImageFlattened.Cast(Enums.BandFormat.Uchar);

		var blurBuf    = blurhashImageActual.WriteToMemory();
		var blurPixels = MemoryMarshal.Cast<byte, Rgb24>(blurBuf).AsSpan2D(blurhashImage.Height, blurhashImage.Width);
		return BlurhashHelper.Encode(blurPixels, 7, 7);
	}

	public IImageInfo Identify(byte[] input) => new VipsImageInfo(Image.NewFromBuffer(input));

	private static MemoryStream EncodeWebp(byte[] buf, ImageFormat.Webp opts)
	{
		using var image  = Thumbnail(buf, opts.TargetRes);
		var       stream = new MemoryStream();
		image.WebpsaveStream(stream, opts.Quality, opts.Mode == ImageFormat.Webp.Compression.Lossless,
		                     nearLossless: opts.Mode == ImageFormat.Webp.Compression.NearLossless);
		return stream;
	}

	private static MemoryStream EncodeAvif(byte[] buf, ImageFormat.Avif opts)
	{
		using var image  = Thumbnail(buf, opts.TargetRes);
		var       stream = new MemoryStream();
		image.HeifsaveStream(stream, opts.Quality, lossless: opts.Mode == ImageFormat.Avif.Compression.Lossless,
		                     bitdepth: opts.BitDepth,
		                     compression: Enums.ForeignHeifCompression.Av1);
		return stream;
	}

	private static MemoryStream EncodeJxl(byte[] buf, ImageFormat.Jxl opts)
	{
		using var image  = Thumbnail(buf, opts.TargetRes);
		var       stream = new MemoryStream();
		image.JxlsaveStream(stream, q: opts.Quality, lossless: opts.Mode == ImageFormat.Jxl.Compression.Lossless,
		                    effort: opts.Effort);
		return stream;
	}

	private static Image StripMetadata(Image image)
	{
		return image.Mutate(mutable =>
		{
			mutable.Autorot();
			foreach (var field in mutable.GetFields())
			{
				if (field is "icc-profile-data") continue;
				mutable.Remove(field);
			}
		});
	}

	private static Image Thumbnail(byte[] buf, int targetRes)
	{
		using var image = Image.ThumbnailBuffer(buf, targetRes, height: targetRes, size: Enums.Size.Down);
		return StripMetadata(image);
	}

	private void VipsLogDelegate(string domain, Enums.LogLevelFlags _, string message) =>
		_logger.LogWarning("{domain} - {message}", domain, message);

	[SuppressUnmanagedCodeSecurity]
	[DllImport("libvips.42", EntryPoint = "vips_image_get_n_pages", CallingConvention = CallingConvention.Cdecl)]
	private static extern int GetPageCount(Image image);

	private class VipsImageInfo(Image image) : IImageInfo
	{
		public int    Width      => image.Width;
		public int    Height     => image.Height;
		public bool   IsAnimated => GetPageCount(image) > 1;
		public string MimeType   => throw new NotImplementedException(); //TODO

		public static implicit operator VipsImageInfo(Image src) => new(src);
	}
}