OnscreenTranslate/src/OnscreenTranslator/DirectXOverlay.cs

using System.Diagnostics;
using System.Drawing;
using System.IO;
using System.Net;
using System.Text.RegularExpressions;
using LibreTranslate.Net.Enhanced.Models;
using Overlay.NET.Common;
using Overlay.NET.Directx;
using Process.NET.Windows;
using TesseractOCR;
using Image = TesseractOCR.Pix.Image;
using ImageFormat = System.Drawing.Imaging.ImageFormat;
using Point = System.Drawing.Point;
using Rectangle = System.Drawing.Rectangle;

namespace OnscreenTranslator;

public class DirectXOverlay : DirectXOverlayPlugin
{
    private const string TessdataPath = "C:/Users/luis-/Downloads/tessdata-4.1.0/tessdata-4.1.0";

    private static readonly Regex LineRegex =
        new(@"class=['""]ocr_(?:line|textline|caption|header)['""][^>]*?title=['""]bbox (\d+) (\d+) (\d+) (\d+)");

    private static readonly Regex WordRegex = new(@"class=['""]ocrx_word['""][^>]*>(.*?)</span>", RegexOptions.Singleline);
    private static readonly Regex TagRegex = new("<[^>]+>");

    private readonly TickEngine _tickEngine = new();
    public readonly ISettings<OverlaySettings> Settings = new SerializableSettings<OverlaySettings>();
    private readonly LibreTranslate.Net.Enhanced.LibreTranslate _libreTranslate = new("http://127.0.0.1:5000", "");
    private int _backgroundBrush;
    private int _displayFps;
    private Engine? _engine;
    private TranslationLanguage? _engineLanguage;
    private bool _excludedFromCapture;
    private int _font;
    private int _hugeFont;
    private int _i;
    private int _interiorBrush;
    private int _redBrush;
    private int _redOpacityBrush;
    private float _rotation;
    private int _textBrush;
    private bool _translating;
    private Stopwatch _watch;

    // Replaced as a whole when a translation finishes and never mutated, so rendering can't see a half-built list
    private IReadOnlyList<Element> _translatedElements = [];

    public TranslationLanguage SourceLanguage { get; set; } = TranslationLanguage.Japanese;
    public TranslationLanguage TargetLanguage { get; set; } = TranslationLanguage.German;

    /// <summary>
    ///     Skip LibreTranslate and show Tesseract's text as-is, for checking OCR results.
    ///     Applies to the next TriggerTranslation.
    /// </summary>
    public bool ShowOcrText { get; set; }

    public override void Initialize(IWindow targetWindow)
    {
        // Set target window by calling the base method
        base.Initialize(targetWindow);

        // For demo, show how to use settings
        var current = Settings.Current;
        var type = GetType();

        if (current.UpdateRate == 0)
            current.UpdateRate = 1000 / 60;

        current.Author = GetAuthor(type);
        current.Description = GetDescription(type);
        current.Identifier = GetIdentifier(type);
        current.Name = GetName(type);
        current.Version = GetVersion(type);

        // File is made from above info
        Settings.Save();
        Settings.Load();
        //Console.Title = @"OverlayExample";

        OverlayWindow = new DirectXOverlayWindow(targetWindow.Handle, false);
        // Keep the overlay out of the screenshots fed to Tesseract
        _excludedFromCapture = OverlayWindow.ExcludeFromCapture();
        _watch = Stopwatch.StartNew();

        _redBrush = OverlayWindow.Graphics.CreateBrush(0x7FFF0000);
        _redOpacityBrush = OverlayWindow.Graphics.CreateBrush(Color.FromArgb(80, 255, 0, 0));
        _interiorBrush = OverlayWindow.Graphics.CreateBrush(0x7FFFFF00);
        _backgroundBrush = OverlayWindow.Graphics.CreateBrush(Color.FromArgb(200, 0, 0, 0));
        _textBrush = OverlayWindow.Graphics.CreateBrush(Color.FromArgb(255, 255, 255, 255));

        _font = OverlayWindow.Graphics.CreateFont("Arial", 20);
        _hugeFont = OverlayWindow.Graphics.CreateFont("Arial", 50, true);

        _rotation = 0.0f;
        _displayFps = 0;
        _i = 0;
        // Set up update interval and register events for the tick engine.

        _tickEngine.PreTick += OnPreTick;
        _tickEngine.Tick += OnTick;
    }

    private void OnPreTick(object sender, EventArgs e)
    {
        var targetWindowIsActivated = TargetWindow.IsActivated;
        if (!targetWindowIsActivated && OverlayWindow.IsVisible)
        {
            _watch.Stop();
            ClearScreen();
            OverlayWindow.Hide();
        }
        else if (targetWindowIsActivated && !OverlayWindow.IsVisible)
        {
            OverlayWindow.Show();
        }
    }

    private void OnTick(object sender, EventArgs e)
    {
        if (!OverlayWindow.IsVisible) return;

        OverlayWindow.Update();
        InternalRender();
    }

    public override void Enable()
    {
        _tickEngine.Interval = Settings.Current.UpdateRate.Milliseconds();
        _tickEngine.IsTicking = true;
        base.Enable();
    }

    public override void Disable()
    {
        _tickEngine.IsTicking = false;
        base.Disable();
    }

    public override void Update()
    {
        _tickEngine.Pulse();
    }

    protected void InternalRender()
    {
        if (!_watch.IsRunning) _watch.Start();

        OverlayWindow.Graphics.BeginScene();
        OverlayWindow.Graphics.ClearScene();

        RenderElements(_translatedElements);

        if (_watch.ElapsedMilliseconds > 1000)
        {
            _i = _displayFps;
            _displayFps = 0;
            _watch.Restart();
        }

        else
        {
            _displayFps++;
        }

        OverlayWindow.Graphics.EndScene();
    }

    /// <summary>
    ///     Captures the target window, reads its text and replaces the overlay with the translation.
    ///     Must be called from the UI thread; calls made while a translation is running are ignored.
    /// </summary>
    public async Task TriggerTranslation()
    {
        if (_translating || OverlayWindow == null) return;
        _translating = true;

        try
        {
            using var screenshot = await CaptureTargetWindow();
            var source = SourceLanguage;
            var target = TargetLanguage;

            // OCR takes a while, keep it off the UI thread so rendering continues
            var html = await Task.Run(() => GetScreenHTML(screenshot, source));
            var untranslatedElements = GetScreenElements(html, source);
            var translatedElements = ShowOcrText
                ? ShowUntranslated(untranslatedElements)
                : await TranslateElements(untranslatedElements, source, target);

            if (OverlayWindow == null) return; // disposed while translating

            foreach (var element in translatedElements)
            {
                // Cover the original line, growing the box if the translation is longer
                var size = OverlayWindow.Graphics.MeasureText(element.Output, _font);
                element.BoxWidth = Math.Max(element.Width, (int)Math.Ceiling(size.Width));
                element.BoxHeight = Math.Max(element.Height, (int)Math.Ceiling(size.Height));
            }

            _translatedElements = translatedElements;
            // DrawText caches a layout per distinct string; drop the ones from the previous translation
            OverlayWindow.Graphics.DeleteLayoutContainer();
        }
        finally
        {
            _translating = false;
        }
    }

    private string GetScreenHTML(MemoryStream screenshot, TranslationLanguage language)
    {
        // Loading the models is slow, so the engine is kept until the source language changes
        if (_engine == null || _engineLanguage != language)
        {
            _engine?.Dispose();
            _engine = new Engine(TessdataPath, language.OcrLanguages.ToList());
            _engineLanguage = language;
        }

        using var img = Image.LoadFromMemory(screenshot);
        using var page = _engine.Process(img);
        return page.HOcrText();
    }

    /// <summary>
    ///     Turns hOCR into one element per text line, positioned at the line's bounding box.
    /// </summary>
    private static List<Element> GetScreenElements(string html, TranslationLanguage language)
    {
        var elements = new List<Element>();
        var lines = LineRegex.Matches(html);

        for (var i = 0; i < lines.Count; i++)
        {
            var line = lines[i];
            var end = i + 1 < lines.Count ? lines[i + 1].Index : html.Length;
            var lineHtml = html.Substring(line.Index, end - line.Index);

            // Word text may be wrapped in <strong>/<em> and contain entities like &amp;
            var words = WordRegex.Matches(lineHtml)
                .Select(word => WebUtility.HtmlDecode(TagRegex.Replace(word.Groups[1].Value, "")).Trim())
                .Where(word => word.Length > 0);
            var text = string.Join(language.JoinWordsWithSpace ? " " : "", words);
            if (text.Length == 0) continue;

            var x0 = int.Parse(line.Groups[1].Value);
            var y0 = int.Parse(line.Groups[2].Value);
            elements.Add(new Element
            {
                X = x0,
                Y = y0,
                Width = int.Parse(line.Groups[3].Value) - x0,
                Height = int.Parse(line.Groups[4].Value) - y0,
                Input = text
            });
        }

        return elements;
    }

    private async Task<List<Element>> TranslateElements(List<Element> elements, TranslationLanguage source,
        TranslationLanguage target)
    {
        var translated = new List<Element>();

        foreach (var element in elements)
        {
            var targetText = await _libreTranslate.TranslateAsync(new Translate
            {
                Source = source.Code,
                Target = target.Code,
                Text = element.Input
            });
            if (string.IsNullOrWhiteSpace(targetText)) continue;

            element.Output = targetText;
            translated.Add(element);
        }

        return translated;
    }

    private static List<Element> ShowUntranslated(List<Element> elements)
    {
        foreach (var element in elements) element.Output = element.Input;
        return elements;
    }

    private void RenderElements(IReadOnlyList<Element> elements)
    {
        const int padding = 4;

        foreach (var element in elements)
        {
            OverlayWindow.Graphics.FillRectangle(element.X - padding, element.Y - padding,
                element.BoxWidth + padding * 2, element.BoxHeight + padding * 2, _backgroundBrush);
            OverlayWindow.Graphics.DrawText(element.Output, _font, _textBrush, element.X, element.Y);
        }
    }

    public override void Dispose()
    {
        _tickEngine.IsTicking = false;
        _tickEngine.PreTick -= OnPreTick;
        _tickEngine.Tick -= OnTick;

        OverlayWindow?.Dispose();
        OverlayWindow = null;

        _engine?.Dispose();
        _engine = null;

        base.Dispose();
    }

    public void ClearScreen()
    {
        OverlayWindow.Graphics.BeginScene();
        OverlayWindow.Graphics.ClearScene();
        OverlayWindow.Graphics.EndScene();
    }

    private async Task<MemoryStream> CaptureTargetWindow()
    {
        // Without capture exclusion (Windows before 10 2004) the overlay would end up in its own screenshot
        var hideOverlay = !_excludedFromCapture && OverlayWindow.IsVisible;
        if (hideOverlay)
        {
            OverlayWindow.Hide();
            await Task.Delay(50); // let DWM compose a frame without the overlay
        }

        try
        {
            // Use the overlay's bounds (GetWindowRect of the target), not TargetWindow.X/Y/Width/Height:
            // Process.NET reports the restored placement rect, which is wrong for snapped or maximized windows
            OverlayWindow.Update();
            var bounds = new Rectangle(OverlayWindow.X,
                OverlayWindow.Y,
                OverlayWindow.Width,
                OverlayWindow.Height);

            using var bitmap = new Bitmap(bounds.Width, bounds.Height);
            using (var g = Graphics.FromImage(bitmap))
            {
                g.CopyFromScreen(new Point(bounds.Left, bounds.Top), Point.Empty, bounds.Size);
            }

            // bitmap.Save("C://test.jpg", ImageFormat.Jpeg);
            var mem = new MemoryStream();
            bitmap.Save(mem, ImageFormat.Bmp);
            mem.Position = 0;
            return mem;
        }
        finally
        {
            if (hideOverlay) OverlayWindow.Show();
        }
    }

    private class Element
    {
        public int X;
        public int Y;
        public int Width;
        public int Height;
        public int BoxWidth;
        public int BoxHeight;
        public string Input = "";
        public string Output = "";
    }
}