benscode-StitcherApi/Services/Utilities/StreamingProcessor.cs

using System.Collections.Concurrent;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.Formats.Png;
using SixLabors.ImageSharp.PixelFormats;
using SixLabors.ImageSharp.Processing;
using StitcherApi.Services.Utilities;

namespace StitcherApi.Services.Streaming;

public class StreamingProcessor
{
    private readonly string _assetPath;
    private readonly ILogger<StreamingProcessor> _logger;
    private const int TILE_SIZE = 720;
    private const int BAND_HEIGHT = TILE_SIZE;

    public StreamingProcessor(string assetPath, ILogger<StreamingProcessor> logger)
    {
        _assetPath = assetPath;
        _logger = logger;
    }

    public async Task<byte[]> ProcessAsync(StitchRequest request)
    {
        _logger.LogDebug(
            "Starting streaming processor for crop size {W}x{H}",
            request.CropW,
            request.CropH
        );

        ConcurrentDictionary<(int, int), Image<Rgba32>> tileCache = await LoadRequiredTilesAsync(
            request
        );

        using MemoryStream memoryStream = new MemoryStream();
        PngEncoder pngEncoder = new PngEncoder { CompressionLevel = PngCompressionLevel.BestSpeed };

        for (int bandY = 0; bandY < request.CropH; bandY += BAND_HEIGHT)
        {
            int currentBandHeight = Math.Min(BAND_HEIGHT, request.CropH - bandY);
            using Image<Rgba32> bandImage = new Image<Rgba32>(request.CropW, currentBandHeight);
            ComposeBandOptimized(bandImage, bandY, request, tileCache);
            await bandImage.SaveAsPngAsync(memoryStream, pngEncoder);
        }

        return memoryStream.ToArray();
    }

    private async Task<ConcurrentDictionary<(int, int), Image<Rgba32>>> LoadRequiredTilesAsync(
        StitchRequest r
    )
    {
        ConcurrentDictionary<(int, int), Image<Rgba32>> cache = new();
        List<Task> loadTasks = new();

        for (int row = r.StartTileRow; row <= r.EndTileRow; row++)
        {
            for (int col = r.StartTileCol; col <= r.EndTileCol; col++)
            {
                int tileRow = row;
                int tileCol = col;
                loadTasks.Add(
                    Task.Run(async () =>
                    {
                        string tileFilePath = Path.Combine(
                            _assetPath,
                            TileHelper.GetTileFileName(tileRow, tileCol)
                        );
                        if (!File.Exists(tileFilePath))
                            throw new FileNotFoundException($"Asset not found: {tileFilePath}");
                        cache[(tileRow, tileCol)] = await Image.LoadAsync<Rgba32>(tileFilePath);
                    })
                );
            }
        }
        await Task.WhenAll(loadTasks);
        return cache;
    }

    private void ComposeBandOptimized(
        Image<Rgba32> bandImage,
        int bandY,
        StitchRequest r,
        ConcurrentDictionary<(int, int), Image<Rgba32>> tileCache
    )
    {
        int bandAbsoluteY = r.CropY + bandY;
        int currentTileRow = r.MinRow + (bandAbsoluteY / TILE_SIZE);

        for (int tileCol = r.StartTileCol; tileCol <= r.EndTileCol; tileCol++)
        {
            if (tileCache.TryGetValue((currentTileRow, tileCol), out Image<Rgba32>? tileImage))
            {
                int tileOriginX = (tileCol - r.MinCol) * TILE_SIZE;
                int tileOriginY = (currentTileRow - r.MinRow) * TILE_SIZE;

                Rectangle bandRect = new Rectangle(
                    r.CropX,
                    bandAbsoluteY,
                    r.CropW,
                    bandImage.Height
                );
                Rectangle tileRect = new Rectangle(tileOriginX, tileOriginY, TILE_SIZE, TILE_SIZE);
                Rectangle intersection = Rectangle.Intersect(tileRect, bandRect);

                if (!intersection.IsEmpty)
                {
                    Point sourcePoint = new Point(
                        intersection.X - tileOriginX,
                        intersection.Y - tileOriginY
                    );
                    Point destPoint = new Point(
                        intersection.X - r.CropX,
                        intersection.Y - bandAbsoluteY
                    );
                    Rectangle cropRectangle = new Rectangle(sourcePoint, intersection.Size);

                    bandImage.Mutate(ctx =>
                        ctx.DrawImage(tileImage, destPoint, cropRectangle, new GraphicsOptions())
                    );
                }
            }
        }
    }
}