File size: 3,727 Bytes
3c0d3e1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | use anyhow::{anyhow, Context, Result};
use image::imageops::FilterType;
use std::{
env,
io::{Read, Write},
net::TcpStream,
path::PathBuf,
time::Instant,
};
const TARGET: u32 = 800;
fn usage() -> ! {
eprintln!("usage: optimized_client <image-path> [--server http://localhost:18082]");
std::process::exit(2);
}
fn parse_args() -> (PathBuf, String) {
let mut args = env::args().skip(1);
let Some(path) = args.next() else { usage() };
let mut server = "http://localhost:18082".to_string();
while let Some(arg) = args.next() {
match arg.as_str() {
"--server" => server = args.next().unwrap_or_else(|| usage()),
_ => usage(),
}
}
(
PathBuf::from(path),
server.trim_end_matches('/').to_string(),
)
}
fn parse_http_server(server: &str) -> Result<(String, String)> {
let Some(rest) = server.strip_prefix("http://") else {
return Err(anyhow!(
"example only supports plain http://server:port URLs"
));
};
let authority = rest.split('/').next().unwrap_or(rest);
if authority.is_empty() {
return Err(anyhow!("empty server authority"));
}
Ok((authority.to_string(), authority.to_string()))
}
fn post_octet_stream(server: &str, path_and_query: &str, body: Vec<u8>) -> Result<String> {
let (addr, host) = parse_http_server(server)?;
let mut stream = TcpStream::connect(&addr).with_context(|| format!("connect {addr}"))?;
let request_head = format!(
"POST {path_and_query} HTTP/1.1\r\nHost: {host}\r\nContent-Type: application/octet-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
body.len()
);
stream.write_all(request_head.as_bytes())?;
stream.write_all(&body)?;
stream.flush()?;
let mut response = Vec::new();
stream.read_to_end(&mut response)?;
let response = String::from_utf8(response).context("response was not UTF-8")?;
let Some((head, body)) = response.split_once("\r\n\r\n") else {
return Err(anyhow!("malformed HTTP response"));
};
let status = head.lines().next().unwrap_or_default();
if !status.contains(" 200 ") {
return Err(anyhow!("request failed: {status}\n{body}"));
}
Ok(body.to_string())
}
fn main() -> Result<()> {
let (path, server) = parse_args();
let image = image::open(&path)
.with_context(|| format!("open {}", path.display()))?
.to_rgb8();
let (original_width, original_height) = image.dimensions();
// This must match the server/model contract. The service still receives original dimensions
// so PP-DocLayout boxes are returned in original page coordinates.
let resized = image::imageops::resize(&image, TARGET, TARGET, FilterType::Triangle);
let plane = (TARGET * TARGET) as usize;
let mut chw_u8 = vec![0u8; plane * 3];
for y in 0..TARGET as usize {
for x in 0..TARGET as usize {
let px = resized.get_pixel(x as u32, y as u32).0;
let idx = y * TARGET as usize + x;
chw_u8[idx] = px[0];
chw_u8[plane + idx] = px[1];
chw_u8[2 * plane + idx] = px[2];
}
}
let path_and_query = format!(
"/v1/layout_chw_u8?width={TARGET}&height={TARGET}&original_width={original_width}&original_height={original_height}"
);
let start = Instant::now();
let body = post_octet_stream(&server, &path_and_query, chw_u8)?;
let elapsed = start.elapsed();
let value: serde_json::Value = serde_json::from_str(&body).context("decode response JSON")?;
eprintln!("request_ms={:.2}", elapsed.as_secs_f64() * 1000.0);
println!("{}", serde_json::to_string_pretty(&value)?);
Ok(())
}
|