use std::path::PathBuf;
use anyhow::{Context, Result, bail};
use inferencelayer::tableformer::TableFormer;
use inferencelayer::vision::decode_rgb8;
use serde_json::json;
fn main() -> Result<()> {
let usage = "usage: tableformer <model-dir> (--page <png> --bbox l,t,r,b | --crop <png>)";
let mut args = std::env::args().skip(1);
let dir = PathBuf::from(args.next().context(usage)?);
let mut page: Option<PathBuf> = None;
let mut crop: Option<PathBuf> = None;
let mut bbox: Option<[f32; 4]> = None;
while let Some(a) = args.next() {
match a.as_str() {
"--page" => page = Some(PathBuf::from(args.next().context("--page needs a path")?)),
"--crop" => crop = Some(PathBuf::from(args.next().context("--crop needs a path")?)),
"--bbox" => {
let s = args.next().context("--bbox needs l,t,r,b")?;
let v: Vec<f32> = s
.split(',')
.map(|x| x.trim().parse())
.collect::<Result<_, _>>()?;
if v.len() != 4 {
bail!(
"--bbox needs exactly 4 comma-separated numbers, got {}",
v.len()
);
}
bbox = Some([v[0], v[1], v[2], v[3]]);
}
other => bail!("unknown arg {other}\n{usage}"),
}
}
let model = TableFormer::load(&dir).context("load tableformer")?;
let (rgb, w, h, table_bbox) = match (page, crop) {
(Some(p), _) => {
let (rgb, w, h) = decode_rgb8(&std::fs::read(&p)?)?;
let bb = bbox.context("--page requires --bbox l,t,r,b")?;
(rgb, w, h, bb)
}
(None, Some(c)) => {
let (rgb, w, h) = decode_rgb8(&std::fs::read(&c)?)?;
(rgb, w, h, [0.0, 0.0, w as f32, h as f32])
}
(None, None) => bail!("{usage}"),
};
let out = model.recognize(&rgb, w, h, table_bbox);
let cells: Vec<_> = out
.cell_classes
.iter()
.zip(&out.bboxes)
.map(|(&class, b)| json!({ "class": class, "cxcywh": b }))
.collect();
let result = json!({
"otsl": out.otsl_seq,
"num_cells": out.bboxes.len(),
"cells": cells,
});
println!("{}", serde_json::to_string_pretty(&result)?);
Ok(())
}