supper/src/ocr.rs

use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
};

use anyhow::Result;
use image::RgbImage;
use serde::{Deserialize, Serialize};

use crate::{
    config::{Config, LearnedConfig},
    image_processing::{extract_region, filter_to_white, hash_image, Region},
};

#[derive(Serialize, Deserialize, Debug)]
pub struct OcrRegion {
    pub confidence: f64,
    pub value: String,
}

#[derive(Serialize, Deserialize, Debug)]
struct OcrResult {
    regions: Vec<OcrRegion>,
    error: Option<String>,
}

async fn run_ocr(image: &RgbImage, url: &str) -> Result<Option<String>> {
    let client = reqwest::Client::new();
    let response = client
        .post(url)
        .body(crate::image_processing::to_png_bytes(image))
        .send()
        .await?;

    if !response.status().is_success() {
        eprintln!("failed to run OCR query");
        anyhow::bail!("failed to run OCR query")
    }
    let result: OcrResult = response.json().await?;
    let result = if result.regions.is_empty() {
        None
    } else {
        let mut buffer = String::new();
        for r in &result.regions {
            buffer += &r.value;
        }
        Some(buffer)
    };
    Ok(result)
}

#[tokio::main(flavor = "current_thread")]
pub async fn ocr_all_regions(
    image: &RgbImage,
    config: Arc<Config>,
    learned: Arc<LearnedConfig>,
) -> HashMap<String, Option<String>> {
    let results = Arc::new(Mutex::new(HashMap::new()));

    let mut handles = Vec::new();
    for region in &config.ocr_regions {
        let filtered_image = extract_region(image, region);
        let region = region.clone();
        let results = results.clone();
        let config = config.clone();
        let learned = learned.clone();
        handles.push(tokio::spawn(async move {
            let mut image = filtered_image;
            filter_to_white(&mut image);
            let hash = hash_image(&image);
            let value = if let Some(learned_value) = learned.learned_images.get(&hash) {
                Some(learned_value.clone())
            } else {
                run_ocr(&image, &config.ocr_server_endpoint)
                    .await
                    .unwrap_or(None)
            };
            results.lock().unwrap().insert(region.name, value);
        }));
    }
    for handle in handles {
        handle.await.expect("failed to join task in OCR");
    }

    let results = results.lock().unwrap().clone();
    results
}