117 lines
5.4 KiB
Python
117 lines
5.4 KiB
Python
"""导图图片 OCR：PowerShell WinRT Windows.Media.Ocr（zh-Hans-CN），坐标还原阅读顺序"""
|
||
import re
|
||
import subprocess
|
||
import tempfile
|
||
from pathlib import Path
|
||
|
||
from .config import OCR_DIR
|
||
|
||
# WinRT StorageFile 无法处理含中文的路径 → 中间文件一律用纯 ASCII 临时目录
|
||
ASCII_TMP = Path(tempfile.gettempdir()) / "yawen_ocr"
|
||
|
||
PS_TEMPLATE = r"""$ErrorActionPreference = 'Stop'
|
||
Add-Type -AssemblyName System.Runtime.WindowsRuntime
|
||
$null = [Windows.Storage.StorageFile, Windows.Storage, ContentType = WindowsRuntime]
|
||
$null = [Windows.Media.Ocr.OcrEngine, Windows.Foundation, ContentType = WindowsRuntime]
|
||
$null = [Windows.Foundation.IAsyncOperation`1, Windows.Foundation, ContentType = WindowsRuntime]
|
||
$null = [Windows.Graphics.Imaging.BitmapDecoder, Windows.Foundation, ContentType = WindowsRuntime]
|
||
$asTaskGeneric = ([System.WindowsRuntimeSystemExtensions].GetMethods() | Where-Object { $_.Name -eq 'AsTask' -and $_.GetParameters().Count -eq 1 -and $_.GetParameters()[0].ParameterType.Name -eq 'IAsyncOperation`1' })[0]
|
||
function Await($WinRtTask, $ResultType) {
|
||
$asTask = $asTaskGeneric.MakeGenericMethod($ResultType)
|
||
$netTask = $asTask.Invoke($null, @($WinRtTask))
|
||
try { $netTask.GetAwaiter().GetResult() }
|
||
catch {
|
||
$ex = $_.Exception
|
||
while ($ex.InnerException) { $ex = $ex.InnerException }
|
||
Write-Output "AWAIT-ERR: $($ex.GetType().FullName): $($ex.Message)"
|
||
throw
|
||
}
|
||
}
|
||
$imgPath = '{{IMG}}'
|
||
$outFile = '{{OUT}}'
|
||
$file = Await ([Windows.Storage.StorageFile]::GetFileFromPathAsync($imgPath)) ([Windows.Storage.StorageFile])
|
||
$stream = Await ($file.OpenAsync([Windows.Storage.FileAccessMode]::Read)) ([Windows.Storage.Streams.IRandomAccessStream])
|
||
$decoder = Await ([Windows.Graphics.Imaging.BitmapDecoder]::CreateAsync($stream)) ([Windows.Graphics.Imaging.BitmapDecoder])
|
||
$bitmap = Await ($decoder.GetSoftwareBitmapAsync()) ([Windows.Graphics.Imaging.SoftwareBitmap])
|
||
$engine = [Windows.Media.Ocr.OcrEngine]::TryCreateFromUserProfileLanguages()
|
||
if ($null -eq $engine) {
|
||
$engine = [Windows.Media.Ocr.OcrEngine]::TryCreateFromLanguage([Windows.Globalization.Language]::new('zh-Hans-CN'))
|
||
}
|
||
$result = Await ($engine.RecognizeAsync($bitmap)) ([Windows.Media.Ocr.OcrResult])
|
||
$sb = New-Object System.Text.StringBuilder
|
||
foreach ($line in $result.Lines) {
|
||
$firstWord = $line.Words | Select-Object -First 1
|
||
$x = [int]$firstWord.BoundingRect.X
|
||
$y = [int]$firstWord.BoundingRect.Y
|
||
$text = $line.Text
|
||
$null = $sb.AppendLine("[$x,$y] $text")
|
||
}
|
||
[System.IO.File]::WriteAllText($outFile, $sb.ToString(), [System.Text.Encoding]::UTF8)
|
||
Write-Host "DONE lines=$($result.Lines.Count)"
|
||
"""
|
||
|
||
def _reorder(lines: list[tuple[int, int, str]]) -> str:
|
||
"""按坐标还原导图阅读顺序：行优先（y 容差 60px），行内按 x"""
|
||
if not lines:
|
||
return ""
|
||
lines.sort(key=lambda t: t[1])
|
||
rows: list[list[tuple[int, int, str]]] = []
|
||
cur_y, cur_row = None, []
|
||
for x, y, text in lines:
|
||
if cur_y is None or abs(y - cur_y) <= 60:
|
||
cur_row.append((x, y, text))
|
||
cur_y = y if cur_y is None else min(cur_y, y)
|
||
else:
|
||
rows.append(sorted(cur_row, key=lambda t: t[0]))
|
||
cur_row, cur_y = [(x, y, text)], y
|
||
if cur_row:
|
||
rows.append(sorted(cur_row, key=lambda t: t[0]))
|
||
return "\n".join(" ".join(t[2] for t in row) for row in rows)
|
||
|
||
|
||
def _prepare(img_path: str) -> str:
|
||
"""预处理：EXIF 转正 + 缩放到 OCR 引擎可处理尺寸（≤2000px），存纯 ASCII 路径"""
|
||
try:
|
||
from PIL import Image, ImageOps
|
||
im = Image.open(img_path)
|
||
im = ImageOps.exif_transpose(im)
|
||
w, h = im.size
|
||
if max(w, h) > 2000:
|
||
r = 2000 / max(w, h)
|
||
im = im.resize((int(w * r), int(h * r)), Image.LANCZOS)
|
||
ASCII_TMP.mkdir(parents=True, exist_ok=True)
|
||
small = ASCII_TMP / f"{Path(img_path).stem}_small.jpg"
|
||
im.convert("RGB").save(small, quality=92)
|
||
return str(small)
|
||
except Exception:
|
||
return img_path # PIL 失败就用原图
|
||
|
||
def ocr_image(img_path: str) -> dict:
|
||
"""OCR 一张图片，返回排序后的文本"""
|
||
target = _prepare(img_path)
|
||
img = Path(target)
|
||
OCR_DIR.mkdir(parents=True, exist_ok=True)
|
||
out = OCR_DIR / f"{img.stem}_ocr.txt"
|
||
ps = OCR_DIR / f"{img.stem}_ocr.ps1"
|
||
ps.write_text(PS_TEMPLATE.replace("{{IMG}}", str(img))
|
||
.replace("{{OUT}}", str(out)),
|
||
encoding="utf-8-sig") # BOM：PowerShell 5.1 按 UTF-8 解析脚本
|
||
try:
|
||
r = subprocess.run(["powershell", "-NoProfile", "-ExecutionPolicy", "Bypass",
|
||
"-File", str(ps)], capture_output=True, text=True,
|
||
encoding="gbk", errors="replace", timeout=180)
|
||
except Exception as e:
|
||
return {"ok": False, "error": f"OCR 进程失败: {e}"}
|
||
if not out.exists():
|
||
return {"ok": False, "error": (r.stderr or r.stdout or "OCR 无输出")[:300]}
|
||
raw = out.read_text(encoding="utf-8", errors="ignore")
|
||
lines = []
|
||
for ln in raw.splitlines():
|
||
m = re.match(r"\[(\d+),(\d+)\] (.*)", ln)
|
||
if m:
|
||
lines.append((int(m.group(1)), int(m.group(2)), m.group(3).strip()))
|
||
text = _reorder(lines)
|
||
if not text.strip():
|
||
return {"ok": False, "error": "OCR 未识别到文字（图片可能不清晰）"}
|
||
return {"ok": True, "text": text, "lines": len(lines)}
|