koishi-plugin-ocr
Version:
本地 OCR 服务，调用 PaddleOCR 或 RapidOCR 进行识别（开坑中，一堆bug，勿用）
277 lines (275 loc) • 10.9 kB
JavaScript
var __create = Object.create;
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __getProtoOf = Object.getPrototypeOf;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __name = (target, value) => __defProp(target, "name", { value, configurable: true });
var __export = (target, all) => {
for (var name2 in all)
__defProp(target, name2, { get: all[name2], enumerable: true });
};
var __copyProps = (to, from, except, desc) => {
if (from && typeof from === "object" || typeof from === "function") {
for (let key of __getOwnPropNames(from))
if (!__hasOwnProp.call(to, key) && key !== except)
__defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
}
return to;
};
var __toESM = (mod, isNodeMode, target) => (target = mod != null ? __create(__getProtoOf(mod)) : {}, __copyProps(
// If the importer is in node compatibility mode or this is not an ESM
// file that has been converted to a CommonJS file using a Babel-
// compatible transform (i.e. "__esModule" has not been set), then set
// "default" to the CommonJS "module.exports" for node compatibility.
isNodeMode || !mod || !mod.__esModule ? __defProp(target, "default", { value: mod, enumerable: true }) : target,
mod
));
var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
// src/index.ts
var src_exports = {};
__export(src_exports, {
Config: () => Config,
apply: () => apply,
name: () => name
});
module.exports = __toCommonJS(src_exports);
var import_koishi = require("koishi");
var import_paddleocrjson = __toESM(require("paddleocrjson"));
var import_os = require("os");
var import_path = __toESM(require("path"));
var import_fs = __toESM(require("fs"));
var import_http = __toESM(require("http"));
var import_https = __toESM(require("https"));
var name = "ocr";
var defaultExePath = /* @__PURE__ */ __name(() => {
const osType = (0, import_os.platform)();
if (osType === "win32") return "PaddleOCR-json.exe";
else return "./PaddleOCR-json";
}, "defaultExePath");
var Config = import_koishi.Schema.object({
timeout: import_koishi.Schema.number().default(3e4).description("OCR识别超时时间（毫秒）"),
exePath: import_koishi.Schema.string().default(defaultExePath()).description("PaddleOCR-json可执行文件路径"),
cwd: import_koishi.Schema.string().description("PaddleOCR-json工作目录，默认为basedir下的data/PaddleOCR-json目录"),
args: import_koishi.Schema.array(String).default([]).description('PaddleOCR-json启动参数，例如 ["-port=9985", "-addr=loopback"]'),
debug: import_koishi.Schema.boolean().default(false).description("是否开启调试模式")
});
var Ocr = class extends import_koishi.Service {
constructor(ctx, config) {
super(ctx, "ocr");
this.config = config;
}
static {
__name(this, "Ocr");
}
static [import_koishi.Service.provide] = "ocr";
static inject = ["ocr"];
ocr;
pid = null;
initialized = false;
ocrInitComplete = false;
async start() {
if (!this.config.cwd) {
this.config.cwd = import_path.default.join(this.ctx.baseDir, "data/PaddleOCR-json");
}
try {
if (!import_fs.default.existsSync(this.config.cwd)) {
import_fs.default.mkdirSync(this.config.cwd, { recursive: true });
this.logger.info(`创建工作目录: ${this.config.cwd}`);
}
} catch (error) {
this.logger.error(`创建工作目录失败: ${error.message}`);
}
const exeFullPath = import_path.default.isAbsolute(this.config.exePath) ? this.config.exePath : import_path.default.join(this.config.cwd, this.config.exePath);
if (!import_fs.default.existsSync(exeFullPath)) {
this.logger.warn(`可执行文件不存在: ${exeFullPath}`);
}
try {
this.ocr = new import_paddleocrjson.default(this.config.exePath, this.config.args, {
cwd: this.config.cwd
}, this.config.debug);
if (this.config.debug) {
this.ocr.stdout.on("data", (chunk) => this.logger.info(chunk.toString()));
this.ocr.stderr.on("data", (data) => this.logger.warn(data.toString()));
}
this.ocr.on("error", (error) => {
this.logger.error(`OCR进程错误: ${error.message}`);
this.initialized = this.ocrInitComplete = false;
});
this.ocr.on("exit", (code) => {
this.logger.info(`OCR进程退出，退出码: ${code}`);
this.initialized = this.ocrInitComplete = false;
});
this.ocr.on("init", (pid, addr, port) => {
this.logger.info(`OCR初始化完成！PID: ${pid}, 地址: ${addr}, 端口: ${port}`);
this.pid = pid;
this.ocrInitComplete = true;
});
this.initialized = true;
return new Promise((resolve) => {
if (this.ocrInitComplete) {
resolve();
return;
}
this.ocr.once("init", () => resolve());
setTimeout(() => {
if (!this.ocrInitComplete) {
this.logger.warn("OCR服务初始化超时，但继续启动");
}
resolve();
}, 1e4);
});
} catch (error) {
this.logger.error(`OCR初始化失败: ${error.message}`);
this.initialized = false;
throw error;
}
}
async stop() {
if (this.ocr) {
try {
this.logger.info("正在停止OCR服务...");
this.ocr.terminate();
if (this.pid) {
await new Promise((resolve) => setTimeout(resolve, 1e3));
try {
this.logger.info(`OCR进程(PID:${this.pid})应该已终止`);
} catch (e) {
this.logger.warn(`尝试检查进程状态时出错: ${e.message}`);
}
}
this.ocr.removeAllListeners();
this.ocr = null;
this.initialized = false;
this.ocrInitComplete = false;
this.pid = null;
this.logger.info("OCR服务已停止并清理完成");
} catch (error) {
this.logger.warn(`终止OCR进程时发生错误: ${error.message}`);
}
}
}
isReady() {
return this.initialized && this.ocr && this.ocrInitComplete;
}
/**
* 将Element元素转换为统一的base64格式
*/
async getImageBase64(imageElement) {
const url = "url" in imageElement && typeof imageElement.url === "string" ? imageElement.url : imageElement?.type === "image" && imageElement?.attrs?.url ? imageElement.attrs.url : null;
if (!url) throw new Error("无效的图像元素");
return this.fetchImageFromUrl(url);
}
// 辅助方法：从URL获取图像并转为base64
async fetchImageFromUrl(imageUrl) {
if (imageUrl.startsWith("data:image") && imageUrl.includes("base64,")) {
return imageUrl;
}
if (imageUrl.startsWith("http://") || imageUrl.startsWith("https://")) {
return new Promise((resolve, reject) => {
const isHttps = imageUrl.startsWith("https://");
const requestModule = isHttps ? import_https.default : import_http.default;
const options = {
headers: {
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
},
timeout: 1e4
};
const req = requestModule.get(imageUrl, options, (res) => {
if (res.statusCode !== 200) {
reject(new Error(`请求失败，状态码: ${res.statusCode}`));
return;
}
const chunks = [];
res.on("data", (chunk) => chunks.push(Buffer.from(chunk)));
res.on("end", () => {
const buffer = Buffer.concat(chunks);
const base64 = buffer.toString("base64");
const mimeType = res.headers["content-type"] || "image/png";
resolve(`data:${mimeType};base64,${base64}`);
});
});
req.on("error", (err) => {
reject(new Error(`获取图像失败: ${err.message}`));
});
req.on("timeout", () => {
req.destroy();
reject(new Error("请求超时"));
});
req.end();
});
}
try {
const imagePath = import_path.default.isAbsolute(imageUrl) ? imageUrl : import_path.default.join(this.config.cwd, imageUrl);
if (!import_fs.default.existsSync(imagePath)) {
throw new Error(`文件不存在: ${imagePath}`);
}
const imageBuffer = import_fs.default.readFileSync(imagePath);
const base64 = imageBuffer.toString("base64");
const ext = import_path.default.extname(imagePath).toLowerCase();
const mimeType = ext === ".png" ? "image/png" : ext === ".jpg" || ext === ".jpeg" ? "image/jpeg" : ext === ".gif" ? "image/gif" : "image/png";
return `data:${mimeType};base64,${base64}`;
} catch (error) {
throw new Error(`图像处理失败: ${error.message}`);
}
}
async recognizeText(imageElement) {
if (!this.isReady()) {
if (this.initialized && this.ocr && !this.ocrInitComplete) {
await new Promise((resolve) => setTimeout(resolve, 5e3));
if (!this.isReady()) {
throw new Error("OCR服务未就绪");
}
} else {
throw new Error("OCR服务未就绪");
}
}
try {
const base64Image = await this.getImageBase64(imageElement);
const result = await this.ocr.flush({
image_base64: base64Image,
timeout: this.config.timeout
});
if (Array.isArray(result)) {
return result.map((item) => item.text).join("\n");
}
return JSON.stringify(result);
} catch (error) {
this.logger.error(`OCR识别失败: ${error.message}`);
throw new Error("OCR识别失败: " + error.message);
}
}
};
function apply(ctx, config) {
ctx.plugin(Ocr, config);
ctx.command("ocr [image:image]", "图像文字识别").option("image", "-i <url:string> 指定图片URL", { authority: 1 }).action(async ({ options, session }) => {
if (!ctx.ocr || !ctx.ocr.isReady()) {
return "OCR服务未就绪，请联系管理员检查配置";
}
let imageElement = null;
if (options.image) {
imageElement = (0, import_koishi.h)("image", { url: options.image });
} else if (session.content) {
const images = import_koishi.h.parse(session.content).filter((node) => node.type === "image");
if (images.length > 0) {
imageElement = images[0];
}
}
if (!imageElement) {
return "请提供图片";
}
try {
const result = await ctx.ocr.recognizeText(imageElement);
return result || "未识别到任何文字";
} catch (error) {
return `识别失败：${error.message}`;
}
});
}
__name(apply, "apply");
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
Config,
apply,
name
});