xSaherElmAiApi/Providers/Extractors/XImageFileContentExtractor.cs

using System;
using System.IO;
using System.Linq;
using System.Threading;
using xAiModels.Models;
using System.Threading.Tasks;
using xAiApi.Interfaces.Extractors;
using xAiApi.Configurations;
using xCommons.Extensions;
using xExceptions.Constants;
using Tesseract;

namespace xAiApi.Providers.Extractors
{
    public class XImageFileContentExtractor : IXImageFileContentExtractor
    {
        /// <summary>
        /// Supported MIME Types ...
        /// </summary>
        private static readonly string[] SupportedMimeTypes =
        [
            "image/png",
            "image/jpeg",
            "image/jpg",
            "image/gif",
            "image/webp",
            "image/bmp"
        ];

        /// <summary>
        /// AI Api Configuration ...
        /// </summary>
        private readonly XAiApiConfiguration configuration;

        public XImageFileContentExtractor(
            XAiApiConfiguration configuration
        )
        {
            //
            if (configuration.IsNullOrDefault() ||
                configuration.OCR.IsNullOrDefault() ||
                configuration.OCR.DataPath.IsNullOrEmpty() ||
                configuration.OCR.DefaultLanguage.IsNullOrEmpty())
            {
                XException.InvalidConfiguration.Throw();
            }

            //
            this.configuration = configuration;
        }

        /// <summary>
        /// Check if this extractor supports the specified MIME type ...
        /// </summary>
        public bool CanExtract(string mimeType)
        {
            //
            return SupportedMimeTypes.Contains(
                mimeType?.ToLowerInvariant() ?? string.Empty
            );
        }

        /// <summary>
        /// Extract text content from file stream ...
        /// </summary>
        public async Task<string> ExtractAsync(
            Stream fileStream,
            string mimeType,
            CancellationToken cancellationToken = default
        )
        {
            //
            var result = await ExtractRichAsync(
                fileStream,
                "image",
                mimeType,
                cancellationToken
            );

            //
            return result.Text;
        }

        /// <summary>
        /// Extract content from stream as Rich Result ...
        /// </summary>
        /// <param name="fileStream"></param>
        /// <param name="fileName"></param>
        /// <param name="mimeType"></param>
        /// <param name="cancellationToken"></param>
        /// <returns></returns>
        public async Task<XFileExtractionResult> ExtractRichAsync(
            Stream fileStream,
            string fileName,
            string mimeType,
            CancellationToken cancellationToken = default
        )
        {
            //
            var result = new XFileExtractionResult
            {
                FileName = fileName,
                MimeType = mimeType
            };

            //
            using var memoryStream = new MemoryStream();
            await fileStream.CopyToAsync(memoryStream, cancellationToken);
            var imageBytes = memoryStream.ToArray();

            //
            result.Images.Add(new XExtractedImage
            {
                Bytes = imageBytes,
                MimeType = mimeType,
                Description = $"Attached image: {fileName}"
            });

            //
            // OCR Fallback ...
            if (configuration.OCR.EnableOcrFallback)
            {
                //
                try
                {
                    //
                    var ocrText = await PerformOcrAsync(imageBytes, cancellationToken);
                    if (!string.IsNullOrWhiteSpace(ocrText))
                    {
                        result.Text = ocrText;
                    }
                }
                catch
                {
                    // OCR failed, but image is still available for Vision
                }
            }

            //
            return result;
        }

        /// <summary>
        /// Do OCR on Image for Text Extraction ...
        /// </summary>
        private async Task<string> PerformOcrAsync(
            byte[] imageBytes,
            CancellationToken cancellationToken
        )
        {
            //
            return await Task.Run(() =>
            {
                //
                // Validate Path ...
                if (!Directory.Exists(configuration.OCR.DataPath))
                {
                    //
                    throw new DirectoryNotFoundException(
                        $"TessData directory not found at: {configuration.OCR.DataPath}. " +
                        "Please download traineddata files from https://github.com/tesseract-ocr/tessdata"
                    );
                }
            
                //
                using var engine = new TesseractEngine(
                    datapath: configuration.OCR.DataPath,
                    language: configuration.OCR.DefaultLanguage,
                    engineMode: configuration.OCR.EngineMode
                );

                //
                using var pix = Pix.LoadFromMemory(imageBytes);

                //
                cancellationToken.ThrowIfCancellationRequested();

                //
                using var page = engine.Process(pix);

                //
                var text = page.GetText();

                //
                text = text?.Trim();

                //
                return text ?? string.Empty;
            }, cancellationToken);
        }
    }
}