feat: OCR

main/src/main.ts

@@ -36,10 +36,10 @@ app.on('ready', async () => {
     async () => {
       const overlay = new OverlayWindow(eventPipe, logger, poeWindow, httpProxy)
       new OverlayVisibility(eventPipe, overlay, gameConfig)
-      const shortcuts = new Shortcuts(logger, overlay, poeWindow, gameConfig, eventPipe)
+      const shortcuts = await Shortcuts.create(logger, overlay, poeWindow, gameConfig, eventPipe)
       eventPipe.onEventAnyClient('CLIENT->MAIN::update-host-config', (cfg) => {
         overlay.updateOpts(cfg.overlayKey, cfg.windowTitle)
-        shortcuts.updateActions(cfg.shortcuts, cfg.stashScroll, cfg.restoreClipboard)
+        shortcuts.updateActions(cfg.shortcuts, cfg.stashScroll, cfg.restoreClipboard, cfg.language)
         gameLogWatcher.restart(cfg.clientLog)
         gameConfig.readConfig(cfg.gameConfig)
         appUpdater.updateOpts(!cfg.disableUpdateDownload)

main/src/shortcuts/Shortcuts.ts

@@ -4,6 +4,7 @@ import { isModKey, KeyToElectron, mergeTwoHotkeys } from '../../../ipc/KeyToCode
 import { typeInChat, stashSearch } from './text-box'
 import { WidgetAreaTracker } from '../windowing/WidgetAreaTracker'
 import { HostClipboard } from './HostClipboard'
+import { OcrWorker } from '../vision/link-main'
 import type { ShortcutAction } from '../../../ipc/types'
 import type { Logger } from '../RemoteLogger'
 import type { OverlayWindow } from '../windowing/OverlayWindow'
@@ -20,12 +21,25 @@ export class Shortcuts {
   private areaTracker: WidgetAreaTracker
   private clipboard: HostClipboard
 
-  constructor (
+  static async create (
+    logger: Logger,
+    overlay: OverlayWindow,
+    poeWindow: GameWindow,
+    gameConfig: GameConfig,
+    server: ServerEvents
+  ) {
+    const ocrWorker = await OcrWorker.create()
+    const shortcuts = new Shortcuts(logger, overlay, poeWindow, gameConfig, server, ocrWorker)
+    return shortcuts
+  }
+
+  private constructor (
     private logger: Logger,
     private overlay: OverlayWindow,
     private poeWindow: GameWindow,
     private gameConfig: GameConfig,
-    private server: ServerEvents
+    private server: ServerEvents,
+    private ocrWorker: OcrWorker
   ) {
     this.areaTracker = new WidgetAreaTracker(server, overlay)
     this.clipboard = new HostClipboard(logger)
@@ -69,9 +83,10 @@ export class Shortcuts {
     })
   }
 
-  updateActions (actions: ShortcutAction[], stashScroll: boolean, restoreClipboard: boolean) {
+  updateActions (actions: ShortcutAction[], stashScroll: boolean, restoreClipboard: boolean, language: string) {
     this.stashScroll = stashScroll
     this.clipboard.updateOptions(restoreClipboard)
+    this.ocrWorker.updateOptions(language)
 
     const copyItemShortcut = mergeTwoHotkeys('Ctrl + C', this.gameConfig.showModsKey)
     if (copyItemShortcut !== 'Ctrl + C') {
@@ -154,6 +169,27 @@ export class Shortcuts {
             (entry.keepModKeys) ? entry.shortcut.split(' + ').filter(key => isModKey(key)) : undefined,
             this.gameConfig.showModsKey
           )
+        } else if (entry.action.type === 'ocr-text' && entry.action.target === 'heist-gems') {
+          if (process.platform !== 'win32') return
+
+          const { action } = entry
+          const pressTime = Date.now()
+          const imageData = this.poeWindow.screenshot()
+          this.ocrWorker.findHeistGems({
+            width: this.poeWindow.bounds.width,
+            height: this.poeWindow.bounds.height,
+            data: imageData
+          }).then(result => {
+            this.server.sendEventTo('last-active', {
+              name: 'MAIN->CLIENT::ocr-text',
+              payload: {
+                target: action.target,
+                pressTime,
+                ocrTime: result.elapsed,
+                paragraphs: result.recognized.map(p => p.text)
+              }
+            })
+          }).catch(() => {})
         }
       })
 

main/src/vision/HeistGemFinder.ts

@@ -0,0 +1,125 @@
+import fs from 'fs/promises'
+import Bmp from '@wokwi/bmp-ts'
+import * as Bindings from './wasm-bindings'
+import { cv, tessApi } from './wasm-bindings'
+import {
+  findNonZeroWeights,
+  groupWeightedPoints,
+  findLines,
+  hsvToU8,
+  timeIt,
+  ImageData
+} from './utils'
+
+const REFERENCE_HEIGHT = 600
+const TEMPLATE_THRESHOLD = 0.75
+const LINE_Y_DIST_TOLERANCE = 4
+const TEXT_HSV_MIN = hsvToU8(173, 31, 31)
+const TEXT_HSV_MAX = hsvToU8(180, 100, 100)
+
+interface OcrResult {
+  elapsed: number
+  matches: number
+  clusters: number
+  linesMin: number
+  linesMax: number
+  recognized: Array<{ text: string, confidence: number }>
+}
+
+export class HeistGemFinder {
+  private constructor (
+    private readonly needleMat: any,
+    private readonly hsvMin: any,
+    private readonly hsvMax: any
+  ) {}
+
+  static async create (binDir: string): Promise<HeistGemFinder> {
+    const needleImg = Bmp.decode(await fs.readFile(binDir + '/heist-lock.bmp'), { toRGBA: true })
+    const needleMat = Bindings.cvMatFromImage(needleImg)
+    cv.cvtColor(needleMat, needleMat, cv.COLOR_RGBA2GRAY)
+
+    const hsvMin = new cv.Mat(3, 1, cv.CV_8U)
+    hsvMin.data.set(TEXT_HSV_MIN)
+    const hsvMax = new cv.Mat(3, 1, cv.CV_8U)
+    hsvMax.data.set(TEXT_HSV_MAX)
+
+    return new HeistGemFinder(needleMat, hsvMin, hsvMax)
+  }
+
+  ocrScreenshot (screenshot: ImageData): OcrResult {
+    let elapsed = 0
+    const colorMat = Bindings.cvMatFromImage(screenshot)
+
+    const scale = screenshot.height / REFERENCE_HEIGHT
+    const graySize = new cv.Size(Math.floor(screenshot.width / scale), REFERENCE_HEIGHT)
+    const grayMat = new cv.Mat()
+    elapsed += timeIt(() => {
+      if (scale > 2.1) {
+        cv.resize(colorMat, grayMat, new cv.Size(graySize.width * 2, graySize.height * 2), 0, 0, cv.INTER_LINEAR)
+        cv.resize(grayMat, grayMat, new cv.Size(graySize.width, graySize.height), 0, 0, cv.INTER_LINEAR)
+      } else {
+        cv.resize(colorMat, grayMat, new cv.Size(graySize.width, graySize.height), 0, 0, cv.INTER_LINEAR)
+      }
+      cv.cvtColor(grayMat, grayMat, cv.COLOR_BGR2GRAY)
+    })
+
+    const { needleMat } = this
+    const matchesMat = grayMat
+    let matches: ReturnType<typeof findNonZeroWeights>
+    elapsed += timeIt(() => {
+      cv.matchTemplate(grayMat, needleMat, matchesMat, cv.TM_CCOEFF_NORMED)
+      cv.threshold(matchesMat, matchesMat, TEMPLATE_THRESHOLD, 1, cv.THRESH_TOZERO)
+      matches = findNonZeroWeights(matchesMat)
+    })
+    matchesMat.delete()
+    const clusteredMatches = groupWeightedPoints(matches!, Math.hypot(needleMat.cols, needleMat.rows))
+    const lines = findLines(clusteredMatches, LINE_Y_DIST_TOLERANCE)
+
+    const recognizedLines: OcrResult['recognized'] = []
+    for (const line of lines) {
+      const topLeft = new cv.Point(
+        (Math.min(line[0].x, line[1].x) + needleMat.cols) * scale,
+        (Math.min(line[0].y, line[1].y) - 1) * scale)
+      const bottomRight = new cv.Point(
+        Math.max(line[0].x, line[1].x) * scale,
+        (Math.max(line[0].y, line[1].y) + needleMat.rows) * scale)
+
+      const roiSize = new cv.Size(bottomRight.x - topLeft.x, bottomRight.y - topLeft.y)
+      const roiRect = new cv.Rect(topLeft, roiSize)
+      const roiColor = colorMat.roi(roiRect)
+      const roiHsv = new cv.Mat()
+
+      elapsed += timeIt(() => {
+        cv.cvtColor(roiColor, roiHsv, cv.COLOR_BGR2HSV_FULL)
+        cv.inRange(roiHsv, this.hsvMin, this.hsvMax, roiHsv)
+        cv.bitwise_not(roiHsv, roiHsv)
+      })
+      roiColor.delete()
+
+      Bindings.ocrSetImage(roiHsv.data, roiHsv.cols, roiHsv.rows, roiHsv.channels())
+      roiHsv.delete()
+      tessApi.SetVariable('tessedit_pageseg_mode', '7') // single line mode
+      elapsed += timeIt(() => {
+        tessApi.Recognize()
+      })
+      const text = tessApi.GetUTF8Text().trim()
+      const confidence = tessApi.MeanTextConf()
+      if (text.length > 0 && confidence > 30) {
+        recognizedLines.push({ text, confidence })
+      }
+    }
+    colorMat.delete()
+
+    const linesWeight = lines.flatMap(([p0, p1]) => ([p0.weight, p1.weight]))
+    const results = {
+      elapsed,
+      matches: matches!.length,
+      clusters: clusteredMatches.length,
+      linesMin: Math.min(...linesWeight),
+      linesMax: Math.max(...linesWeight),
+      recognized: recognizedLines
+    }
+    // console.log(results)
+    return results
+  }
+}

main/src/vision/link-main.ts

@@ -0,0 +1,42 @@
+import { Worker } from 'worker_threads'
+import * as Comlink from 'comlink'
+import nodeEndpoint from 'comlink/dist/umd/node-adapter'
+import type { WorkerAPI } from './link-worker'
+import type { ImageData } from './utils'
+import { app } from 'electron'
+import path from 'path'
+
+export class OcrWorker {
+  private binDir = path.join(app.getPath('userData'), 'apt-data/cv-ocr')
+  private api: Comlink.Remote<WorkerAPI>
+  private lang = ''
+
+  private constructor () {
+    const worker = new Worker(__dirname + '/vision.js')
+    this.api = Comlink.wrap<WorkerAPI>(nodeEndpoint(worker))
+  }
+
+  static async create () {
+    const worker = new OcrWorker()
+    try {
+      await worker.api.init(worker.binDir)
+    } catch {}
+    return worker
+  }
+
+  async updateOptions (lang: string) {
+    try {
+      if (lang !== this.lang) {
+        await this.api.changeLanguage(lang, this.binDir)
+      }
+    } catch {} finally {
+      this.lang = lang
+    }
+  }
+
+  async findHeistGems (image: ImageData) {
+    const result = await this.api.findHeistGems(
+      Comlink.transfer(image, [image.data.buffer]))
+    return result
+  }
+}

main/src/vision/link-worker.ts

@@ -0,0 +1,28 @@
+import { parentPort } from 'worker_threads'
+import * as Comlink from 'comlink'
+import nodeEndpoint from 'comlink/dist/umd/node-adapter'
+import * as Bindings from './wasm-bindings'
+import { HeistGemFinder } from './HeistGemFinder'
+import { ImageData } from './utils'
+
+let _heistGems: HeistGemFinder
+let _changeLangPromise = Promise.resolve()
+
+const WorkerBody = {
+  async init (binDir: string) {
+    await Bindings.init(binDir)
+    _heistGems = await HeistGemFinder.create(binDir)
+  },
+  async changeLanguage (lang: string, binDir: string) {
+    await _changeLangPromise
+    _changeLangPromise = Bindings.changeLanguage(lang, binDir)
+    await _changeLangPromise
+  },
+  async findHeistGems (screenshot: ImageData) {
+    await _changeLangPromise
+    return _heistGems.ocrScreenshot(screenshot)
+  }
+}
+Comlink.expose(WorkerBody, nodeEndpoint(parentPort!))
+
+export type WorkerAPI = Comlink.Remote<typeof WorkerBody>

main/src/vision/utils.ts

@@ -0,0 +1,75 @@
+import { cv } from './wasm-bindings'
+
+export interface ImageData {
+  width: number
+  height: number
+  data: Uint8Array
+}
+
+export interface WeightedPoint {
+  x: number
+  y: number
+  weight: number
+}
+
+export type LinePoints = [WeightedPoint, WeightedPoint]
+
+export function findNonZeroWeights (matchResult: any): WeightedPoint[] {
+  const locations = new cv.Mat()
+  cv.findNonZero(matchResult, locations)
+  const weights = Array<WeightedPoint>(locations.rows)
+  for (let i = 0; i < locations.rows; ++i) {
+    const x = locations.intAt(i, 0)
+    const y = locations.intAt(i, 1)
+    const weight = matchResult.floatAt(y, x)
+    weights[i] = { x, y, weight }
+  }
+  locations.delete()
+  return weights
+}
+
+export function groupWeightedPoints (weights: WeightedPoint[], radius: number): WeightedPoint[] {
+  // similar to non-maximum suppression
+  const maxWeighted: WeightedPoint[] = []
+  for (const point of weights) {
+    const closeIdx = maxWeighted.findIndex(maxPoint => {
+      const dist = Math.hypot(point.x - maxPoint.x, point.y - maxPoint.y)
+      return dist < radius
+    })
+    if (closeIdx === -1) {
+      maxWeighted.push(point)
+    } else if (point.weight > maxWeighted[closeIdx].weight) {
+      maxWeighted[closeIdx] = point
+    }
+  }
+  return maxWeighted
+}
+
+export function findLines (points: WeightedPoint[], yTolerance: number): LinePoints[] {
+  points.sort((a, b) => a.x - b.x)
+  const lines: LinePoints[] = []
+  for (let idxA = 0; idxA < points.length; ++idxA) {
+    for (let idxB = idxA + 1; idxB < points.length; ++idxB) {
+      const pointA = points[idxA]
+      const pointB = points[idxB]
+      if (Math.abs(pointA.y - pointB.y) > yTolerance) continue
+      lines.push([pointA, pointB])
+      break
+    }
+  }
+  return lines
+}
+
+export function hsvToU8 (h: number, s: number, v: number) {
+  return [
+    Math.round(h * 255 / 360),
+    Math.round(s * 255 / 100),
+    Math.round(v * 255 / 100),
+  ]
+}
+
+export function timeIt (syncFn: () => void): number {
+  const startTime = performance.now()
+  syncFn()
+  return performance.now() - startTime
+}

main/src/vision/wasm-bindings.ts

@@ -0,0 +1,56 @@
+import fs from 'fs/promises'
+import type { ImageData } from './utils'
+
+let tessModule: any
+export let tessApi: any
+export let cv: any
+
+const langMap = new Map([
+  ['en', 'eng'],
+  ['ru', 'rus'],
+  // ['cmn-Hant', 'chi_tra'],
+])
+
+export async function init (binDir: string) {
+  if (process.platform !== 'win32') {
+    // so far only tested on Windows with BGRA images
+    throw new Error('Unsupported platform')
+  }
+
+  const tessInstantiate = (await import('file://' + binDir + '/tesseract-core-simd.js')).default
+  tessModule = await tessInstantiate()
+  tessApi = new tessModule.TessBaseAPI()
+
+  const cvPromise = (await import('file://' + binDir + '/opencv.js')).default
+  cv = await cvPromise
+}
+
+export async function changeLanguage (lang: string, binDir: string) {
+  if (!langMap.has(lang)) {
+    throw new Error('Unsupported language')
+  }
+  lang = langMap.get(lang)!
+  const langData = await fs.readFile(binDir + `/${lang}.traineddata`)
+  tessModule.FS.writeFile(`${lang}.traineddata`, langData)
+  if (tessApi.Init(null, lang, tessModule.OEM_DEFAULT)) {
+    throw new Error('Could not initialize tesseract.')
+  }
+  tessModule.FS.unlink(`${lang}.traineddata`)
+}
+
+export function ocrSetImage (data: Uint8Array, width: number, height: number, bpp: number) {
+  const imgPtr = tessModule._malloc(data.byteLength)
+  tessModule.HEAPU8.set(data, imgPtr)
+  if (bpp === 0) {
+    tessApi.SetImage(imgPtr, width, height, 0, Math.ceil(width / 8))
+  } else {
+    tessApi.SetImage(imgPtr, width, height, bpp, width * bpp)
+  }
+  tessModule._free(imgPtr)
+}
+
+export function cvMatFromImage (img: ImageData) {
+  const mat = new cv.Mat(img.height, img.width, cv.CV_8UC4)
+  mat.data.set(img.data)
+  return mat
+}

main/src/windowing/GameWindow.ts

@@ -44,4 +44,8 @@ export class GameWindow extends EventEmitter {
       cb(e.hasAccess)
     })
   }
+
+  screenshot () {
+    return OverlayController.screenshot()
+  }
 }