AI: Allow to use defaults when having a custom vision.yml #127 photoprism#5011

Signed-off-by: Michael Mayer <michael@photoprism.app>

Author: lastzero

internal/ai/vision/model.go

@@ -25,6 +25,7 @@ var (
 // Model represents a computer vision model configuration.
 type Model struct {
 	Type          ModelType             `yaml:"Type,omitempty" json:"type,omitempty"`
+	Default       bool                  `yaml:"Default,omitempty" json:"default,omitempty"`
 	Name          string                `yaml:"Name,omitempty" json:"name,omitempty"`
 	Version       string                `yaml:"Version,omitempty" json:"version,omitempty"`
 	System        string                `yaml:"System,omitempty" json:"system,omitempty"`

internal/ai/vision/models.go

@@ -9,6 +9,7 @@ import (
 var (
 	NasnetModel = &Model{
 		Type:       ModelTypeLabels,
+		Default:    true,
 		Name:       "nasnet",
 		Version:    VersionMobile,
 		Resolution: 224,
@@ -39,6 +40,7 @@ var (
 	}
 	NsfwModel = &Model{
 		Type:       ModelTypeNsfw,
+		Default:    true,
 		Name:       "nsfw",
 		Version:    VersionLatest,
 		Resolution: 224,
@@ -61,6 +63,7 @@ var (
 	}
 	FacenetModel = &Model{
 		Type:       ModelTypeFace,
+		Default:    true,
 		Name:       "facenet",
 		Version:    VersionLatest,
 		Resolution: 160,

internal/ai/vision/options.go

@@ -45,6 +45,33 @@ func (c *Options) Load(fileName string) error {
 		return err
 	}
 
+	// 1. Ensure that there is at least one configuration for each model type,
+	//    so that adding a copy of the default configuration to the vision.yml file
+	//    is not required. We could alternatively require a model to included in
+	//    the "vision.yml" file, but set the defaults if the "Default" flag is set.
+	// 2. Use the default "Thresholds" if no custom thresholds are configured.
+
+	for i, model := range c.Models {
+		if !model.Default {
+			continue
+		}
+
+		switch model.Type {
+		case ModelTypeLabels:
+			c.Models[i] = NasnetModel
+		case ModelTypeNsfw:
+			c.Models[i] = NsfwModel
+		case ModelTypeFace:
+			c.Models[i] = FacenetModel
+		case ModelTypeCaption:
+			c.Models[i] = CaptionModel
+		}
+	}
+
+	if c.Thresholds.Confidence <= 0 || c.Thresholds.Confidence > 100 {
+		c.Thresholds.Confidence = DefaultThresholds.Confidence
+	}
+
 	return nil
 }
 

internal/ai/vision/testdata/vision.yml

@@ -1,5 +1,6 @@
 Models:
 - Type: labels
+  Default: true
   Name: nasnet
   Version: mobile
   Resolution: 224
@@ -19,6 +20,7 @@ Models:
       Name: predictions/Softmax
       Outputs: 1000
 - Type: nsfw
+  Default: true
   Name: nsfw
   Version: latest
   Resolution: 224
@@ -33,6 +35,7 @@ Models:
       Name: nsfw_cls_model/final_prediction
       Outputs: 5
 - Type: face
+  Default: true
   Name: facenet
   Version: latest
   Resolution: 160