Skip to content

Commit fb2e098

Browse files
authored
Test OCR slice boundary handling (#369)
1 parent 9ced0fa commit fb2e098

3 files changed

Lines changed: 65 additions & 23 deletions

File tree

TelegramSearchBot.Test/Service/AI/OCR/LongImageOcrSlicerTests.cs

Lines changed: 47 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -54,18 +54,39 @@ public void PlanSlices_WithoutBlankBand_UsesOverlapAndCoversImage() {
5454
}
5555

5656
[Fact]
57-
public void MergeResults_RemovesRepeatedBoundaryLines() {
57+
public void PlanSlices_SingleBlankRowStillUsesOverlap() {
58+
using var image = CreateDenseImage(800, 7000);
59+
FillRows(image, LongImageOcrSlicer.TargetSliceHeight, LongImageOcrSlicer.TargetSliceHeight + 1, SKColors.White);
60+
61+
var slices = LongImageOcrSlicer.PlanSlices(image);
62+
63+
Assert.True(slices[1].HasOverlap);
64+
AssertCompleteCoverage(slices, image.Height);
65+
}
66+
67+
[Fact]
68+
public void MergeResults_RemovesRepeatedBoundaryLinesFromOverlappingSlices() {
5869
var result = LongImageOcrSlicer.MergeResults([
59-
"第一段 重叠块一 重叠块二",
60-
"重叠块一 重叠块二 第二段",
61-
"第三段"
70+
(new OcrImageSlice(0, 3072, false), "第一段 重叠块一 重叠块二"),
71+
(new OcrImageSlice(2976, 3072, true), "重叠块一 重叠块二 第二段"),
72+
(new OcrImageSlice(6048, 952, false), "第三段")
6273
]);
6374

6475
Assert.Equal("第一段 重叠块一 重叠块二 第二段 第三段", result);
6576
}
6677

6778
[Fact]
68-
public async Task ExecuteAsync_LongImage_RunsSlicesInOrderAndMergesResults() {
79+
public void MergeResults_PreservesRepeatedBoundaryTokensWithoutOverlap() {
80+
var result = LongImageOcrSlicer.MergeResults([
81+
(new OcrImageSlice(0, 3072, false), "第一段 2026"),
82+
(new OcrImageSlice(3072, 3928, false), "2026 第二段")
83+
]);
84+
85+
Assert.Equal("第一段 2026 2026 第二段", result);
86+
}
87+
88+
[Fact]
89+
public async Task ExecuteAsync_LongImage_PreservesRepeatedTokensAcrossSafeCut() {
6990
var calls = new List<(int Width, int Height)>();
7091
var responses = new Queue<string>(["第一段 重叠块", "重叠块 第二段"]);
7192
var service = new PaddleOCRService(
@@ -82,6 +103,19 @@ public async Task ExecuteAsync_LongImage_RunsSlicesInOrderAndMergesResults() {
82103
Assert.Equal(2, calls.Count);
83104
Assert.All(calls, call => Assert.Equal(1000, call.Width));
84105
Assert.All(calls, call => Assert.InRange(call.Height, 1, LongImageOcrSlicer.LongImageHeightThreshold));
106+
Assert.Equal("第一段 重叠块 重叠块 第二段", result);
107+
}
108+
109+
[Fact]
110+
public async Task ExecuteAsync_OverlappingSlicesRemoveRepeatedBoundaryTokens() {
111+
var responses = new Queue<string>(["第一段 重叠块", "重叠块 第二段"]);
112+
var service = new PaddleOCRService(
113+
Mock.Of<IConnectionMultiplexer>(),
114+
payload => Task.FromResult(responses.Dequeue()));
115+
using var stream = Encode(CreateDenseImage(1000, 5000));
116+
117+
var result = await service.ExecuteAsync(stream);
118+
85119
Assert.Equal("第一段 重叠块 第二段", result);
86120
}
87121

@@ -121,11 +155,11 @@ public async Task ExecuteAsync_WhenSliceFails_StopsProcessing() {
121155
}
122156

123157
private static SKBitmap CreateDenseImage(int width, int height) {
124-
var image = new SKBitmap(width, height);
125-
for (var y = 0; y < height; y++) {
126-
for (var x = 0; x < width; x++) {
127-
image.SetPixel(x, y, x % 16 < 8 ? SKColors.Black : SKColors.White);
128-
}
158+
var image = CreateWhiteImage(width, height);
159+
using var canvas = new SKCanvas(image);
160+
using var paint = new SKPaint { Color = SKColors.Black };
161+
for (var x = 0; x < width; x += 16) {
162+
canvas.DrawRect(x, 0, 8, height, paint);
129163
}
130164
return image;
131165
}
@@ -137,11 +171,9 @@ private static SKBitmap CreateWhiteImage(int width, int height) {
137171
}
138172

139173
private static void FillRows(SKBitmap image, int start, int end, SKColor color) {
140-
for (var y = start; y < end; y++) {
141-
for (var x = 0; x < image.Width; x++) {
142-
image.SetPixel(x, y, color);
143-
}
144-
}
174+
using var canvas = new SKCanvas(image);
175+
using var paint = new SKPaint { Color = color };
176+
canvas.DrawRect(0, start, image.Width, end - start, paint);
145177
}
146178

147179
private static MemoryStream Encode(SKBitmap image) {

TelegramSearchBot/Service/AI/OCR/LongImageOcrSlicer.cs

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ internal static class LongImageOcrSlicer {
1616
private const int AnalysisWidth = 256;
1717
private const int MinimumSliceHeight = 1024;
1818
private const int BrightnessDifferenceThreshold = 24;
19+
private const int CutDensityRadius = 3;
1920
private const double SafeRowDensityThreshold = 0.025;
2021

2122
internal static IReadOnlyList<OcrImageSlice> PlanSlices(SKBitmap image) {
@@ -33,7 +34,7 @@ internal static IReadOnlyList<OcrImageSlice> PlanSlices(SKBitmap image) {
3334
var searchStart = Math.Max(top + MinimumSliceHeight, targetBottom - CutSearchRadius);
3435
var searchEnd = Math.Min(image.Height - 1, targetBottom + CutSearchRadius);
3536
var cut = FindBestCut(rowDensities, searchStart, searchEnd);
36-
var safeCut = rowDensities[cut] <= SafeRowDensityThreshold;
37+
var safeCut = AverageDensity(rowDensities, cut, CutDensityRadius) <= SafeRowDensityThreshold;
3738
var hasOverlap = slices.Count > 0 && slices[^1].Top + slices[^1].Height > top;
3839

3940
slices.Add(new OcrImageSlice(top, cut - top, hasOverlap));
@@ -44,13 +45,22 @@ internal static IReadOnlyList<OcrImageSlice> PlanSlices(SKBitmap image) {
4445
return slices;
4546
}
4647

47-
internal static string MergeResults(IEnumerable<string> results) {
48+
internal static string MergeResults(IEnumerable<(OcrImageSlice Slice, string Text)> results) {
4849
var mergedTokens = new List<string>();
50+
var previousSliceHadText = false;
4951

50-
foreach (var result in results.Where(result => !string.IsNullOrWhiteSpace(result))) {
51-
var nextTokens = SplitTokens(result);
52-
var duplicateTokenCount = FindDuplicateBoundary(mergedTokens, nextTokens);
52+
foreach (var (slice, text) in results) {
53+
if (string.IsNullOrWhiteSpace(text)) {
54+
previousSliceHadText = false;
55+
continue;
56+
}
57+
58+
var nextTokens = SplitTokens(text);
59+
var duplicateTokenCount = slice.HasOverlap && previousSliceHadText
60+
? FindDuplicateBoundary(mergedTokens, nextTokens)
61+
: 0;
5362
mergedTokens.AddRange(nextTokens.Skip(duplicateTokenCount));
63+
previousSliceHadText = true;
5464
}
5565

5666
return string.Join(' ', mergedTokens);
@@ -97,7 +107,7 @@ private static int FindBestCut(double[] rowDensities, int searchStart, int searc
97107
var bestScore = double.MaxValue;
98108

99109
for (var y = searchStart; y <= searchEnd; y++) {
100-
var density = AverageDensity(rowDensities, y, 3);
110+
var density = AverageDensity(rowDensities, y, CutDensityRadius);
101111
var distancePenalty = Math.Abs(y - target) / ( double ) Math.Max(1, searchEnd - searchStart) * 0.01;
102112
var score = density + distancePenalty;
103113
if (score < bestScore) {

TelegramSearchBot/Service/AI/OCR/PaddleOCRService.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,9 @@ public async Task<string> ExecuteAsync(Stream file) {
4343
return await ExecuteSliceAsync(image, slices[0]);
4444
}
4545

46-
var results = new List<string>(slices.Count);
46+
var results = new List<(OcrImageSlice Slice, string Text)>(slices.Count);
4747
foreach (var slice in slices) {
48-
results.Add(await ExecuteSliceAsync(image, slice));
48+
results.Add((slice, await ExecuteSliceAsync(image, slice)));
4949
}
5050

5151
return LongImageOcrSlicer.MergeResults(results);

0 commit comments

Comments
 (0)