Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
263 changes: 142 additions & 121 deletions src/UI/Logic/Ocr/NOcrDb.cs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,56 @@ public class NOcrDb
public List<NOcrChar> OcrCharacters = new();
public List<NOcrChar> OcrCharactersExpanded = new();

private enum ExpandedMatchPhase
{
Exact,
Relaxed,
}

private sealed class ExpandedMatchCandidate
{
public NOcrChar Template { get; }
public int DbIndex { get; }
public ExpandedMatchPhase Phase { get; }
public int ActualWidth { get; }
public int ActualHeight { get; }
public double ActualWidthPercent { get; }
public int WidthDelta { get; }
public int HeightDelta { get; }
public double WidthPercentDelta { get; }
public bool ForegroundMatched { get; }
public bool BackgroundMatched { get; }
public bool SizeMatched { get; }

public ExpandedMatchCandidate(
NOcrChar template,
int dbIndex,
ExpandedMatchPhase phase,
int actualWidth,
int actualHeight,
double actualWidthPercent,
int widthDelta,
int heightDelta,
double widthPercentDelta,
bool foregroundMatched,
bool backgroundMatched,
bool sizeMatched)
{
Template = template;
DbIndex = dbIndex;
Phase = phase;
ActualWidth = actualWidth;
ActualHeight = actualHeight;
ActualWidthPercent = actualWidthPercent;
WidthDelta = widthDelta;
HeightDelta = heightDelta;
WidthPercentDelta = widthPercentDelta;
ForegroundMatched = foregroundMatched;
BackgroundMatched = backgroundMatched;
SizeMatched = sizeMatched;
}
}

public int TotalCharacterCount => OcrCharacters.Count + OcrCharactersExpanded.Count;

private const string Version = "V2";
Expand Down Expand Up @@ -141,148 +191,119 @@ public void Remove(NOcrChar ocrChar)
return null;
}

var w = targetItem.NikseBitmap.Width;
var exactCandidates = GetExpandedMatchCandidates(nikseBitmap, targetItem, listIndex, list, ExpandedMatchPhase.Exact);
if (exactCandidates.Count > 0)
{
return exactCandidates[0].Template;
}

var relaxedCandidates = GetExpandedMatchCandidates(nikseBitmap, targetItem, listIndex, list, ExpandedMatchPhase.Relaxed);
if (relaxedCandidates.Count > 0)
{
return relaxedCandidates[0].Template;
}

return null;
}

private List<ExpandedMatchCandidate> GetExpandedMatchCandidates(NikseBitmap2 nikseBitmap, ImageSplitterItem2 targetItem, int listIndex, List<ImageSplitterItem2> list, ExpandedMatchPhase phase)
{
var candidates = new List<ExpandedMatchCandidate>();
var targetWidth = targetItem.NikseBitmap?.Width ?? 0;

for (var i = 0; i < OcrCharactersExpanded.Count; i++)
{
var oc = OcrCharactersExpanded[i];
if (oc.ExpandCount > 1 && oc.Width > w && targetItem.X + oc.Width < nikseBitmap.Width)
if (oc.ExpandCount <= 1 || oc.Width <= targetWidth || targetItem.X + oc.Width >= nikseBitmap.Width)
{
var ok = true;
var index = 0;
while (index < oc.LinesForeground.Count && ok)
{
var op = oc.LinesForeground[index];
foreach (var point in op.GetPoints())
{
var p = new OcrPoint(point.X + targetItem.X, point.Y + targetItem.Y - oc.MarginTop);
if (p.X >= 0 && p.Y >= 0 && p.X < nikseBitmap.Width && p.Y < nikseBitmap.Height)
{
var a = nikseBitmap.GetAlpha(p.X, p.Y);
if (a <= 150)
{
ok = false;
break;
}
}
else if (p.X >= 0 && p.Y >= 0)
{
ok = false;
break;
}
}
continue;
}

index++;
}
var candidate = TryCreateExpandedCandidate(nikseBitmap, targetItem, listIndex, list, phase, i, oc);
if (candidate != null)
{
candidates.Add(candidate);
}
}

index = 0;
while (index < oc.LinesBackground.Count && ok)
{
var op = oc.LinesBackground[index];
foreach (var point in op.GetPoints())
{
var p = new OcrPoint(point.X + targetItem.X, point.Y + targetItem.Y - oc.MarginTop);
if (p.X >= 0 && p.Y >= 0 && p.X < nikseBitmap.Width && p.Y < nikseBitmap.Height)
{
var a = nikseBitmap.GetAlpha(p.X, p.Y);
if (a > 150)
{
ok = false;
break;
}
}
else if (p.X >= 0 && p.Y >= 0)
{
ok = false;
break;
}
}
return candidates;
}

index++;
}
private ExpandedMatchCandidate? TryCreateExpandedCandidate(NikseBitmap2 nikseBitmap, ImageSplitterItem2 targetItem, int listIndex, List<ImageSplitterItem2> list, ExpandedMatchPhase phase, int dbIndex, NOcrChar oc)
{
var foregroundMatched = TryMatchExpandedLineSet(nikseBitmap, targetItem, oc, oc.LinesForeground, phase, true);
if (!foregroundMatched)
{
return null;
}

if (ok)
{
var size = GetTotalSize(listIndex, list, oc.ExpandCount);
if (Math.Abs(size.X - oc.Width) < 3 && Math.Abs(size.Y - oc.Height) < 3)
{
return oc;
}
}
}
var backgroundMatched = TryMatchExpandedLineSet(nikseBitmap, targetItem, oc, oc.LinesBackground, phase, false);
if (!backgroundMatched)
{
return null;
}

for (var i = 0; i < OcrCharactersExpanded.Count; i++)
var size = GetTotalSize(listIndex, list, oc.ExpandCount);
var actualWidthPercent = GetWidthPercent(size.X, size.Y);
var widthDelta = Math.Abs(size.X - oc.Width);
var heightDelta = Math.Abs(size.Y - oc.Height);
var widthPercentDelta = Math.Abs(actualWidthPercent - oc.WidthPercent);
var sizeMatched = phase == ExpandedMatchPhase.Exact
? widthDelta < 3 && heightDelta < 3
: widthPercentDelta < 15 && widthDelta < 25 && heightDelta < 20;

if (!sizeMatched)
{
var oc = OcrCharactersExpanded[i];
if (oc.ExpandCount > 1 && oc.Width > w && targetItem.X + oc.Width < nikseBitmap.Width)
{
var ok = true;
var index = 0;
while (index < oc.LinesForeground.Count && ok)
{
var op = oc.LinesForeground[index];
foreach (var point in op.ScaledGetPoints(oc, oc.Width, oc.Height - 1))
{
var p = new OcrPoint(point.X + targetItem.X, point.Y + targetItem.Y - oc.MarginTop);
if (p.X >= 0 && p.Y >= 0 && p.X < nikseBitmap.Width && p.Y < nikseBitmap.Height)
{
var a = nikseBitmap.GetAlpha(p.X, p.Y);
if (a <= 150)
{
ok = false;
break;
}
}
else if (p.X >= 0 && p.Y >= 0)
{
ok = false;
break;
}
}
return null;
}

index++;
}
return new ExpandedMatchCandidate(
oc,
dbIndex,
phase,
size.X,
size.Y,
actualWidthPercent,
widthDelta,
heightDelta,
widthPercentDelta,
foregroundMatched,
backgroundMatched,
sizeMatched);
}

private static bool TryMatchExpandedLineSet(NikseBitmap2 nikseBitmap, ImageSplitterItem2 targetItem, NOcrChar oc, List<NOcrLine> lines, ExpandedMatchPhase phase, bool isForeground)
{
foreach (var op in lines)
{
IEnumerable<OcrPoint> points = phase == ExpandedMatchPhase.Exact
? op.GetPoints()
: op.ScaledGetPoints(oc, oc.Width, oc.Height - 1);

index = 0;
while (index < oc.LinesBackground.Count && ok)
foreach (var point in points)
{
var p = new OcrPoint(point.X + targetItem.X, point.Y + targetItem.Y - oc.MarginTop);
if (p.X >= 0 && p.Y >= 0 && p.X < nikseBitmap.Width && p.Y < nikseBitmap.Height)
{
var op = oc.LinesBackground[index];
foreach (var point in op.ScaledGetPoints(oc, oc.Width, oc.Height - 1))
var a = nikseBitmap.GetAlpha(p.X, p.Y);
if (isForeground ? a <= 150 : a > 150)
{
var p = new OcrPoint(point.X + targetItem.X, point.Y + targetItem.Y - oc.MarginTop);
if (p.X >= 0 && p.Y >= 0 && p.X < nikseBitmap.Width && p.Y < nikseBitmap.Height)
{
var a = nikseBitmap.GetAlpha(p.X, p.Y);
if (a > 150)
{
ok = false;
break;
}
}
else if (p.X >= 0 && p.Y >= 0)
{
ok = false;
break;
}
return false;
}

index++;
}

if (ok)
else if (p.X >= 0 && p.Y >= 0)
{
var size = GetTotalSize(listIndex, list, oc.ExpandCount);
var widthPercent = size.Y * 100.0 / size.X;
if (Math.Abs(widthPercent - oc.WidthPercent) < 15 &&
Math.Abs(size.X - oc.Width) < 25 && Math.Abs(size.Y - oc.Height) < 20)
{
return oc;
}
return false;
}
}
}

return null;
return true;
}

private static double GetWidthPercent(int width, int height)
{
return width == 0 ? double.PositiveInfinity : height * 100.0 / width;
}

private static OcrPoint GetTotalSize(int listIndex, List<ImageSplitterItem2> items, int count)
Expand Down Expand Up @@ -548,4 +569,4 @@ public static List<string> GetDatabases()
.OrderBy(p => p)
.ToList()!;
}
}
}
87 changes: 87 additions & 0 deletions tests/UI/Logic/Ocr/NOcrDbExpandedMatchTests.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
using Nikse.SubtitleEdit.Logic.Ocr;

namespace UITests.Logic.Ocr;

public class NOcrDbExpandedMatchTests
{
[Fact]
public void GetMatchExpanded_WhenMultipleExactCandidatesMatch_ReturnsFirstCandidateInDbOrder()
{
var db = CreateDb();
var firstExact = CreateExpandedChar("first-exact", width: 5, height: 3);
var secondExact = CreateExpandedChar("second-exact", width: 5, height: 3);
db.OcrCharactersExpanded = new List<NOcrChar> { firstExact, secondExact };
var letters = CreateLetters();

var match = db.GetMatchExpanded(CreateParentBitmap(), letters[0], 0, letters);

Assert.Same(firstExact, match);
}

[Fact]
public void GetMatchExpanded_WhenExactAndRelaxedCandidatesMatch_PrefersExactPhase()
{
var db = CreateDb();
var relaxed = CreateExpandedChar("relaxed", width: 10, height: 6);
var exact = CreateExpandedChar("exact", width: 5, height: 3);
db.OcrCharactersExpanded = new List<NOcrChar> { relaxed, exact };
var letters = CreateLetters();

var match = db.GetMatchExpanded(CreateParentBitmap(), letters[0], 0, letters);

Assert.Same(exact, match);
}

[Fact]
public void GetMatchExpanded_WhenOnlyRelaxedCandidatesMatch_ReturnsFirstCandidateInDbOrder()
{
var db = CreateDb();
var firstRelaxed = CreateExpandedChar("first-relaxed", width: 10, height: 6);
var secondRelaxed = CreateExpandedChar("second-relaxed", width: 10, height: 6);
db.OcrCharactersExpanded = new List<NOcrChar> { firstRelaxed, secondRelaxed };
var letters = CreateLetters();

var match = db.GetMatchExpanded(CreateParentBitmap(), letters[0], 0, letters);

Assert.Same(firstRelaxed, match);
}

private static NOcrDb CreateDb()
{
var fileName = Path.Combine(Path.GetTempPath(), Path.GetRandomFileName() + ".nocr");
return new NOcrDb(fileName);
}

private static NOcrChar CreateExpandedChar(string text, int width, int height, IEnumerable<NOcrLine>? foregroundLines = null, IEnumerable<NOcrLine>? backgroundLines = null)
{
var item = new NOcrChar(text)
{
Width = width,
Height = height,
MarginTop = 0,
ExpandCount = 2,
};
item.LinesForeground.AddRange(foregroundLines ?? new[]
{
new NOcrLine(new OcrPoint(0, 0), new OcrPoint(0, 0)),
});
item.LinesBackground.AddRange(backgroundLines ?? []);
return item;
}

private static NikseBitmap2 CreateParentBitmap()
{
var bitmap = new NikseBitmap2(30, 20);
bitmap.SetAlpha(5, 5, byte.MaxValue);
return bitmap;
}

private static List<ImageSplitterItem2> CreateLetters()
{
return new List<ImageSplitterItem2>
{
new ImageSplitterItem2(5, 5, new NikseBitmap2(2, 3)),
new ImageSplitterItem2(8, 5, new NikseBitmap2(2, 3)),
};
}
}