From b286a0b0f1fcdb511d2dbb8886039cfb0182c89b Mon Sep 17 00:00:00 2001 From: Jake Mannens Date: Fri, 1 Sep 2023 13:03:57 +1000 Subject: Merged OCR functionality --- Controllers/MediaController.cs | 3 -- HBContext.cs | 1 + Media.cs | 15 +++++- Pages/TagDefinitions.razor | 7 +-- Pages/ViewMedia.razor | 26 ++++++--- Program.cs | 3 +- Server.csproj | 2 + Services/OcrService.cs | 117 +++++++++++++++++++++++++++++++++++++++++ Services/SearchService.cs | 60 ++++++++++++++++++--- Util.cs | 99 ++++++++++++++++++++++++++++++++++ appsettings.Development.json | 3 +- tessdata/eng.traineddata | Bin 0 -> 15400601 bytes wwwroot/styles/global.css | 11 ++++ 13 files changed, 321 insertions(+), 26 deletions(-) create mode 100644 Services/OcrService.cs create mode 100644 tessdata/eng.traineddata diff --git a/Controllers/MediaController.cs b/Controllers/MediaController.cs index 85dfc65..8070199 100644 --- a/Controllers/MediaController.cs +++ b/Controllers/MediaController.cs @@ -1,9 +1,6 @@ using HyperBooru.Services; using HyperBooru.Util; -using ImageMagick; using Microsoft.AspNetCore.Mvc; -using MimeDetective; -using System.Security.Cryptography; namespace HyperBooru.Controllers; diff --git a/HBContext.cs b/HBContext.cs index f6bc15c..15dad6d 100644 --- a/HBContext.cs +++ b/HBContext.cs @@ -15,6 +15,7 @@ public class HBContext : DbContext { public DbSet Tags { get; set; } public DbSet Media { get; set; } public DbSet UploadedFiles { get; set; } + public DbSet OcrData { get; set; } private IConfigService config; diff --git a/Media.cs b/Media.cs index e2598a9..2a4dab6 100644 --- a/Media.cs +++ b/Media.cs @@ -13,6 +13,7 @@ public class Media : HBObject { public string? LongDescription { get; set; } public int Width { get; set; } public int Height { get; set; } + public virtual OcrData? OcrData { get; set; } public virtual List UploadedFiles { get; set; } = new(); public bool IsIngest => Tags @@ -26,7 +27,7 @@ public class Media : HBObject { return UploadedFiles .OrderBy(f => f.UploadTime) - .First()?.Filename; + .First()?.Filename ?? Guid.ToString().ToUpper(); } } } @@ -40,4 +41,16 @@ public class UploadedFile : HBObject { public DateTime? LastWriteTime { get; set; } public DateTime? CreateTime { get; set; } public virtual Media Media { get; set; } +} + +public class OcrData { + [Key] + [DatabaseGenerated(DatabaseGeneratedOption.Identity)] + public int OcrDataId { get; set; } + [ForeignKey("ObjectId")] + public int MediaId { get; set; } + public string Text { get; set; } + public string SearchableText { get; set; } + public DateTime Timestamp { get; set; } + public virtual Media Media { get; set; } } \ No newline at end of file diff --git a/Pages/TagDefinitions.razor b/Pages/TagDefinitions.razor index f5339e7..f32e803 100644 --- a/Pages/TagDefinitions.razor +++ b/Pages/TagDefinitions.razor @@ -43,12 +43,7 @@ @(", ") } } - -@* @(string.Join(", ", tagDef.ImplicitTags - .Where(it => it.Source == TagSource.UserTag) - .Select(it => it.Name) - .Order())) -*@ + PromptToEdit(tagDef))>Edit diff --git a/Pages/ViewMedia.razor b/Pages/ViewMedia.razor index bb6a207..eb49b15 100644 --- a/Pages/ViewMedia.razor +++ b/Pages/ViewMedia.razor @@ -62,17 +62,18 @@ - @if(media.IsIngest) { - - } else { - - } + @if(infoEditMode) { } else { } + @if(media.IsIngest) { + + } else { + + } @@ -85,6 +86,17 @@ + + @if(media.OcrData is null) { +

This media item hasn't been scanned yet!

+ } else { + @media.OcrData?.Text + } + + + + + m.Tags) .ThenInclude(t => t.TagDefinition) .Include(m => m.UploadedFiles) + .Include(m => m.OcrData) .First(m => m.Guid == MediaId); title = media.DisplayName ?? "Media View"; diff --git a/Program.cs b/Program.cs index 90375b0..564ab30 100644 --- a/Program.cs +++ b/Program.cs @@ -14,13 +14,14 @@ public class Program { builder.Services.AddRazorPages(); builder.Services.AddServerSideBlazor(); - // Add out custom services + // Add our custom services builder.Services.AddSingleton(); builder.Services.AddDbContextFactory(); builder.Services.AddScoped(); builder.Services.AddScoped(); builder.Services.AddScoped(); builder.Services.AddSingleton(); + builder.Services.AddHostedService(); var app = builder.Build(); diff --git a/Server.csproj b/Server.csproj index fab6521..c14aa24 100644 --- a/Server.csproj +++ b/Server.csproj @@ -30,6 +30,8 @@ + + diff --git a/Services/OcrService.cs b/Services/OcrService.cs new file mode 100644 index 0000000..2f65e43 --- /dev/null +++ b/Services/OcrService.cs @@ -0,0 +1,117 @@ +using HyperBooru.Util; +using Microsoft.EntityFrameworkCore; +using System.Diagnostics; +using System.Runtime.InteropServices; +using System.Text.RegularExpressions; +using Tesseract; + +namespace HyperBooru.Services; + +public class OcrService : IHostedService { + private readonly TimeSpan ProcessInterval = TimeSpan.FromMinutes(30); + private readonly TimeSpan StartupDelay = TimeSpan.FromSeconds(30); + + private readonly Regex SpaceRegex = new(@"[^0-9a-z]+", RegexOptions.Compiled); + + private Task? task; + private CancellationTokenSource cts = new(); + + private Timer timer; + + private IServiceScopeFactory scopeFactory; + private ILogger logger; + private IDbContextFactory dbFactory; + + public OcrService( + IServiceScopeFactory scopeFactory, + ILogger logger, + IDbContextFactory dbFactory) { + + this.scopeFactory = scopeFactory; + this.logger = logger; + this.dbFactory = dbFactory; + + timer = new((object? state) => { + if(task is not null && !task.IsCompleted) + return; + cts = new(); + task = ProcessAllAsync(cts.Token); + }); + } + + public Task StartAsync(CancellationToken ct) { + logger.LogInformation("Service starting..."); + timer.Change(StartupDelay, ProcessInterval); + return Task.CompletedTask; + } + + public Task StopAsync(CancellationToken ct) { + logger.LogInformation("Service stopping..."); + timer.Change(Timeout.Infinite, Timeout.Infinite); + cts.Cancel(); + return Task.CompletedTask; + } + + async Task ProcessAllAsync(CancellationToken ct) { + using var scope = scopeFactory.CreateScope(); + var mediaService = scope.ServiceProvider + .GetRequiredService(); + + using var db = dbFactory.CreateDbContext(); + Guid[] guids = db.Media + .Include(m => m.OcrData) + .Where(m => m.OcrData == null) + .Where(m => m.MimeType.Contains("image/")) + .Select(m => m.Guid) + .ToArray(); + db.Dispose(); + + logger.LogInformation($"Performing OCR pass on {guids.Count()} media items"); + + var factory = new TaskFactory(new LimitedConcurrencyTaskScheduler()); + var tasks = new List(); + + var stopwatch = Stopwatch.StartNew(); + + foreach(var guid in guids) + tasks.Add(factory.StartNew(() => Process(guid, mediaService), ct)); + + await Task.WhenAll(tasks); + stopwatch.Stop(); + + var time = stopwatch.Elapsed.ToStringHumanReadable(); + logger.LogInformation( + $"Performed OCR pass on {guids.Count()} media items in {time}"); + } + + private void Process(Guid media, IMediaService mediaService) { + logger.LogDebug($"Performing OCR on media item {media}"); + + using var db = dbFactory.CreateDbContext(); + var m = db.Media + .Include(m => m.OcrData) + .First(m => m.Guid == media); + + OcrData o = m.OcrData ?? new(); + + using var engine = new TesseractEngine("tessdata", "eng", EngineMode.Default); + using var image = Pix.LoadFromFile(mediaService.GetPath(m)); + engine.SetVariable("debug_file", NullFile); + + o.Timestamp = DateTime.UtcNow; + o.Text = engine.Process(image).GetText().Trim(); + o.SearchableText = SpaceRegex.Replace(o.Text.ToLower(), " ").Trim(); + + m.OcrData = o; + db.SaveChanges(); + } + + private string NullFile { + get { + if(RuntimeInformation.IsOSPlatform(OSPlatform.Windows)) + return "NUL"; + else + return "/dev/null"; + } + } +} diff --git a/Services/SearchService.cs b/Services/SearchService.cs index e8e497d..bb2963d 100644 --- a/Services/SearchService.cs +++ b/Services/SearchService.cs @@ -24,33 +24,77 @@ public class SearchService : ISearchService { query = query.ToLower(); + int[] descriptionResults = SearchDescription(query); + int[] ocrResults = SearchOcr(query); + var matchedTag = db.TagDefinitions .FirstOrDefault(td => td.Name.ToLower() == query); int[] tags; - if(matchedTag is not null) { tags = tagService .TagsThatImply(matchedTag) .Select(td => td.ObjectId) .ToArray(); } else { - // TODO: expand scope to all tags that imply + // TODO: Expand scope to all tags that imply tags = db.TagDefinitions .Where(td => td.Name.ToLower().Contains(query)) .Select(td => td.ObjectId) .ToArray(); } + int[] tagResults = SearchTags(tags); + + int[] mediaIds = descriptionResults + .Union(ocrResults) + .Union(tagResults) + .OrderDescending() + .ToArray(); + return db.Media .Include(m => m.Tags) - .AsEnumerable() - .Where(m => m.Tags.IntersectBy(tags, t => t.TagDefinitionId).Any()) - .Concat(db.Media + .Where(m => mediaIds.Contains(m.ObjectId)) + .ToArray(); + } + + // TODO: Make asynchronous + private int[] SearchTags(int[] tags) { + return Task.Run(() => { + using var db = dbFactory.CreateDbContext(); + return db.Media + .Include(m => m.Tags) + .AsEnumerable() + .Where(m => m.Tags.IntersectBy(tags, t => t.TagDefinitionId).Any()) + .Select(m => m.ObjectId) + .ToArray(); + }).GetAwaiter().GetResult(); + } + + // TODO: Make asynchronous + private int[] SearchDescription(string query) { + return Task.Run(() => { + using var db = dbFactory.CreateDbContext(); + query = query.ToLower(); + return db.Media .Where(m => (m.ShortDescription != null && m.ShortDescription.ToLower().Contains(query)) || - (m.LongDescription != null && m.LongDescription.ToLower().Contains(query)))) - .DistinctBy(m => m.ObjectId) - .ToArray(); + (m.LongDescription != null && m.LongDescription.ToLower().Contains(query))) + .Select(m => m.ObjectId) + .ToArray(); + }).GetAwaiter().GetResult(); + } + + // TODO: Make asynchronous + private int[] SearchOcr(string query) { + return Task.Run(() => { + using var db = dbFactory.CreateDbContext(); + query = query.ToLower(); + return db.OcrData + .Include(o => o.Media) + .Where(o => o.SearchableText.Contains(query)) + .Select(o => o.Media.ObjectId) + .ToArray(); + }).GetAwaiter().GetResult(); } } diff --git a/Util.cs b/Util.cs index 31a2e84..6af6c81 100644 --- a/Util.cs +++ b/Util.cs @@ -18,4 +18,103 @@ public static class Extensions { double n = x / Math.Pow(10, exp / 3 * 3); return $"{Math.Round(n, 2 - (exp % 3))} {suffix}B"; } + + public static string ToStringHumanReadable(this TimeSpan t) { + if(t.TotalMilliseconds < 1000) + return string.Format("{0:0}ms", t.TotalMilliseconds); + if(t.TotalSeconds < 60) + return string.Format("{0:0.00}s", t.TotalSeconds); + if(t.TotalMinutes < 60) + return string.Format("{0:0}m{0:0}s", t.TotalMinutes, t.Seconds); + if(t.TotalHours < 24) + return string.Format("{0:0}h{0:0}m", t.TotalHours, t.Minutes); + return string.Format("{0:0.00}d", t.TotalDays); + } +} + +public class LimitedConcurrencyTaskScheduler : TaskScheduler { + public sealed override int MaximumConcurrencyLevel => + maxConcurrency; + + private int maxConcurrency; + + [ThreadStatic] + private static bool threadIsProcessingItems; + + private readonly LinkedList tasks = new(); + + private int delegatesQueuedOrRunning = 0; + + public LimitedConcurrencyTaskScheduler() { + maxConcurrency = Environment.ProcessorCount; + } + + public LimitedConcurrencyTaskScheduler(int maxConcurrency) { + if(maxConcurrency < 1) + throw new ArgumentOutOfRangeException("maxConcurrency must be greater than 0"); + this.maxConcurrency = (int) maxConcurrency; + } + + protected sealed override void QueueTask(Task task) { + lock(tasks) { + tasks.AddLast(task); + if(delegatesQueuedOrRunning < maxConcurrency) { + delegatesQueuedOrRunning++; + NotifyThreadPoolOfPendingWork(); + } + } + } + + private void NotifyThreadPoolOfPendingWork() { + ThreadPool.UnsafeQueueUserWorkItem(_ => { + threadIsProcessingItems = true; + try { + while(true) { + Task item; + lock(tasks) { + if(tasks.Count == 0) { + delegatesQueuedOrRunning--; + break; + } else { + item = tasks.First.Value; + tasks.RemoveFirst(); + } + } + TryExecuteTask(item); + } + } finally { + threadIsProcessingItems = false; + } + }, null); + } + + protected sealed override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued) { + if(!threadIsProcessingItems) + return false; + + if(taskWasPreviouslyQueued) + return TryDequeue(task) ? TryExecuteTask(task) : false; + else + return TryExecuteTask(task); + } + + protected sealed override bool TryDequeue(Task task) { + lock(tasks) { + return tasks.Remove(task); + } + } + + protected sealed override IEnumerable GetScheduledTasks() { + bool lockTaken = false; + try { + Monitor.TryEnter(tasks, ref lockTaken); + if(lockTaken) + return tasks; + else + throw new NotSupportedException(); + } finally { + if(lockTaken) + Monitor.Exit(tasks); + } + } } diff --git a/appsettings.Development.json b/appsettings.Development.json index 770d3e9..6860045 100644 --- a/appsettings.Development.json +++ b/appsettings.Development.json @@ -3,7 +3,8 @@ "Logging": { "LogLevel": { "Default": "Information", - "Microsoft.AspNetCore": "Warning" + "Microsoft.AspNetCore": "Warning", + "HyperBooru.Services.OcrService": "Debug" } } } diff --git a/tessdata/eng.traineddata b/tessdata/eng.traineddata new file mode 100644 index 0000000..176dc32 Binary files /dev/null and b/tessdata/eng.traineddata differ diff --git a/wwwroot/styles/global.css b/wwwroot/styles/global.css index c0dbe3f..b694fe4 100644 --- a/wwwroot/styles/global.css +++ b/wwwroot/styles/global.css @@ -59,6 +59,17 @@ a.nondecorated:hover { color: #999; } +code { + background: #222; + border-radius: 10px; + box-sizing: border-box; + font-family: 'Lucida Console'; + font-size: 8pt; + overflow-y: auto; + padding: 20px; + white-space: pre-line; +} + button, input[type=submit] { color: white; background: var(--col-button-pri); -- cgit v1.3