server/rust-batcher/examples/optimized_client.rs

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(())
}