Add test endpoint (#31)

xnorpx · Dec 26, 2024 · 612cc69 · 612cc69
1 parent cce6f6c
commit 612cc69
Show file tree

Hide file tree

Showing 6 changed files with 218 additions and 41 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -15,6 +15,7 @@ axum = { version = "0", default-features = false, features = [
     "multipart",
     "tokio",
 ] }
+base64 = "0"
 bytes = "1"
 clap = { version = "4", default-features = false, features = [
     "color",
@@ -33,6 +34,7 @@ image = "0"
 imageproc = "0"
 indicatif = "0"
 jpeg-encoder = "0"
+mime = "0"
 ndarray = "0"
 num_cpus = "1"
 ort = { version = "2.0.0-rc.9", default-features = false, features = [

diff --git a/src/image.rs b/src/image.rs
@@ -259,3 +259,22 @@ impl Resizer {
         Ok(())
     }
 }
+
+pub fn draw_boundary_boxes_on_encoded_image(
+    data: Bytes,
+    predictions: &[Prediction],
+) -> anyhow::Result<Bytes> {
+    let mut image = Image::default();
+    decode_jpeg(None, data, &mut image)?;
+    let dynamic_image_with_boundary_box =
+        create_dynamic_image_maybe_with_boundary_box(Some(predictions), &image, 20)?;
+    let mut encoded_image = Vec::new();
+    let encoder = Encoder::new(&mut encoded_image, 100);
+    encoder.encode(
+        dynamic_image_with_boundary_box.as_rgb8().unwrap(),
+        dynamic_image_with_boundary_box.width() as u16,
+        dynamic_image_with_boundary_box.height() as u16,
+        ColorType::Rgb,
+    )?;
+    Ok(Bytes::from(encoded_image))
+}
diff --git a/src/server.rs b/src/server.rs
@@ -1,6 +1,9 @@
-use crate::api::{
-    StatusUpdateResponse, VersionInfo, VisionCustomListResponse, VisionDetectionRequest,
-    VisionDetectionResponse,
+use crate::{
+    api::{
+        StatusUpdateResponse, VersionInfo, VisionCustomListResponse, VisionDetectionRequest,
+        VisionDetectionResponse,
+    },
+    image::draw_boundary_boxes_on_encoded_image,
 };
 use askama::Template;
 use axum::{
@@ -11,7 +14,10 @@ use axum::{
     routing::{get, post},
     Json, Router,
 };
+use base64::{engine::general_purpose, Engine as _};
+use bytes::Bytes;
 use chrono::Utc;
+use mime::IMAGE_JPEG;
 use reqwest;
 use serde::Deserialize;
 use std::{
@@ -65,6 +71,8 @@ pub async fn run_server(
         .route("/v1/vision/custom/list", post(v1_vision_custom_list))
         .route("/stats", get(stats_handler))
         .with_state(server_state.clone())
+        .route("/test", get(show_form).post(handle_upload))
+        .with_state(server_state.clone())
         .route("/favicon.ico", get(favicon_handler))
         .fallback(fallback_handler)
         .layer(DefaultBodyLimit::max(THIRTY_MEGABYTES));
@@ -190,33 +198,16 @@ async fn stats_handler(State(server_state): State<Arc<ServerState>>) -> impl Int
     )
 }
 
-async fn favicon_handler() -> impl IntoResponse {
-    StatusCode::NO_CONTENT
+async fn show_form() -> impl IntoResponse {
+    let template = TestTemplate { image_data: None };
+    (
+        [(CACHE_CONTROL, "no-store, no-cache, must-revalidate")],
+        template.into_response(),
+    )
 }
 
-struct BlueOnyxError(anyhow::Error);
-
-impl IntoResponse for BlueOnyxError {
-    fn into_response(self) -> Response {
-        (
-            StatusCode::INTERNAL_SERVER_ERROR,
-            Json(VisionDetectionResponse {
-                success: false,
-                message: "".into(),
-                error: Some(self.0.to_string()),
-                predictions: vec![],
-                count: 0,
-                command: "".into(),
-                module_id: "".into(),
-                execution_provider: "".into(),
-                can_useGPU: false,
-                inference_ms: 0_i32,
-                process_ms: 0_i32,
-                analysis_round_trip_ms: 0_i32,
-            }),
-        )
-            .into_response()
-    }
+async fn favicon_handler() -> impl IntoResponse {
+    StatusCode::NO_CONTENT
 }
 
 async fn fallback_handler(req: Request<Body>) -> impl IntoResponse {
@@ -235,16 +226,6 @@ async fn fallback_handler(req: Request<Body>) -> impl IntoResponse {
 
     (StatusCode::NOT_FOUND, "Endpoint not implemented")
 }
-
-impl<E> From<E> for BlueOnyxError
-where
-    E: Into<anyhow::Error>,
-{
-    fn from(err: E) -> Self {
-        Self(err.into())
-    }
-}
-
 #[allow(unused)]
 #[derive(Debug, Deserialize)]
 struct VersionJson {
@@ -355,3 +336,115 @@ impl Metrics {
         self.avg_ms(self.total_analysis_round_trip_ms)
     }
 }
+
+#[derive(Template)]
+#[template(path = "test.html")]
+struct TestTemplate<'a> {
+    image_data: Option<&'a str>,
+}
+
+async fn handle_upload(
+    State(server_state): State<Arc<ServerState>>,
+    mut multipart: Multipart,
+) -> impl IntoResponse {
+    let request_start_time = Instant::now();
+    while let Some(field) = multipart
+        .next_field()
+        .await
+        .map_err(|_| StatusCode::BAD_REQUEST)?
+    {
+        let name = field.name().unwrap_or("").to_string();
+        if name == "image" {
+            let content_type = field
+                .content_type()
+                .map(|ct| ct.to_string())
+                .unwrap_or_default();
+            if content_type != IMAGE_JPEG.to_string() {
+                return Err(StatusCode::BAD_REQUEST);
+            }
+
+            let data: Bytes = field.bytes().await.map_err(|_| StatusCode::BAD_REQUEST)?;
+            let vision_request = VisionDetectionRequest {
+                min_confidence: 0., // This will be set to None and will use server default
+                image_data: data.clone(),
+                image_name: "image.jpg".to_string(),
+            };
+
+            let (sender, receiver) = tokio::sync::oneshot::channel();
+            if let Err(err) = server_state.sender.send((vision_request, sender)) {
+                error!(?err, "Failed to send request to detection worker");
+            }
+            let result = timeout(Duration::from_secs(30), receiver).await;
+
+            let mut vision_response = match result {
+                Ok(Ok(response)) => response,
+                Ok(Err(err)) => {
+                    error!("Failed to receive vision detection response: {:?}", err);
+                    return Err(StatusCode::INTERNAL_SERVER_ERROR);
+                }
+                Err(_) => {
+                    error!("Timeout while waiting for vision detection response");
+                    return Err(StatusCode::INTERNAL_SERVER_ERROR);
+                }
+            };
+
+            let data =
+                draw_boundary_boxes_on_encoded_image(data, vision_response.predictions.as_slice())
+                    .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+
+            let encoded = general_purpose::STANDARD.encode(&data);
+            let data_url = format!("data:image/jpeg;base64,{}", encoded);
+
+            let template = TestTemplate {
+                image_data: Some(&data_url),
+            };
+
+            vision_response.analysis_round_trip_ms =
+                request_start_time.elapsed().as_millis() as i32;
+
+            {
+                let mut metrics = server_state.metrics.lock().await;
+                metrics.update_metrics(&vision_response);
+            }
+            return Ok((
+                [(CACHE_CONTROL, "no-store, no-cache, must-revalidate")],
+                template.into_response(),
+            ));
+        }
+    }
+    Err(StatusCode::BAD_REQUEST)
+}
+
+struct BlueOnyxError(anyhow::Error);
+
+impl IntoResponse for BlueOnyxError {
+    fn into_response(self) -> Response {
+        (
+            StatusCode::INTERNAL_SERVER_ERROR,
+            Json(VisionDetectionResponse {
+                success: false,
+                message: "".into(),
+                error: Some(self.0.to_string()),
+                predictions: vec![],
+                count: 0,
+                command: "".into(),
+                module_id: "".into(),
+                execution_provider: "".into(),
+                can_useGPU: false,
+                inference_ms: 0_i32,
+                process_ms: 0_i32,
+                analysis_round_trip_ms: 0_i32,
+            }),
+        )
+            .into_response()
+    }
+}
+
+impl<E> From<E> for BlueOnyxError
+where
+    E: Into<anyhow::Error>,
+{
+    fn from(err: E) -> Self {
+        Self(err.into())
+    }
+}
diff --git a/src/worker.rs b/src/worker.rs
@@ -48,9 +48,9 @@ impl DetectorWorker {
                 Some(image_name)
             };
 
-            let detect_result = self
-                .detector
-                .detect(image_data, image_name, Some(min_confidence));
+            let min_confidence = (min_confidence > 0.01).then_some(min_confidence);
+
+            let detect_result = self.detector.detect(image_data, image_name, min_confidence);
 
             let detect_response = match detect_result {
                 Ok(detect_result) => VisionDetectionResponse {

diff --git a/templates/test.html b/templates/test.html
@@ -0,0 +1,61 @@
+<!-- templates/test.html -->
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <title>Blue Onyx - Test</title>
+    <style>
+        body {
+            font-family: Arial, sans-serif;
+            background-color: #f0f8ff;
+            margin: 50px;
+            text-align: center;
+        }
+        h1 {
+            color: #4682b4;
+        }
+        form {
+            margin: 20px auto;
+            padding: 20px;
+            border: 2px solid #4682b4;
+            border-radius: 10px;
+            width: 300px;
+            background-color: #ffffff;
+        }
+        input[type="file"] {
+            width: 90%;
+            padding: 10px;
+            margin: 10px 0;
+        }
+        input[type="submit"] {
+            padding: 10px 20px;
+            background-color: #4682b4;
+            color: #ffffff;
+            border: none;
+            border-radius: 5px;
+            cursor: pointer;
+        }
+        img {
+            margin-top: 20px;
+            max-width: 80%;
+            height: auto;
+            border: 2px solid #4682b4;
+            border-radius: 10px;
+        }
+    </style>
+</head>
+<body>
+    <h1>Blue Onyx</h1>
+
+    <form action="/test" method="post" enctype="multipart/form-data">
+        <input type="file" name="image" accept="image/jpeg" required>
+        <br>
+        <input type="submit" value="Upload">
+    </form>
+
+    {% if image_data.is_some() %}
+        <h2>Processed Image:</h2>
+        <img src="{{ image_data.unwrap() }}" alt="Processed Image">
+    {% endif %}
+</body>
+</html>