use crate::
{
error::Result,
vision::{ Vision, ImageInput, ClassificationResult },
environment::HuggingFaceEnvironment,
components::input::BinaryClassificationInput,
};
use serde::Deserialize;
#[ derive( Debug, Deserialize ) ]
#[ serde( untagged ) ]
enum ClassificationResponse
{
Single( Vec< ClassificationResult > ),
Batch( Vec< Vec< ClassificationResult > > ),
}
impl< E > Vision< E >
where
E : HuggingFaceEnvironment + crate::environment::EnvironmentInterface + Send + Sync + 'static + Clone,
{
#[ inline ]
pub async fn classify_image(
&self,
image : ImageInput,
model : impl AsRef< str >
) -> Result< Vec< ClassificationResult > >
{
let request = BinaryClassificationInput
{
inputs : image.to_base64(),
};
let endpoint = format!( "/models/{}", model.as_ref() );
let url = self.client.environment.endpoint_url( &endpoint )?;
let response : ClassificationResponse = self.client
.post( url.as_str(), &request )
.await?;
match response
{
ClassificationResponse::Single( results ) => Ok( results ),
ClassificationResponse::Batch( batch ) =>
{
Ok( batch.into_iter().next().unwrap_or_default() )
}
}
}
}