Something went wrong. Try again.
A privacy-first, self-hosted, fully open source personal knowledge management software, written in typescript and golang. (PERSONAL FORK)
Something went wrong. Try again.
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476// SiYuan - Refactor your thinking// Copyright (c) 2020-present, b3log.org//// This program is free software: you can redistribute it and/or modify// it under the terms of the GNU Affero General Public License as published by// the Free Software Foundation, either version 3 of the License, or// (at your option) any later version.//// This program is distributed in the hope that it will be useful,// but WITHOUT ANY WARRANTY; without even the implied warranty of// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the// GNU Affero General Public License for more details.//// You should have received a copy of the GNU Affero General Public License// along with this program. If not, see <https://www.gnu.org/licenses/>.
package util
import ( "bytes" "context" "errors" "fmt" "os" "os/exec" "path/filepath" "runtime/debug" "strconv" "strings" "sync" "sync/atomic" "time"
"github.com/88250/go-humanize" "github.com/88250/gulu" "github.com/88250/lute/ast" "github.com/88250/lute/html" "github.com/siyuan-note/filelock" "github.com/siyuan-note/logging")
var ( TesseractBin = "tesseract" TesseractEnabled bool TesseractMaxSize = 2 * 1000 * uint64(1000) TesseractLangs []string
assetsTexts = map[string]string{} assetsTextsLock = sync.Mutex{} assetsTextsChanged = atomic.Bool{})
func CleanNotExistAssetsTexts() { assetsTextsLock.Lock() defer assetsTextsLock.Unlock()
assetsPath := GetDataAssetsAbsPath() var toRemoves []string for asset, _ := range assetsTexts { assetAbsPath := strings.TrimPrefix(asset, "assets") assetAbsPath = filepath.Join(assetsPath, assetAbsPath) if !filelock.IsExist(assetAbsPath) { toRemoves = append(toRemoves, asset) } }
for _, asset := range toRemoves { delete(assetsTexts, asset) assetsTextsChanged.Store(true) } return}
func LoadAssetsTexts() { assetsPath := GetDataAssetsAbsPath() assetsTextsPath := filepath.Join(assetsPath, "ocr-texts.json") if !filelock.IsExist(assetsTextsPath) { return }
start := time.Now() data, err := filelock.ReadFile(assetsTextsPath) if err != nil { logging.LogErrorf("read assets texts failed: %s", err) return }
assetsTextsLock.Lock() if err = gulu.JSON.UnmarshalJSON(data, &assetsTexts); err != nil { logging.LogErrorf("unmarshal assets texts failed: %s", err) if err = filelock.Remove(assetsTextsPath); err != nil { logging.LogErrorf("removed corrupted assets texts failed: %s", err) } return } assetsTextsLock.Unlock() debug.FreeOSMemory()
if elapsed := time.Since(start).Seconds(); 2 < elapsed { logging.LogWarnf("read assets texts [%s] to [%s], elapsed [%.2fs]", humanize.BytesCustomCeil(uint64(len(data)), 2), assetsTextsPath, elapsed) } return}
func SaveAssetsTexts() { if !assetsTextsChanged.Load() { return }
start := time.Now()
assetsPath := GetDataAssetsAbsPath() assetsTextsPath := filepath.Join(assetsPath, "ocr-texts.json")
assetsTextsLock.Lock() // OCR 功能未开启且 ocr-texts.json 不存在时，如果 assetsTexts 为空则不创建文件 if !TesseractEnabled && !filelock.IsExist(assetsTextsPath) && 0 == len(assetsTexts) { assetsTextsLock.Unlock() assetsTextsChanged.Store(false) return } data, err := gulu.JSON.MarshalIndentJSON(assetsTexts, "", " ") if err != nil { logging.LogErrorf("marshal assets texts failed: %s", err) assetsTextsLock.Unlock() return } assetsTextsLock.Unlock()
if err = filelock.WriteFile(assetsTextsPath, data); err != nil { logging.LogErrorf("write assets texts failed: %s", err) return } debug.FreeOSMemory()
if elapsed := time.Since(start).Seconds(); 2 < elapsed { logging.LogWarnf("save assets texts [size=%s] to [%s], elapsed [%.2fs]", humanize.BytesCustomCeil(uint64(len(data)), 2), assetsTextsPath, elapsed) }
assetsTextsChanged.Store(false)}
func SetAssetText(asset, text string) { assetsTextsLock.Lock() oldText, ok := assetsTexts[asset] assetsTexts[asset] = text assetsTextsLock.Unlock() if !ok || oldText != text { assetsTextsChanged.Store(true) }}
func ExistsAssetText(asset string) (ret bool) { assetsTextsLock.Lock() _, ret = assetsTexts[asset] assetsTextsLock.Unlock() return}
func OcrAsset(asset string) (ret []map[string]interface{}, err error) { if !TesseractEnabled { err = fmt.Errorf(Langs[Lang][266]) return }
assetsPath := GetDataAssetsAbsPath() assetAbsPath := strings.TrimPrefix(asset, "assets") assetAbsPath = filepath.Join(assetsPath, assetAbsPath) ret = Tesseract(assetAbsPath) assetsTextsLock.Lock() ocrText := GetOcrJsonText(ret) assetsTexts[asset] = ocrText assetsTextsLock.Unlock() if "" != ocrText { assetsTextsChanged.Store(true) } return}
func GetAssetText(asset string) (ret string) { assetsTextsLock.Lock() ret = assetsTexts[asset] assetsTextsLock.Unlock() return}
func RemoveAssetText(asset string) { assetsTextsLock.Lock() delete(assetsTexts, asset) assetsTextsLock.Unlock() assetsTextsChanged.Store(true)}
var tesseractExts = []string{ ".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".gif", ".webp", ".pbm", ".pgm", ".ppm", ".pnm",}
func IsTesseractExtractable(p string) bool { lowerName := strings.ToLower(p) for _, ext := range tesseractExts { if strings.HasSuffix(lowerName, ext) { return true } } return false}
// tesseractOCRLock 用于 Tesseract OCR 加锁串行执行提升稳定性 https://github.com/siyuan-note/siyuan/issues/7265var tesseractOCRLock = sync.Mutex{}
func Tesseract(imgAbsPath string) (ret []map[string]interface{}) { if ContainerStd != Container || !TesseractEnabled { return }
defer logging.Recover() tesseractOCRLock.Lock() defer tesseractOCRLock.Unlock()
if !IsTesseractExtractable(imgAbsPath) { return }
info, err := os.Stat(imgAbsPath) if err != nil { return }
if TesseractMaxSize < uint64(info.Size()) { return }
defer logging.Recover()
timeout := 7000 timeoutEnv := os.Getenv("SIYUAN_TESSERACT_TIMEOUT") if "" != timeoutEnv { if timeoutParsed, parseErr := strconv.Atoi(timeoutEnv); nil == parseErr { timeout = timeoutParsed } else { logging.LogWarnf("parse tesseract timeout [%s] failed: %s", timeoutEnv, parseErr) } } ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Millisecond) defer cancel()
cmd := exec.CommandContext(ctx, TesseractBin, "-c", "debug_file=/dev/null", imgAbsPath, "stdout", "-l", strings.Join(TesseractLangs, "+"), "tsv") gulu.CmdAttr(cmd) output, err := cmd.CombinedOutput() if errors.Is(ctx.Err(), context.DeadlineExceeded) { logging.LogWarnf("tesseract [path=%s, size=%d] timeout [%dms]", imgAbsPath, info.Size(), timeout) return }
if err != nil { logging.LogWarnf("tesseract [path=%s, size=%d] failed: %s", imgAbsPath, info.Size(), err) return }
tsv := string(output) //logging.LogInfof("tesseract [path=%s] success [%s]", imgAbsPath, tsv)
// 按行分割 TSV 数据 tsv = strings.ReplaceAll(tsv, "\r", "") lines := strings.Split(tsv, "\n")
// 解析 TSV 数据 跳过标题行，从第二行开始处理 for _, line := range lines[1:] { if line == "" { continue // 跳过空行 } // 分割每列数据 fields := strings.Split(line, "\t") // 将字段名和字段值映射到一个 map 中 dataMap := make(map[string]interface{}) headers := strings.Split(lines[0], "\t") for i, header := range headers { if i < len(fields) { dataMap[header] = fields[i] } else { dataMap[header] = "" } } ret = append(ret, dataMap) }
tsv = RemoveInvalid(tsv) tsv = RemoveRedundantSpace(tsv) msg := fmt.Sprintf("OCR [%s] [%s]", html.EscapeString(info.Name()), html.EscapeString(GetOcrJsonText(ret))) PushStatusBar(msg) return}
// GetOcrJsonText 提取并连接所有 text 字段的函数func GetOcrJsonText(jsonData []map[string]interface{}) (ret string) { for _, dataMap := range jsonData { // 检查 text 字段是否存在 if text, ok := dataMap["text"]; ok { // 确保 text 是字符串类型 if textStr, ok := text.(string); ok { ret += " " + strings.ReplaceAll(textStr, "\r", "") } } } ret = RemoveInvalid(ret) ret = RemoveRedundantSpace(ret) return ret}
var tesseractInited = atomic.Bool{}
func WaitForTesseractInit() { for { if tesseractInited.Load() { return } time.Sleep(time.Second) }}
func InitTesseract() { ver := getTesseractVer() if "" == ver { tesseractInited.Store(true) return }
langs := getTesseractLangs() if 1 > len(langs) { logging.LogWarnf("no tesseract langs found, disabling tesseract-ocr") TesseractEnabled = false tesseractInited.Store(true) return }
maxSizeVal := os.Getenv("SIYUAN_TESSERACT_MAX_SIZE") if "" != maxSizeVal { if maxSize, parseErr := strconv.ParseUint(maxSizeVal, 10, 64); nil == parseErr { TesseractMaxSize = maxSize } }
// Supports via environment var `SIYUAN_TESSERACT_ENABLED=false` to close OCR https://github.com/siyuan-note/siyuan/issues/9619 if enabled := os.Getenv("SIYUAN_TESSERACT_ENABLED"); "" != enabled { if enabledBool, parseErr := strconv.ParseBool(enabled); nil == parseErr { TesseractEnabled = enabledBool if !enabledBool { logging.LogInfof("tesseract-ocr disabled by env") tesseractInited.Store(true) return } } }
TesseractLangs = filterTesseractLangs(langs) logging.LogInfof("tesseract-ocr enabled [ver=%s, maxSize=%s, langs=%s]", ver, humanize.BytesCustomCeil(TesseractMaxSize, 2), strings.Join(TesseractLangs, "+")) tesseractInited.Store(true)}
func filterTesseractLangs(langs []string) (ret []string) { ret = []string{}
envLangsVal := os.Getenv("SIYUAN_TESSERACT_LANGS") if "" != envLangsVal { envLangs := strings.Split(envLangsVal, "+") for _, lang := range langs { if gulu.Str.Contains(lang, envLangs) { ret = append(ret, lang) } } } else { for _, lang := range langs { if "eng" == lang || strings.HasPrefix(lang, "chi") || "fra" == lang || "spa" == lang || "deu" == lang || "rus" == lang || "jpn" == lang || "osd" == lang { ret = append(ret, lang) } } } return ret}
func getTesseractVer() (ret string) { if ContainerStd != Container { return }
cmd := exec.Command(TesseractBin, "--version") gulu.CmdAttr(cmd) data, err := cmd.CombinedOutput() if err != nil { errMsg := strings.ToLower(err.Error()) if strings.Contains(errMsg, "executable file not found") || strings.Contains(errMsg, "no such file or directory") { // macOS 端 Tesseract OCR 安装后不识别 https://github.com/siyuan-note/siyuan/issues/7107 TesseractBin = "/usr/local/bin/tesseract" cmd = exec.Command(TesseractBin, "--version") gulu.CmdAttr(cmd) data, err = cmd.CombinedOutput() if err != nil { errMsg = strings.ToLower(err.Error()) if strings.Contains(errMsg, "executable file not found") || strings.Contains(errMsg, "no such file or directory") { TesseractBin = "/opt/homebrew/bin/tesseract" cmd = exec.Command(TesseractBin, "--version") gulu.CmdAttr(cmd) data, err = cmd.CombinedOutput() } } } } if err != nil { return }
if strings.HasPrefix(string(data), "tesseract ") { parts := bytes.Split(data, []byte("\n")) if 0 < len(parts) { ret = strings.TrimPrefix(string(parts[0]), "tesseract ") ret = strings.TrimSpace(ret) TesseractEnabled = true } return } return}
func getTesseractLangs() (ret []string) { if !TesseractEnabled { return nil }
cmd := exec.Command(TesseractBin, "--list-langs") gulu.CmdAttr(cmd) data, err := cmd.CombinedOutput() if err != nil { return nil }
parts := bytes.Split(data, []byte("\n")) if 0 < len(parts) { parts = parts[1:] } for _, part := range parts { part = bytes.TrimSpace(part) if 0 == len(part) { continue } ret = append(ret, string(part)) } return}
var ( NodeOCRQueue []string NodeOCRQueueLock = sync.Mutex{})
func PushNodeOCRQueue(n *ast.Node) { if nil == n { return }
NodeOCRQueueLock.Lock() defer NodeOCRQueueLock.Unlock() NodeOCRQueue = append(NodeOCRQueue, n.ID)}