captchaforge 0.2.39

//! YOLOv8 object detection for CAPTCHA challenge screenshots.
//!
//! Runs ONNX inference via `ort` with CUDA/TensorRT acceleration
//! and automatic CPU fallback. Targets reCAPTCHA / hCaptcha image
//! grids where the challenge asks the user to select tiles containing
//! a specific object class.

use anyhow::{Context, Result};
use image::{imageops, DynamicImage, GenericImageView};
use std::path::Path;
use tracing::{debug, info};

pub const YOLO_INPUT_SIZE: u32 = 640;
const CONF_THRESHOLD: f32 = 0.25;
const NMS_THRESHOLD: f32 = 0.45;

/// A single object detection result.
#[derive(Debug, Clone, PartialEq)]
pub struct Detection {
    /// Normalized bounding box [x1, y1, x2, y2] in 0..1 range.
    pub bbox: [f32; 4],
    /// COCO class name (e.g. "traffic light", "bus", "bicycle").
    pub class: String,
    /// Detection confidence 0.0–1.0.
    pub confidence: f32,
}

/// YOLOv8 detector backed by ONNX Runtime.
pub struct YoloDetector {
    session: ort::session::Session,
}

impl YoloDetector {
    /// Load a YOLOv8 ONNX model from disk.
    pub fn load(model_path: &Path) -> Result<Self> {
        let session = ort::session::Session::builder()
            .map_err(|e| anyhow::anyhow!("create ONNX session builder: {e}"))?
            .with_optimization_level(ort::session::builder::GraphOptimizationLevel::Level3)
            .map_err(|e| anyhow::anyhow!("set graph optimization: {e}"))?
            .with_execution_providers([
                ort::execution_providers::CUDAExecutionProvider::default().build(),
                ort::execution_providers::TensorRTExecutionProvider::default().build(),
                ort::execution_providers::CPUExecutionProvider::default().build(),
            ])
            .map_err(|e| anyhow::anyhow!("set execution providers: {e}"))?
            .commit_from_file(model_path)
            .map_err(|e| anyhow::anyhow!("load ONNX model from {}: {e}", model_path.display()))?;

        info!(model = %model_path.display(), "YOLOv8 detector loaded");
        Ok(Self { session })
    }

    /// Run detection on an image, returning all detections above the
    /// confidence threshold after NMS.
    pub fn detect(&mut self, image: &DynamicImage) -> Result<Vec<Detection>> {
        let (orig_w, orig_h) = image.dimensions();

        // Resize to YOLO input size with padding (letterbox).
        let (resized, pad_x, pad_y, scale) = letterbox(image, YOLO_INPUT_SIZE);
        let rgb = resized.to_rgb8();

        // NCHW float tensor, normalized to [0, 1].
        let (w, h) = (YOLO_INPUT_SIZE as usize, YOLO_INPUT_SIZE as usize);
        let mut pixel_data = vec![0.0f32; 3 * h * w];
        for y in 0..h {
            for x in 0..w {
                let pixel = rgb.get_pixel(x as u32, y as u32);
                pixel_data[y * w + x] = pixel[0] as f32 / 255.0;
                pixel_data[h * w + y * w + x] = pixel[1] as f32 / 255.0;
                pixel_data[2 * h * w + y * w + x] = pixel[2] as f32 / 255.0;
            }
        }

        let shape = vec![1_i64, 3, h as i64, w as i64];
        let input =
            ort::value::Tensor::from_array((shape, pixel_data)).context("build input tensor")?;

        let outputs = self
            .session
            .run(ort::inputs! { "images" => input })
            .context("run YOLO inference")?;

        let output_view = outputs[0]
            .try_extract_array::<f32>()
            .context("extract output tensor")?;

        let output_flat: Vec<f32> = output_view.iter().copied().collect();
        let shape = output_view.shape().to_vec();

        let raw =
            Self::parse_raw_detections(&output_flat, &shape, pad_x, pad_y, scale, orig_w, orig_h)?;
        debug!(raw = raw.len(), "YOLO raw detections");

        let filtered = nms(&raw);
        debug!(after_nms = filtered.len(), "YOLO after NMS");

        Ok(filtered)
    }

    /// Parse the raw YOLOv8 output tensor.
    ///
    /// YOLOv8 ONNX export produces [1, 84, 8400] (transposed) where:
    /// - 84 = 4 box coords (center_x, center_y, width, height) + 80 COCO class scores
    /// - 8400 = number of anchor boxes
    fn parse_raw_detections(
        flat: &[f32],
        shape: &[usize],
        pad_x: f32,
        pad_y: f32,
        scale: f32,
        orig_w: u32,
        orig_h: u32,
    ) -> Result<Vec<Detection>> {
        if shape.len() != 3 {
            anyhow::bail!("expected 3D output, got {:?}", shape);
        }

        let (_batch, dim_a, dim_b) = (shape[0], shape[1], shape[2]);
        let (num_detections, num_outputs) = if dim_a > dim_b {
            // [1, 8400, 84] format
            (dim_a, dim_b)
        } else {
            // [1, 84, 8400] transposed
            (dim_b, dim_a)
        };

        let is_transposed = dim_a < dim_b;
        let num_classes = num_outputs.saturating_sub(4);

        let mut detections = Vec::with_capacity(num_detections.min(100));

        for i in 0..num_detections {
            let get = |j: usize| -> f32 {
                if is_transposed {
                    flat[j * num_detections + i]
                } else {
                    flat[i * num_outputs + j]
                }
            };

            let cx = get(0);
            let cy = get(1);
            let bw = get(2);
            let bh = get(3);

            // Find best class.
            let mut best_class = 0;
            let mut best_score = 0.0f32;
            for c in 0..num_classes {
                let score = get(4 + c);
                if score > best_score {
                    best_score = score;
                    best_class = c;
                }
            }

            if best_score < CONF_THRESHOLD {
                continue;
            }

            // Convert from padded/letterboxed coordinates back to original image.
            let x1 = ((cx - bw / 2.0) - pad_x) / scale;
            let y1 = ((cy - bh / 2.0) - pad_y) / scale;
            let x2 = ((cx + bw / 2.0) - pad_x) / scale;
            let y2 = ((cy + bh / 2.0) - pad_y) / scale;

            // Normalize to 0..1.
            let nx1 = (x1 / orig_w as f32).clamp(0.0, 1.0);
            let ny1 = (y1 / orig_h as f32).clamp(0.0, 1.0);
            let nx2 = (x2 / orig_w as f32).clamp(0.0, 1.0);
            let ny2 = (y2 / orig_h as f32).clamp(0.0, 1.0);

            let class_name = coco_class_name(best_class);
            detections.push(Detection {
                bbox: [nx1, ny1, nx2, ny2],
                class: class_name.to_string(),
                confidence: best_score,
            });
        }

        Ok(detections)
    }
}

/// Letterbox resize: scale image to fit inside `target_size` while
/// maintaining aspect ratio, then pad with gray.
fn letterbox(image: &DynamicImage, target_size: u32) -> (DynamicImage, f32, f32, f32) {
    let (orig_w, orig_h) = image.dimensions();
    let scale = (target_size as f32 / orig_w as f32).min(target_size as f32 / orig_h as f32);

    let new_w = (orig_w as f32 * scale) as u32;
    let new_h = (orig_h as f32 * scale) as u32;

    let resized = image.resize_exact(new_w, new_h, imageops::FilterType::Triangle);

    let mut padded = DynamicImage::new_rgb8(target_size, target_size);
    let pad_x = (target_size - new_w) / 2;
    let pad_y = (target_size - new_h) / 2;

    imageops::overlay(&mut padded, &resized, pad_x as i64, pad_y as i64);

    (padded, pad_x as f32, pad_y as f32, scale)
}

/// Non-maximum suppression: keep only the best detection per spatial region.
fn nms(detections: &[Detection]) -> Vec<Detection> {
    let mut sorted: Vec<_> = detections.to_vec();
    sorted.sort_by(|a, b| {
        b.confidence
            .partial_cmp(&a.confidence)
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    let mut kept = Vec::with_capacity(sorted.len());
    let mut suppressed = vec![false; sorted.len()];

    for i in 0..sorted.len() {
        if suppressed[i] {
            continue;
        }
        kept.push(sorted[i].clone());
        for j in (i + 1)..sorted.len() {
            if suppressed[j] {
                continue;
            }
            if iou(&sorted[i].bbox, &sorted[j].bbox) > NMS_THRESHOLD {
                suppressed[j] = true;
            }
        }
    }

    kept
}

fn iou(a: &[f32; 4], b: &[f32; 4]) -> f32 {
    let x1 = a[0].max(b[0]);
    let y1 = a[1].max(b[1]);
    let x2 = a[2].min(b[2]);
    let y2 = a[3].min(b[3]);

    let inter_w = (x2 - x1).max(0.0);
    let inter_h = (y2 - y1).max(0.0);
    let inter = inter_w * inter_h;

    let area_a = (a[2] - a[0]) * (a[3] - a[1]);
    let area_b = (b[2] - b[0]) * (b[3] - b[1]);
    let union = area_a + area_b - inter;

    if union <= 0.0 {
        0.0
    } else {
        inter / union
    }
}

/// COCO class names (80 classes, 0-indexed).
fn coco_class_name(id: usize) -> &'static str {
    const NAMES: &[&str] = &[
        "person",
        "bicycle",
        "car",
        "motorcycle",
        "airplane",
        "bus",
        "train",
        "truck",
        "boat",
        "traffic light",
        "fire hydrant",
        "stop sign",
        "parking meter",
        "bench",
        "bird",
        "cat",
        "dog",
        "horse",
        "sheep",
        "cow",
        "elephant",
        "bear",
        "zebra",
        "giraffe",
        "backpack",
        "umbrella",
        "handbag",
        "tie",
        "suitcase",
        "frisbee",
        "skis",
        "snowboard",
        "sports ball",
        "kite",
        "baseball bat",
        "baseball glove",
        "skateboard",
        "surfboard",
        "tennis racket",
        "bottle",
        "wine glass",
        "cup",
        "fork",
        "knife",
        "spoon",
        "bowl",
        "banana",
        "apple",
        "sandwich",
        "orange",
        "broccoli",
        "carrot",
        "hot dog",
        "pizza",
        "donut",
        "cake",
        "chair",
        "couch",
        "potted plant",
        "bed",
        "dining table",
        "toilet",
        "tv",
        "laptop",
        "mouse",
        "remote",
        "keyboard",
        "cell phone",
        "microwave",
        "oven",
        "toaster",
        "sink",
        "refrigerator",
        "book",
        "clock",
        "vase",
        "scissors",
        "teddy bear",
        "hair drier",
        "toothbrush",
    ];
    NAMES.get(id).copied().unwrap_or("unknown")
}

/// Map reCAPTCHA / hCaptcha task text to COCO class names.
///
/// Returns `None` when the task mentions an object not in the COCO
/// dataset (the caller should fall back to VLM in that case).
pub fn task_to_coco_classes(task: &str) -> Option<Vec<&'static str>> {
    let lower = task.to_lowercase();

    // reCAPTCHA / hCaptcha task patterns → COCO classes.
    if lower.contains("traffic light") || lower.contains("traffic lights") {
        return Some(vec!["traffic light"]);
    }
    if lower.contains("bus") || lower.contains("buses") {
        return Some(vec!["bus"]);
    }
    if lower.contains("bicycle") || lower.contains("bicycles") || lower.contains("bike") {
        return Some(vec!["bicycle"]);
    }
    if lower.contains("car")
        || lower.contains("cars")
        || lower.contains("vehicle")
        || lower.contains("vehicles")
    {
        return Some(vec!["car", "truck", "bus", "motorcycle"]);
    }
    if lower.contains("motorcycle") || lower.contains("motorcycles") {
        return Some(vec!["motorcycle"]);
    }
    if lower.contains("train") || lower.contains("trains") {
        return Some(vec!["train"]);
    }
    if lower.contains("boat") || lower.contains("boats") {
        return Some(vec!["boat"]);
    }
    if lower.contains("fire hydrant") || lower.contains("fire hydrants") {
        return Some(vec!["fire hydrant"]);
    }
    if lower.contains("cat") || lower.contains("cats") {
        return Some(vec!["cat"]);
    }
    if lower.contains("dog") || lower.contains("dogs") {
        return Some(vec!["dog"]);
    }
    if lower.contains("person") || lower.contains("people") || lower.contains("pedestrian") {
        return Some(vec!["person"]);
    }
    if lower.contains("bird") || lower.contains("birds") {
        return Some(vec!["bird"]);
    }

    None
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn coco_names_complete() {
        assert_eq!(coco_class_name(0), "person");
        assert_eq!(coco_class_name(9), "traffic light");
        assert_eq!(coco_class_name(79), "toothbrush");
    }

    #[test]
    fn task_mapping_known() {
        assert_eq!(
            task_to_coco_classes("Select all images with traffic lights"),
            Some(vec!["traffic light"])
        );
        assert_eq!(task_to_coco_classes("Select all buses"), Some(vec!["bus"]));
    }

    #[test]
    fn task_mapping_unknown() {
        assert!(task_to_coco_classes("Select all crosswalks").is_none());
        assert!(task_to_coco_classes("Select all chimneys").is_none());
    }

    #[test]
    fn iou_calculation() {
        let a = [0.0, 0.0, 1.0, 1.0];
        let b = [0.5, 0.5, 1.5, 1.5];
        assert!((iou(&a, &b) - 0.1428).abs() < 0.01);
    }

    /// Integration test: load real YOLOv8n ONNX model and run inference.
    /// Skips gracefully if model is not cached and network is unavailable.
    #[test]
    #[cfg(feature = "vision")]
    fn yolov8n_loads_and_runs() {
        use crate::vision::ModelHub;
        let hub = ModelHub::new();
        let rt = tokio::runtime::Runtime::new().unwrap();
        let path = match rt.block_on(async { hub.resolve(super::super::ModelId::YoloV8n).await }) {
            Ok(p) => p,
            Err(e) => {
                eprintln!("Skipping integration test, model unavailable: {e}");
                return;
            }
        };

        let mut detector = YoloDetector::load(&path).expect("load YOLOv8n from disk");

        // Create a synthetic test image (red square on gray background).
        let mut img = image::RgbImage::new(640, 480);
        for pixel in img.pixels_mut() {
            *pixel = image::Rgb([128, 128, 128]);
        }
        for y in 100..300 {
            for x in 100..300 {
                img.put_pixel(x, y, image::Rgb([255, 0, 0]));
            }
        }
        let dyn_img = image::DynamicImage::ImageRgb8(img);

        let detections = detector
            .detect(&dyn_img)
            .expect("run detection without panic");
        // The model runs, we don't assert specific detections on synthetic
        // data because YOLOv8n is trained on real-world COCO images.
        println!(
            "YOLOv8n detections on synthetic image: {}",
            detections.len()
        );
    }
}