Merge branch 'feat/ocr' into feat/ocr-translate

packages/shared/config/constant.ts

@@ -197,6 +197,12 @@ export enum FeedUrl {
   GITHUB_LATEST = 'https://github.com/CherryHQ/cherry-studio/releases/latest/download'
 }
 
+export const tesseractLangs = ['chi_sim', 'chi_tra', 'eng']
+export enum TesseractLangsDownloadUrl {
+  CN = 'https://gitcode.com/beyondkmp/tessdata/releases/download/4.1.0/',
+  GLOBAL = 'https://github.com/tesseract-ocr/tessdata/raw/main/'
+}
+
 export enum UpgradeChannel {
   LATEST = 'latest', // 最新稳定版本
   RC = 'rc', // 公测版本

src/main/ipc.ts

@@ -30,7 +30,7 @@ import { openTraceWindow, setTraceWindowTitle } from './services/NodeTraceServic
 import NotificationService from './services/NotificationService'
 import * as NutstoreService from './services/NutstoreService'
 import ObsidianVaultService from './services/ObsidianVaultService'
-import { ipcOcr } from './services/ocr/OcrService'
+import { ocrService } from './services/ocr/OcrService'
 import { proxyManager } from './services/ProxyManager'
 import { pythonService } from './services/PythonService'
 import { FileServiceManager } from './services/remotefile/FileServiceManager'
@@ -713,5 +713,5 @@ export function registerIpc(mainWindow: BrowserWindow, app: Electron.App) {
   ipcMain.handle(IpcChannel.CodeTools_Run, codeToolsService.run)
 
   // OCR
-  ipcMain.handle(IpcChannel.OCR_ocr, ipcOcr)
+  ipcMain.handle(IpcChannel.OCR_ocr, (_, ...args: Parameters<typeof ocrService.ocr>) => ocrService.ocr(...args))
 }

src/main/services/ocr/OcrService.ts

@@ -1,91 +1,30 @@
-import { loggerService } from '@logger'
-import { MB } from '@shared/config/constant'
-import {
-  ImageFileMetadata,
-  ImageOcrProvider,
-  isBuiltinOcrProvider,
-  isImageFile,
-  isImageOcrProvider,
-  OcrProvider,
-  OcrResult,
-  SupportedOcrFile
-} from '@types'
-import { statSync } from 'fs'
-import { readFile } from 'fs/promises'
+import { BuiltinOcrProviderIds, FileMetadata, OcrProvider, OcrResult, SupportedOcrFile } from '@types'
 
-import { getTesseractWorker } from './tesseract/TesseractService'
+import { tesseractService } from './tesseract/TesseractService'
 
-const logger = loggerService.withContext('main:OcrService')
+type OcrHandler = (file: FileMetadata) => Promise<OcrResult>
 
-/**
- * ocr by tesseract
- * @param file image file or base64 string
- * @returns ocr result
- * @throws {Error}
- */
-const tesseractOcr = async (file: ImageFileMetadata | string): Promise<Tesseract.RecognizeResult> => {
-  try {
-    const worker = await getTesseractWorker()
-    let ret: Tesseract.RecognizeResult
-    if (typeof file === 'string') {
-      ret = await worker.recognize(file)
-    } else {
-      const stat = statSync(file.path)
-      if (stat.size > 50 * MB) {
-        throw new Error('This image is too large (max 50MB)')
-      }
-      const buffer = await readFile(file.path)
-      ret = await worker.recognize(buffer)
+export class OcrService {
+  private registry: Map<string, OcrHandler> = new Map()
+
+  register(providerId: string, handler: OcrHandler): void {
+    this.registry.set(providerId, handler)
+  }
+
+  unregister(providerId: string): void {
+    this.registry.delete(providerId)
+  }
+
+  public async ocr(file: SupportedOcrFile, provider: OcrProvider): Promise<OcrResult> {
+    const handler = this.registry.get(provider.id)
+    if (!handler) {
+      throw new Error(`Provider ${provider.id} is not registered`)
     }
-    return ret
-  } catch (e) {
-    logger.error('Failed to ocr with tesseract.', e as Error)
-    throw e
+    return handler(file)
   }
 }
 
-/**
- * ocr image file
- * @param file image file
- * @param provider ocr provider that supports image ocr
- * @returns ocr result
- * @throws {Error}
- */
-const imageOcr = async (file: ImageFileMetadata, provider: ImageOcrProvider): Promise<OcrResult> => {
-  if (isBuiltinOcrProvider(provider)) {
-    if (provider.id === 'tesseract') {
-      const result = await tesseractOcr(file)
-      return { text: result.data.text }
-    } else {
-      throw new Error(`Unsupported built-in ocr provider: ${provider.id}`)
-    }
-  }
-  throw new Error(`Provider ${provider.id} is not supported.`)
-}
+export const ocrService = new OcrService()
 
-/**
- * ocr a file
- * @param file any supported file
- * @param provider ocr provider
- * @returns ocr result
- * @throws {Error}
- */
-export const ocr = async (file: SupportedOcrFile, provider: OcrProvider): Promise<OcrResult> => {
-  if (isImageFile(file) && isImageOcrProvider(provider)) {
-    return imageOcr(file, provider)
-  } else {
-    throw new Error(`File type and provider capability is not matched, otherwise one of them is not supported.`)
-  }
-}
-
-/**
- * ocr a file
- * @param _ ipc event
- * @param file any supported file
- * @param provider ocr provider
- * @returns ocr result
- * @throws {Error}
- */
-export const ipcOcr = async (_: Electron.IpcMainInvokeEvent, ...args: Parameters<typeof ocr>) => {
-  return ocr(...args)
-}
+// Register built-in providers
+ocrService.register(BuiltinOcrProviderIds.tesseract, tesseractService.ocr.bind(tesseractService))

src/main/services/ocr/tesseract/TesseractService.ts

@@ -1,10 +1,14 @@
 import { loggerService } from '@logger'
+import { getIpCountry } from '@main/utils/ipService'
+import { MB, TesseractLangsDownloadUrl } from '@shared/config/constant'
+import { FileMetadata, ImageFileMetadata, isImageFile, OcrResult } from '@types'
+import { app } from 'electron'
+import fs from 'fs'
+import path from 'path'
 import Tesseract, { createWorker } from 'tesseract.js'
 
 const logger = loggerService.withContext('TesseractService')
 
-let worker: Tesseract.Worker | null = null
-
 // const languageCodeMap: Record<string, string> = {
 //   'af-za': 'afr',
 //   'am-et': 'amh',
@@ -110,20 +114,65 @@ let worker: Tesseract.Worker | null = null
 //   'yi-us': 'yid'
 // }
 
-export const getTesseractWorker = async (): Promise<Tesseract.Worker> => {
-  if (!worker) {
-    // for now, only support limited languages
-    worker = await createWorker(['chi_sim', 'chi_tra', 'eng'], undefined, {
-      // langPath: getCacheDir(),
-      logger: (m) => logger.debug('From worker', m)
-    })
+export class TesseractService {
+  private worker: Tesseract.Worker | null = null
+
+  async getWorker(): Promise<Tesseract.Worker> {
+    if (!this.worker) {
+      // for now, only support limited languages
+      this.worker = await createWorker(['chi_sim', 'chi_tra', 'eng'], undefined, {
+        langPath: await this._getLangPath(),
+        cachePath: await this._getCacheDir(),
+        gzip: false,
+        logger: (m) => logger.debug('From worker', m)
+      })
+    }
+    return this.worker
+  }
+
+  async imageOcr(file: ImageFileMetadata): Promise<OcrResult> {
+    const worker = await this.getWorker()
+    const stat = await fs.promises.stat(file.path)
+    if (stat.size > 50 * MB) {
+      throw new Error('This image is too large (max 50MB)')
+    }
+    const buffer = await fs.promises.readFile(file.path)
+    const result = await worker.recognize(buffer)
+    return { text: result.data.text }
+  }
+
+  async ocr(file: FileMetadata): Promise<OcrResult> {
+    if (!isImageFile(file)) {
+      throw new Error('Only image files are supported currently')
+    }
+    return this.imageOcr(file)
+  }
+
+  private async _getLangPath(): Promise<string> {
+    const country = await getIpCountry()
+    return country.toLowerCase() === 'cn' ? TesseractLangsDownloadUrl.CN : TesseractLangsDownloadUrl.GLOBAL
+  }
+
+  private async _getCacheDir(): Promise<string> {
+    const cacheDir = path.join(app.getPath('userData'), 'tesseract')
+    // use access to check if the directory exists
+    if (
+      !(await fs.promises
+        .access(cacheDir, fs.constants.F_OK)
+        .then(() => true)
+        .catch(() => false))
+    ) {
+      await fs.promises.mkdir(cacheDir, { recursive: true })
+    }
+    return cacheDir
+  }
+
+  async dispose(): Promise<void> {
+    if (this.worker) {
+      await this.worker.terminate()
+      this.worker = null
+    }
   }
-  return worker
 }
 
-export const disposeTesseractWorker = async () => {
-  if (worker) {
-    await worker.terminate()
-    worker = null
-  }
-}
+export const tesseractService = new TesseractService()

src/renderer/src/store/migrate.ts

@@ -2182,7 +2182,7 @@ const migrateConfig = {
       state.ocr.imageProvider = DEFAULT_OCR_PROVIDER.image
       return state
     } catch (error) {
-      logger.error('migrate 136 error', error as Error)
+      logger.error('migrate 137 error', error as Error)
       return state
     }
   }