AI: Add TensorFlow model shape detection #127 #5164

* AI: Added support for non BHWC models

Tensorflow models use BHWC by default, however, if we are using
converted models, we can find that the expected input is BCHW. Now the
input is configurable (although the restriction of being dimesion 4 is
still there) via Shape parameter on the input definition. Also, the
model instrospection will try to deduce the input shape from the model
signature.

* AI: Added more tests for enum parsing

ShapeComponent was missing from the tests

* AI: Modified external tests to the new url

The path has been moved from tensorflow/vision to tensorflow/models

* AI: Moved the builder to the model to reuse it

It should reduce the amount of allocations done

* AI: fixed errors after merge

Mainly incorrect paths and duplicated variables
This commit is contained in:
raystlin
2025-08-16 15:55:59 +02:00
committed by GitHub
parent 2a7351ee9a
commit 519a6ab34a
15 changed files with 502 additions and 157 deletions

View File

@@ -0,0 +1,96 @@
package tensorflow
import (
"path/filepath"
"slices"
"testing"
"github.com/photoprism/photoprism/pkg/fs"
)
var assetsPath = fs.Abs("../../../assets")
var testDataPath = fs.Abs("testdata")
func TestTF1ModelLoad(t *testing.T) {
model, err := SavedModel(
filepath.Join(assetsPath, "models", "nasnet"),
[]string{"photoprism"})
if err != nil {
t.Fatal(err)
}
input, output, err := GetInputAndOutputFromSavedModel(model)
if err == nil {
t.Fatalf("TF1 does not have signatures, but GetInput worked")
}
input, output, err = GuessInputAndOutput(model)
if err != nil {
t.Fatal(err)
}
if input == nil {
t.Fatal("Could not get the input")
} else if output == nil {
t.Fatal("Could not get the output")
} else if input.Shape == nil {
t.Fatal("Could not get the shape")
} else {
t.Logf("Shape: %v", input.Shape)
}
}
func TestTF2ModelLoad(t *testing.T) {
model, err := SavedModel(
filepath.Join(testDataPath, "tf2"),
[]string{"serve"})
if err != nil {
t.Fatal(err)
}
input, output, err := GetInputAndOutputFromSavedModel(model)
if err != nil {
t.Fatal(err)
}
if input == nil {
t.Fatal("Could not get the input")
} else if output == nil {
t.Fatal("Could not get the output")
} else if input.Shape == nil {
t.Fatal("Could not get the shape")
} else if !slices.Equal(input.Shape, DefaultPhotoInputShape()) {
t.Fatalf("Invalid shape calculated. Expected BHWC, got %v",
input.Shape)
}
}
func TestTF2ModelBCHWLoad(t *testing.T) {
model, err := SavedModel(
filepath.Join(testDataPath, "tf2_bchw"),
[]string{"serve"})
if err != nil {
t.Fatal(err)
}
input, output, err := GetInputAndOutputFromSavedModel(model)
if err != nil {
t.Fatal(err)
}
if input == nil {
t.Fatal("Could not get the input")
} else if output == nil {
t.Fatal("Could not get the output")
} else if input.Shape == nil {
t.Fatal("Could not get the shape")
} else if !slices.Equal(input.Shape, []ShapeComponent{
ShapeBatch, ShapeColor, ShapeHeight, ShapeWidth,
}) {
t.Fatalf("Invalid shape calculated. Expected BCHW, got %v",
input.Shape)
}
}