chore: update output names

Gravitalia · Dec 28, 2023 · 84a35cc · 84a35cc
1 parent ac4c094
commit 84a35cc
Showing 1 changed file with 3 additions and 9 deletions.
diff --git a/src/corpus/mod.rs b/src/corpus/mod.rs
@@ -3,15 +3,9 @@ use tract_onnx::prelude::*;
 
 const IMAGE_WIDTH: u32 = 224;
 const IMAGE_HEIGHT: u32 = 224;
-const RESULT: [&str; 8] = [
-    "animals",
-    "flower",
-    "human",
-    "landscape",
+const RESULT: [&str; 2] = [
+    "not_nude",
     "nude",
-    "plant",
-    "sport",
-    "vehicle",
 ];
 
 /// Define a structure to manage the Corpus model.
@@ -70,7 +64,7 @@ pub fn init() -> Result<super::Model> {
     let model = tract_onnx::onnx()
         .model_for_path("./src/corpus/model.onnx")?
         .with_input_fact(0, f32::fact([1, 224, 224, 3]).into())?
-        .with_output_fact(0, InferenceFact::dt_shape(f32::datum_type(), tvec![1, 8]))?
+        .with_output_fact(0, InferenceFact::dt_shape(f32::datum_type(), tvec![1, RESULT.len()]))?
         .into_optimized()?
         .into_runnable()?;