From 337b83450f25fc6dfe7ea382d4d3b4a0b8883af3 Mon Sep 17 00:00:00 2001 From: Kinch Date: Wed, 27 Sep 2023 14:21:41 +0200 Subject: [PATCH] rr: can display ms form results --- Cargo.lock | 1 + lua/rr.example.lua | 64 +++++++++++++++++++++++++++++--- lua_modules/rerun_lua/Cargo.toml | 3 +- lua_modules/rerun_lua/src/lib.rs | 47 ++++++++++++++++++++++- 4 files changed, 107 insertions(+), 8 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7aa0f81..1199bf9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4479,6 +4479,7 @@ dependencies = [ "image", "mlua", "rerun", + "serde", "serde_json", ] diff --git a/lua/rr.example.lua b/lua/rr.example.lua index 9bbcf79..f094eb8 100644 --- a/lua/rr.example.lua +++ b/lua/rr.example.lua @@ -1,10 +1,64 @@ local rr = require("rerun_lua") ---local record = rr.recording("example") ---record:image("54-0.png", "/Users/tbr/Desktop/00054-0.png") - -function bbox(boundings) +function bbox(ms_ocr_bbox) + return rr.rect_xyxy { + ms_ocr_bbox[1], + ms_ocr_bbox[2], + ms_ocr_bbox[5], + ms_ocr_bbox[6], + } end -local json = rr.json("/Users/tbr/Desktop/ms_form.json") +function paragraphs(ms_form_result) + local result = {} + for _, paragraph in ipairs(ms_form_result.analyzeResult.paragraphs) do + for _, region in ipairs(paragraph.boundingRegions) do + table.insert(result, { + text = paragraph.content, + page = region.pageNumber, + bbox = bbox(region.polygon) + }) + end + end + return result +end + +function lines(ms_form_result) + local result = {} + for _, page in ipairs(ms_form_result.analyzeResult.pages) do + for _, line in ipairs(page.lines) do + table.insert(result, { + text = line.content, + page = page.pageNumber, + bbox = bbox(line.polygon) + }) + end + end + return result +end + + + +function words(ms_form_result) + local result = {} + for _, page in ipairs(ms_form_result.analyzeResult.pages) do + for _, word in ipairs(page.words) do + table.insert(result, { + text = word.content, + page = page.pageNumber, + bbox = bbox(word.polygon) + }) + end + end + return result +end + + +local ms_form = rr.json("/Users/tbr/Desktop/ms_form.json") + +local record = rr.recording("example-2") +record:image("54-0", "/Users/tbr/Desktop/00054-0.png") +record:text_objects("54-0/ms-forms/paragraphs",paragraphs(ms_form)) +record:text_objects("54-0/ms-forms/lines",lines(ms_form)) +record:text_objects("54-0/ms-forms/words",words(ms_form)) diff --git a/lua_modules/rerun_lua/Cargo.toml b/lua_modules/rerun_lua/Cargo.toml index 5832763..a1168b6 100644 --- a/lua_modules/rerun_lua/Cargo.toml +++ b/lua_modules/rerun_lua/Cargo.toml @@ -11,4 +11,5 @@ crate-type = ["cdylib"] image = { version = "0.24.7", features = ["png", "jpeg", "tiff"] } mlua = { workspace = true } rerun = "0.8.2" -serde_json = "1.0.107" \ No newline at end of file +serde_json = "1.0.107" +serde = { version = "1.0.188", features = ["derive"] } \ No newline at end of file diff --git a/lua_modules/rerun_lua/src/lib.rs b/lua_modules/rerun_lua/src/lib.rs index e329ad3..3ac7bd1 100644 --- a/lua_modules/rerun_lua/src/lib.rs +++ b/lua_modules/rerun_lua/src/lib.rs @@ -1,12 +1,36 @@ -use std::fmt::format; use std::fs::File; use std::io::BufReader; use std::path::PathBuf; -use mlua::AnyUserData; +use mlua::{AnyUserData, Function, MetaMethod, Table, TableExt, UserData, UserDataRef}; use mlua::prelude::*; use rerun::{MsgSender, RecordingStream, RecordingStreamBuilder}; use rerun::components::{Rect2D, Tensor, Vec4D}; +use serde::{Deserialize, Serialize}; +#[derive(Debug)] +pub struct TextObject { + pub text: String, + pub bbox: Rect2D, +} + +impl UserData for TextObject { + fn add_methods<'lua, M: LuaUserDataMethods<'lua, Self>>(methods: &mut M) { + methods.add_meta_method(MetaMethod::ToString, |_lua, this, _:()| { + Ok(format!("{:?}", this)) + }) + } +} + +impl<'lua> FromLua<'lua> for TextObject { + fn from_lua(value: LuaValue<'lua>, lua: &'lua Lua) -> LuaResult { + let table = value.as_table().ok_or(LuaError::RuntimeError("invalid args for text_object".to_string()))?; + let rect = table.get::<_, UserDataRef>("bbox")?.clone(); + Ok(TextObject { + text: table.get("text")?, + bbox: rect + }) + } +} fn json(lua: &Lua, path: String) -> LuaResult { let file = File::open(path.as_str()) .map_err(|error| LuaError::RuntimeError(format!("could not open json path: {path}. Error: {error}")))?; @@ -29,6 +53,11 @@ fn rect_xyxy(lua: &Lua, points: [f32;4]) -> LuaResult { } fn register_bindings(lua: &Lua) -> LuaResult<()> { + + lua.register_userdata_type::(|reg| { + reg.add_meta_method(MetaMethod::ToString, |lua, this, _:()| Ok(format!("{:?}", this))) + })?; + lua.register_userdata_type::(|reg|{ reg.add_method("image", |lua, this, (entity_path, image_path): (String,String)| { let image_tensor = Tensor::from_image_file(PathBuf::from(image_path.as_str()).as_path()) @@ -41,6 +70,20 @@ fn register_bindings(lua: &Lua) -> LuaResult<()> { .map_err(|error| LuaError::RuntimeError(format!("Could not send message {error}")))?; Ok(()) }); + + reg.add_method("text_objects", |lua, this, (entity_path, objects): (String, Vec)| { + let rects: Vec = objects.into_iter() + .map(|t| t.bbox) + .collect(); + MsgSender::new(entity_path) + .with_timeless(true) + .with_component(rects.as_slice()) + .map_err(|error| LuaError::RuntimeError(format!("Could not send message: {}", error)))? + .send(this) + .map_err(|error| LuaError::RuntimeError(format!("Could not send message {error}")))?; + + Ok(()) + }) })?; Ok(()) }