Skip to content
Merged
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
62 changes: 47 additions & 15 deletions TelegramSearchBot.Test/Service/AI/OCR/LongImageOcrSlicerTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -54,18 +54,39 @@ public void PlanSlices_WithoutBlankBand_UsesOverlapAndCoversImage() {
}

[Fact]
public void MergeResults_RemovesRepeatedBoundaryLines() {
public void PlanSlices_SingleBlankRowStillUsesOverlap() {
using var image = CreateDenseImage(800, 7000);
FillRows(image, LongImageOcrSlicer.TargetSliceHeight, LongImageOcrSlicer.TargetSliceHeight + 1, SKColors.White);

var slices = LongImageOcrSlicer.PlanSlices(image);

Assert.True(slices[1].HasOverlap);
AssertCompleteCoverage(slices, image.Height);
}

[Fact]
public void MergeResults_RemovesRepeatedBoundaryLinesFromOverlappingSlices() {
var result = LongImageOcrSlicer.MergeResults([
"第一段 重叠块一 重叠块二",
"重叠块一 重叠块二 第二段",
"第三段"
(new OcrImageSlice(0, 3072, false), "第一段 重叠块一 重叠块二"),
(new OcrImageSlice(2976, 3072, true), "重叠块一 重叠块二 第二段"),
(new OcrImageSlice(6048, 952, false), "第三段")
]);

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

[Fact]
public async Task ExecuteAsync_LongImage_RunsSlicesInOrderAndMergesResults() {
public void MergeResults_PreservesRepeatedBoundaryTokensWithoutOverlap() {
var result = LongImageOcrSlicer.MergeResults([
(new OcrImageSlice(0, 3072, false), "第一段 2026"),
(new OcrImageSlice(3072, 3928, false), "2026 第二段")
]);

Assert.Equal("第一段 2026 2026 第二段", result);
}

[Fact]
public async Task ExecuteAsync_LongImage_PreservesRepeatedTokensAcrossSafeCut() {
var calls = new List<(int Width, int Height)>();
var responses = new Queue<string>(["第一段 重叠块", "重叠块 第二段"]);
var service = new PaddleOCRService(
Expand All @@ -82,6 +103,19 @@ public async Task ExecuteAsync_LongImage_RunsSlicesInOrderAndMergesResults() {
Assert.Equal(2, calls.Count);
Assert.All(calls, call => Assert.Equal(1000, call.Width));
Assert.All(calls, call => Assert.InRange(call.Height, 1, LongImageOcrSlicer.LongImageHeightThreshold));
Assert.Equal("第一段 重叠块 重叠块 第二段", result);
}

[Fact]
public async Task ExecuteAsync_OverlappingSlicesRemoveRepeatedBoundaryTokens() {
var responses = new Queue<string>(["第一段 重叠块", "重叠块 第二段"]);
var service = new PaddleOCRService(
Mock.Of<IConnectionMultiplexer>(),
payload => Task.FromResult(responses.Dequeue()));
using var stream = Encode(CreateDenseImage(1000, 5000));

var result = await service.ExecuteAsync(stream);

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

Expand Down Expand Up @@ -121,11 +155,11 @@ public async Task ExecuteAsync_WhenSliceFails_StopsProcessing() {
}

private static SKBitmap CreateDenseImage(int width, int height) {
var image = new SKBitmap(width, height);
for (var y = 0; y < height; y++) {
for (var x = 0; x < width; x++) {
image.SetPixel(x, y, x % 16 < 8 ? SKColors.Black : SKColors.White);
}
var image = CreateWhiteImage(width, height);
using var canvas = new SKCanvas(image);
using var paint = new SKPaint { Color = SKColors.Black };
for (var x = 0; x < width; x += 16) {
canvas.DrawRect(x, 0, 8, height, paint);
}
return image;
}
Expand All @@ -137,11 +171,9 @@ private static SKBitmap CreateWhiteImage(int width, int height) {
}

private static void FillRows(SKBitmap image, int start, int end, SKColor color) {
for (var y = start; y < end; y++) {
for (var x = 0; x < image.Width; x++) {
image.SetPixel(x, y, color);
}
}
using var canvas = new SKCanvas(image);
using var paint = new SKPaint { Color = color };
canvas.DrawRect(0, start, image.Width, end - start, paint);
}

private static MemoryStream Encode(SKBitmap image) {
Expand Down
22 changes: 16 additions & 6 deletions TelegramSearchBot/Service/AI/OCR/LongImageOcrSlicer.cs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ internal static class LongImageOcrSlicer {
private const int AnalysisWidth = 256;
private const int MinimumSliceHeight = 1024;
private const int BrightnessDifferenceThreshold = 24;
private const int CutDensityRadius = 3;
private const double SafeRowDensityThreshold = 0.025;

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

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

internal static string MergeResults(IEnumerable<string> results) {
internal static string MergeResults(IEnumerable<(OcrImageSlice Slice, string Text)> results) {
var mergedTokens = new List<string>();
var previousSliceHadText = false;

foreach (var result in results.Where(result => !string.IsNullOrWhiteSpace(result))) {
var nextTokens = SplitTokens(result);
var duplicateTokenCount = FindDuplicateBoundary(mergedTokens, nextTokens);
foreach (var (slice, text) in results) {
if (string.IsNullOrWhiteSpace(text)) {
previousSliceHadText = false;
continue;
}

var nextTokens = SplitTokens(text);
var duplicateTokenCount = slice.HasOverlap && previousSliceHadText
? FindDuplicateBoundary(mergedTokens, nextTokens)
: 0;
mergedTokens.AddRange(nextTokens.Skip(duplicateTokenCount));
previousSliceHadText = true;
}

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

for (var y = searchStart; y <= searchEnd; y++) {
var density = AverageDensity(rowDensities, y, 3);
var density = AverageDensity(rowDensities, y, CutDensityRadius);
var distancePenalty = Math.Abs(y - target) / ( double ) Math.Max(1, searchEnd - searchStart) * 0.01;
var score = density + distancePenalty;
if (score < bestScore) {
Expand Down
4 changes: 2 additions & 2 deletions TelegramSearchBot/Service/AI/OCR/PaddleOCRService.cs
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,9 @@ public async Task<string> ExecuteAsync(Stream file) {
return await ExecuteSliceAsync(image, slices[0]);
}

var results = new List<string>(slices.Count);
var results = new List<(OcrImageSlice Slice, string Text)>(slices.Count);
foreach (var slice in slices) {
results.Add(await ExecuteSliceAsync(image, slice));
results.Add((slice, await ExecuteSliceAsync(image, slice)));
}

return LongImageOcrSlicer.MergeResults(results);
Expand Down
Loading