EAimTY 的博客
肥宅
首页
关于
主题
LeetCode 笔记
用 Rust 写一个用于 OCR 的 Telegram Bot
Aug 09, 2021

最近一段时间在学 Rust,想写一些简单的小工具来巩固一下对语言基本特性的认识。之前用其它语言写过 Telegram bot,所以我想通过用 Rust 写 Telegram bot 来更深刻地了解 Rust 的独特之处。

Rust 相比于其它流行的语言网络上的资源比较少,中文内容更是寥寥无几。虽然我的 Rust 连入门都谈不上,代码里可能会有不少不合理的地方,但是还是想把过程记录一下,供他人参考,希望可以为 Rust 社区做一些微不足道的贡献,也算是抛砖引玉吧。

成品 bot 在 EAimTY/eaimty_bot。这是个用来练手的小项目,不止有 OCR 功能,还有一些其它的杂七杂八的功能,其中 OCR 部分在 这里

仅仅写这样一个简单的 OCR bot 并不能了解到太多 Rust 优秀的地方,毕竟写的只是像脚本一样的线性处理过程。但总归是写了点东西,对之后更深入的学习还是有点帮助的。

设计逻辑

先来设计一下这个 OCR Bot 的逻辑:

用户通过 /ocr 命令触发流程开始,Bot 发送一条带按钮的消息要求选择 OCR 的目标语言,在用户选择目标语言后显示选择的语言并提供“重新选择”选项,Bot 接受之后用户发送的消息,如果收到的是图片就用 Tesseract 进行处理,最后发送 OCR 结果,整个流程结束。
在整个流程中需要确保只有触发流程开始的用户才可以通过点击按钮选择 OCR 目标语言,并且只接受触发流程开始的用户发送的图片。

由于整个流程中需要接受并处理 3 种不同的请求,单个 handler 肯定是没法完成的,所以我们将整个流程分为 3 个部分:/ocr 命令触发,处理用户点击按钮选择语言,处理用户发送的图片,因此需要引入 session 保存状态。

选择框架与库

用 Rust 实现的 Bot API 框架列表可以在 Telegram 官方的 Bot Code Examples#Rust 中找到。
我选择的框架是 carapax。虽然这个框架的使用者不如 teloxide 或是 telebot 多,但我觉得它用起来更简单方便。

对于 OCR,最好的开源实现肯定是 Tesseract。Tesseract 在 Rust 下的 bindings 有 tesseract-systesseractleptess,我这里用的是 leptess。

编译 Bot 程序的机器上必须装 Tesseract 和 Leptonica,否则会编译失败。运行 Bot 程序的机器上必须装 Tesseract 和需要 OCR 的 Tesseract 数据包,比如 eng、jpn、chi_sim 和 chi_tra。

有了 Telegram Bot API 框架和 OCR 库,就可以开始动手了。

开始动手

首先去 Telegram 找 @BotFather 新建一个机器人,拿到 API Token。

然后新建一个 Cargo 项目,参照 carapax 的 examples 把大体框架写出来:

struct Context {
    api: Api,
    session_manager: SessionManager<FilesystemBackend>,
}

#[tokio::main]
async fn main() {
    // 在这里设置 API TOKEN
    let mut config = Config::new("API_TOKEN");
    // 需要代理的话取消下面一行的注释,在这里设置代理
    // config = config.proxy("PROXY").expect("Failed to set proxy");
    let api = Api::new(config).expect("Failed to create API");
    // 创建临时目录供 session 存储信息
    let tmpdir = tempdir().expect("Failed to create temp directory");
    let backend = FilesystemBackend::new(tmpdir.path());
    // session 的 GC 间隔
    let gc_period = Duration::from_secs(3);
    // 单个 session 的最大时长
    let session_lifetime = Duration::from_secs(86400);
    // 创建一个用来 session GC 的线程
    let mut collector = SessionCollector::new(backend.clone(), gc_period, session_lifetime);
    tokio::spawn(async move {
        collector.run().await
    });
    // 建立 bot 的 dispatcher,把 Context 传给之后的每一个 handler
    let mut dispatcher = Dispatcher::new(Context {
        api: api.clone(),
        session_manager: SessionManager::new(backend),
    });
    // 在下面添加 handler
    dispatcher.add_handler(ocr_command_handler);
    dispatcher.add_handler(ocr_image_handler);
    dispatcher.add_handler(ocr_inlinekeyboard_handler);
    // 以 longpoll 模式运行,适合本机调试
    LongPoll::new(api, dispatcher).run().await;
    // 以 webhook 模式运行,适合生产环境
    // webhook::run_server(([127, 0, 0, 1], PORT), "/", dispatcher).await.expect("Failed to run webhook server");
}

框架会把 &Context 传给每一个 handler, 这样就实现了 api 和 session backend 的传递。

大体的框架写好了,下面可以来写 OCR 功能了。首先需要一个列表来存储 bot 支持的 OCR 语言。

OCR 支持语言列表的存储

建立语言列表存储结构体不单是为了存储语言列表,还是为了实现可以直接调用的获取语言选择按钮列表之类的方法。

我们要用到的 OCR 后端:Tesseract 支持的语言非常多,这里选英文、中文和日语实现支持。
Tesseract 的语言包的名字是 engchi_simchi_trajpn 之类的语言名称缩写,直接把这些名字显示给用户很不友好,所以还需要存储语言的显示名称。既然要存储成对的值,哈希表应该是最好的选择,新建一个语言列表结构体 OcrLangs

struct OcrLangs<'a> {
    langs: HashMap<&'a str, &'a str>,
}

为什么要用 &str 而不是 String?因为这个 bot 不需要动态加载语言列表之类的这些功能,要支持什么语言直接写死进代码。
既然是“写死”,就意味着把语言包名称和显示名称写进二进制文件成为 &'static str,这样之后在调用的时候就不会出现需要在堆上内存里复制来复制去这类操作,同时也避免了所有权问题,一举两得。
由于用的全部都是 &str,所以需要显式地标出生命周期,这里直接用 <'a> 了。

然后来为这个结构体实现新建和添加语言方法:

impl<'a> OcrLangs<'a> {
    fn new() -> Self {
        Self {
            langs: HashMap::new(),
        }
    }

    // 添加 OCR 语言
    fn add(&mut self, lang: &'a str, name: &'a str) {
        self.langs.insert(lang, name);
    }
}

用语言数据包名作为哈希表键名,显示名称作为值。

下一个问题是,把实例化的语言列表存在哪里?由于这个 bot 的结构十分简单,这里就直接将列表存进全局变量了。
然而问题又来了,由于需要用 OcrLangs::new() 新建一个列表,再用 OcrLangs::add() 添加语言,这个全局变量需要是可变的,而 Rust 中 static mut 的修改是 unsafe 的。不仅如此,static 还无法初始化如 HashMap 需要堆来建立的结构,所以需要另寻出路。

lazy_static 是一个很常用的用来创建全局变量的 crate,这里用它提供的 lazy_static! 宏来惰性地初始化语言列表:

lazy_static! {
    static ref OCR_LANGS: OcrLangs<'static> = {
        let mut ocr_langs = OcrLangs::new();
        ocr_langs.add("eng", "English");
        ocr_langs.add("jpn", "日本語");
        ocr_langs.add("chi_sim", "简体中文");
        ocr_langs.add("chi_tra", "繁體中文");
        ocr_langs
    };
}

这样就完成了 OCR 支持语言列表的存储。只有语言列表没什么用,OCR 流程间的每一步都需要传递状态,所以下面来写 OCR 流程状态结构体,用来存储在 session 中。

P.S. 其实还可以把这个列表作为 Context 的成员,这样 dispatcher 就会将语言列表传给每一个 handler。这样做的好处在于可以做成从其它地方(如文件)加载列表,而且不需要再用全局变量了,全局变量在项目比较大的时候会造成负面影响;缺点在于不能用 &str 存储语言包名称和显示名称了,必须用 String

用来存储 OCR 流程状态的结构

定义一个用来存储 OCR 流程状态的结构 Ocr,很简单,只有一个成员 lang 存储用户指定的数据包名。lang 中的 Option 如果是 None 代表用户还没有选择语言,Some("LANG") 代表用户已经选择了目标语言 LANG.

#[derive(Serialize, Deserialize)]
struct Ocr {
    lang: Option<String>,
}

由于这个结构需要在 session 中存储(对于这个 bot 来讲就是存成文件),所以需要实现对它的序列化和反序列化,这里直接用 serde crate 中提供的 derive 宏 SerializeDeserialize

然后先为这个结构定义初始化方法:

impl Ocr {
    fn new() -> Self {
        Self { lang: None }
    }
}

好,这样用来存储的结构就都写完了,下面来写回应用户操作的 handlers。

命令触发

用户触发 /ocr 命令后,bot 回复一个带有语言选择按钮列表的消息供用户选择语言。之前已经定义过语言列表 OcrLangs,就为它定义一个返回语言选择按钮列表的方法。

这个框架用 SendMessagereply_markup 方法用来为消息附加按钮之类的功能,这里用 ReplyMarkup 中的 InlineKeyboardMarkup,通过一个 Vec<Vec<InlineKeyboardButton>> 来建立按钮列表。CallbackData 让每一个按钮被点击时返回特定的 String。

放进 OcrLangsimpl 中:

fn get_langs_keyboard(&self) -> InlineKeyboardMarkup {
    let mut keyboad: Vec<Vec<InlineKeyboardButton>> = Vec::new();
    for (lang, name) in &self.langs {
        keyboad.push(vec![InlineKeyboardButton::new(
            *name,
            InlineKeyboardButtonKind::CallbackData(format!("ocr-{}", lang)),
        )]);
    }
    InlineKeyboardMarkup::from(keyboad)
}

然后是通过 /ocr 命令触发的 handler:

#[handler(command = "/ocr")]
async fn ocr_command_handler(
    context: &Context,
    command: Command,
) -> Result<HandlerResult, ExecuteError> {
    let message = command.get_message();
    let chat_id = message.get_chat_id();
    if let Some(user) = message.get_user() {
        let user_id = user.id;
        // 在 session 中存储 OCR 状态
        let mut session = context
            .session_manager
            .get_session(SessionId::new(chat_id, user_id)).unwrap();
        session.set("ocr", &Ocr::new()).await.unwrap();
        // 发送语言选择信息
        let method = SendMessage::new(chat_id, "请选择 OCR 目标语言").reply_markup(
            ReplyMarkup::InlineKeyboardMarkup(OCR_LANGS.get_langs_keyboard()),
        );
        context.api.execute(method).await?;
    }
    Ok(HandlerResult::Stop)
}

暂时先用 unwrap 处理 Result,之后再完善错误处理。

这里通过 chat id 和 user id 得到 session 的 id,然后在 session 中新建一个键名为 "ocr" 的 Ocr 结构以供验证。

处理按钮点击

如果用户已经选择过了语言,为了防止是误点,还需要有一个“重新选择”按钮供用户“反悔”。为 OcrLangs 加一个方法:

fn get_reselect_keyboard(&self) -> InlineKeyboardMarkup {
    InlineKeyboardMarkup::from(vec![vec![InlineKeyboardButton::new(
        "重新选择",
        InlineKeyboardButtonKind::CallbackData(String::from("ocr-reselect")),
    )]])
}

用户点击按钮的操作要么是选择语言,要么是希望重新选择语言,所以把它抽象为一个枚举:

enum Operation {
    Select(String),
    Reselect,
}

用户选择语言是 Select("LANG"),希望重新选择则是 Reselect

Ocr 加上对应的操作:

fn get(&self) -> Option<String> {
    self.lang.clone()
}

fn set(&mut self, operation: &Operation) {
    match operation {
        Operation::Select(lang) => self.lang = Some(lang.to_string()),
        Operation::Reselect => self.lang = None,
    }
}

这里 get() 到的是语言数据包名称,显示给用户的语言名称还需要把处理一下这个结果,给 OcrLangs 添加一个方法:

// 获取语言的显示名称
fn get_lang_name(&self, lang: &str) -> &str {
    self.langs.get(lang).unwrap_or(&"")
}

用户点击按钮后 Telegram 会返回 CallbackQuery,其中包含之前设置的 CallbackData,所以需要通过匹配 CallbackData 得到用户的操作。

框架会将所有用户点击按钮的操作交给这个 handler 处理,不论是不是 OCR 操作相关的,所以还需要一个用来识别 CallbackQuery 的方法,给 OcrLangs 加一个方法:

fn try_parse_callback(&self, data: String) -> Option<Operation> {
    if data.starts_with("ocr-") {
        let mut data = data[4..].split('-');
        if let Some(lang) = data.next() {
            if let None = data.next() {
                if lang == "reselect" {
                    return Some(Operation::Reselect);
                } else if self.langs.contains_key(lang) {
                    return Some(Operation::Select(lang.to_string()));
                }
            }
        }
    }
    None
}

如果返回的是 Some(Operation) 说明用户点击的是 OCR 语言选择或重新选择按钮,None 说明用户的操作与 OCR 无关。
然后就是这个 handler 本身了:

#[handler]
async fn ocr_inlinekeyboard_handler(
    context: &Context,
    query: CallbackQuery,
) -> Result<HandlerResult, ExecuteError> {
    // 检查非空 query
    if let Some(data) = query.data {
        // 尝试 parse callback data
        if let Some(operation) = OCR_LANGS.try_parse_callback(data) {
            let message = query.message.unwrap();
            let chat_id = message.get_chat_id();
            let user_id = query.from.id;
            // 从 session 获取 OCR 状态
            let mut session = context
                .session_manager
                .get_session(SessionId::new(chat_id, user_id))
                .unwrap();
            let ocr: Option<Ocr> = session
                .get("ocr")
                .await
                .unwrap();
            // 检查该用户是否触发过 /ocr 指令
            if let Some(mut ocr) = ocr {
                // 用户触发过指令,保存用户的目标操作
                ocr.set(&operation);
                session
                    .set("ocr", &ocr)
                    .await
                    .unwrap();
                let method: EditMessageText;
                // 检查用户目标操作是否是选择语言
                if let Operation::Select(lang) = operation {
                    // 用户目标操作是选择语言
                    method = EditMessageText::new(
                        chat_id,
                        message.id,
                        format!("OCR 目标语言:{}\n\n请发送需要识别的图片(需以 Telegram 图片方式发送)", OCR_LANGS.get_lang_name(&lang)),
                    )
                    .reply_markup(OCR_LANGS.get_reselect_keyboard());
                } else {
                    // 用户目标操作是重新选择语言
                    method = EditMessageText::new(chat_id, message.id, "请选择 OCR 目标语言:")
                        .reply_markup(OCR_LANGS.get_langs_keyboard());
                }
                context.api.execute(method).await?;
                // 回应 callback
                let method = AnswerCallbackQuery::new(query.id);
                context.api.execute(method).await?;
            } else {
                // 用户没有触发过指令,以错误提示回应 callback
                let method = AnswerCallbackQuery::new(query.id)
                    .text("如需图片识别,请使用 /ocr 命令")
                    .show_alert(true);
                context.api.execute(method).await?;
            }
            return Ok(HandlerResult::Stop);
        }
    }
    Ok(HandlerResult::Continue)
}

由于这个 handler 会处理所有 CallbackQuery,所以应该一步一步定位这个 query 是否是由执行过 OCR 命令的用户发起的,范围由大到小,并且将有 IO 操作的步骤放在里层,来最小化对性能的影响。

最后的 HandlerResult 告诉 dispatcher 是否要继续顺序执行匹配后面的其它 handlers,HandlerResult::Continue 继续执行匹配,HandlerResult::Stop 反之。

处理图片消息

在处理图片前,需要先从 Telegram 把目标图片下载到服务器上,但是具体要下载到什么地方?在最开始写大体框架时我们就已经创建过一个目录 tmpdir 了。在调用 FilesystemBackend::new() 创建 session 的后端时,当作参数传入了 tmpdir.path(),这一步并没有消耗 tmpdir 本身,所以可以用这个临时目录存放下载的图片。之前提到过,dispatcher 会将 &Context 传给每一个 handler,所以可以把 tmpdir 也放进 Context

struct Context {
    // --snip--
    tmpdir: TempDir,
}

同时也修改创建 dispatcher 处:

// --snip--
let mut dispatcher = Dispatcher::new(Context {
    api: api.clone(),
    session_manager: SessionManager::new(backend),
    tmpdir,
});
// --snip--

终于到了最后一步。在用户触发 /ocr 命令并且选择了语言后,接收用户发送的图片,用 Tesseract 处理之后把结果发给用户。

#[handler]
async fn ocr_image_handler(
    context: &Context,
    message: Message,
) -> Result<HandlerResult, ExecuteError> {
    if let MessageData::Photo { data, .. } = &message.data {
        let chat_id = message.get_chat_id();
        // 获取 Photo 发送者
        if let Some(user) = message.get_user() {
            // 从 session 获取 OCR 状态
            let mut session = context
                .session_manager
                .get_session(SessionId::new(chat_id, user.id))
                .unwrap();
            let ocr: Option<Ocr> = session
                .get("ocr")
                .await
                .unwrap();
            // 检查该用户是否触发过 /ocr 指令
            if let Some(ocr) = ocr {
                // 检查该用户是否已经选择过 OCR 目标语言
                if let Some(lang) = ocr.get() {
                    // 从 session 中移除存储的 OCR 状态
                    session
                        .remove("ocr")
                        .await
                        .unwrap();
                    // 获取图片 URL
                    let file_id = &data.last().unwrap().file_id;
                    let method = GetFile::new(file_id);
                    let photo = context.api.execute(method).await?;
                    let photo_url = photo.file_path.unwrap();
                    // 下载图片
                    let photo_save_path = {
                        let mut path = context.tmpdir.path().to_path_buf().join(file_id);
                        path.set_extension("jpg");
                        path
                    };
                    let mut photo = File::create(&photo_save_path).await.unwrap();
                    let mut stream = context
                        .api
                        .download_file(photo_url)
                        .await
                        .unwrap();
                    while let Some(chunk) = stream.next().await {
                        photo
                            .write_all(&chunk.unwrap())
                            .await
                            .unwrap();
                    }
                    // 使用 LepTess 识别图片
                    let mut leptess =
                        LepTess::new(None, &lang).unwrap();
                    leptess
                        .set_image(photo_save_path)
                        .unwrap();
                    let result = leptess.get_utf8_text().unwrap_or(String::from("识别失败"));
                    // 发送结果
                    let method = SendMessage::new(chat_id, result);
                    context.api.execute(method).await?;
                    return Ok(HandlerResult::Stop);
                }
            }
        }
    }
    Ok(HandlerResult::Continue)
}

这个 handler 会处理每一条收到的消息,判断这条消息是不是收到了图片,如果是的话就检查是否是由执行过 OCR 命令并选择了目标语言的用户发送的,都没有问题就把这张图片下载到 tmpdir 后交给 Tesseract 处理。

这样这个 OCR bot 就基本完成了,下面来做最后的一点收尾工作:异常处理。

异常处理

上面每一个 handler 返回的都是 Result<HandlerResult, ExecuteError>,在文档中查看 DispatcherHandlerResult 时可以找到 这里
可以看到,作为 handler 的 Err 必须要有 Error + Send + Sync + 'static,所以用 anyhow 把所有错误无脑传上去是行不通的(anyhow 实现了 trait From<Error>,所以不能实现 trait Error,否则会出现递归 From),只能手写错误类型。为了方便,这里用 thiserror 生成错误类型:

#[derive(Error, Debug)]
enum OcrError {
    // 无法通过 Telegram Bot API 执行操作的错误
    #[error("{0}")]
    ExecuteError(#[from] carapax::ExecuteError),
    // 无法获取 session 的错误
    #[error("failed to get session")]
    GetSessionError,
    // 无法从 session 读写数据的错误
    #[error("failed to read / write data from session")]
    SessionDataError,
    // 无法操作 IO 的错误
    #[error("{0}")]
    IoError(#[from] std::io::Error),
    // 下载文件的错误
    #[error("failed to download file")]
    FileDownloadError,
    // Tesseract 初始化错误
    #[error("failed to initial Tesseract")]
    TessInitError,
    // 无法读取图片的错误
    #[error("failed to read image for Tesseract")]
    TessReadImageError,
}

然后把上面 handler 的返回值 Result<HandlerResult, ExecuteError> 改为 Result<HandlerResult, OcrError>,并将代码中的 unwrap() 替换成 ?.or_else(|_| return Err(OcrError::错误类型))?,就能向上传递错误到 dispatcher 了。

处理 dispatcher 收到的错误,在 dispatcher.add_handler() 后加上一个错误处理 handler:

dispatcher.set_error_handler(OcrErrorHandler);

不同于其它 handler,错误处理 handler 是一个实现了框架里的 trait ErrorHandler 的结构体。收到错误后,dispatcher 会调用错误处理 handler 实现的 trait ErrorHandlerhandle() 方法,这个结构体本身不需要任何功能:

struct ErrorHandler;

#[async_trait]
impl carapax::ErrorHandler for ErrorHandler {
    async fn handle(&mut self, err: carapax::HandlerError) -> ErrorPolicy {
        // 打印错误至 stderr
        eprintln!("{}", err);
        ErrorPolicy::Stop
    }
}

Rust 原生的 trait 并不支持 async 函数,所以要用 async_trait 这个 crate 实现 async。

写在最后

现在网上吹 Rust 的文章太多了,所以我也没有必要再列一大堆 Rust 的优点在这里。我是把 Rust 作为兴趣来学习的,写这个 OCR bot 的时候,我的 The Book 进度只有前几章。但随着之后继续深入学习,我更能明白为什么这门语言能连续六年都是 most loved language 了。

用 Rust 写东西心里会更有底,因为编译器会把底层的错误过滤掉。只要代码能编译,剩下的问题大多是逻辑错误,而不用像 C/C++ 一样到处找 seg fault 了。

学 Rust 不一定就是为了写 Rust。这门语言的设计是靠总结之前编程遇到的问题完善的,所以学习它的思想更重要,能把 Rust 解决问题的思路应用在其它地方才是最好的。

respond-post-48

添加新评论

请填写称呼
请填写合法的电子邮箱地址
请填写合法的网站地址
请填写内容