fixed predict

This commit is contained in:
Joshua Perry 2024-07-11 11:42:56 +01:00
parent 04de502f42
commit 03ebafe5d4
1 changed files with 5 additions and 6 deletions

View File

@ -1,4 +1,5 @@
use axum::{ use axum::{
extract::Path,
body::{Body, Bytes}, body::{Body, Bytes},
response::{IntoResponse, Json, Response} response::{IntoResponse, Json, Response}
}; };
@ -6,7 +7,7 @@ use crate::{
ApiResult, ApiResult,
account::auth account::auth
}; };
use http::{HeaderMap, StatusCode}; use http::StatusCode;
use pyo3::{ use pyo3::{
prelude::*, prelude::*,
types::PyTuple types::PyTuple
@ -19,7 +20,7 @@ const MODEL: &[u8] = include_bytes!("model/model.keras");
pub fn router() -> axum::Router { pub fn router() -> axum::Router {
axum::Router::new() axum::Router::new()
.route("/", .route("/:api",
axum::routing::post(get_predict) axum::routing::post(get_predict)
) )
} }
@ -30,10 +31,8 @@ struct Prediction {
malignant_probability: f64, malignant_probability: f64,
} }
async fn get_predict(headers: HeaderMap, body: Bytes) -> ApiResult { async fn get_predict(Path(api): Path<String>, body: Bytes) -> ApiResult {
let api = headers["api_key"].to_str()?; match auth(&api).await? {
match auth(api).await? {
Some(_user) => { Some(_user) => {
let prediction = predict(body).await?; let prediction = predict(body).await?;