AI: Allow users to use a custom source for TensorFlow labels photoprism#5011 photoprism#5232

Signed-off-by: Michael Mayer <[email protected]>

Author: lastzero

internal/ai/vision/config.go

@@ -92,13 +92,13 @@ func (c *ConfigValues) Load(fileName string) error {
 
 		switch model.Type {
 		case ModelTypeLabels:
-			c.Models[i] = NasnetModel
+			c.Models[i] = NasnetModel.Clone()
 		case ModelTypeNsfw:
-			c.Models[i] = NsfwModel
+			c.Models[i] = NsfwModel.Clone()
 		case ModelTypeFace:
-			c.Models[i] = FacenetModel
+			c.Models[i] = FacenetModel.Clone()
 		case ModelTypeCaption:
-			c.Models[i] = CaptionModel
+			c.Models[i] = CaptionModel.Clone()
 		}
 
 		if runType != RunAuto {

internal/ai/vision/labels.go

@@ -16,7 +16,7 @@ import (
 var labelsFunc = labelsInternal
 
 // SetLabelsFunc overrides the labels generator. Intended for tests.
-func SetLabelsFunc(fn func(Files, media.Src, string) (classify.Labels, error)) {
+func SetLabelsFunc(fn func(Files, media.Src, entity.Src) (classify.Labels, error)) {
 	if fn == nil {
 		labelsFunc = labelsInternal
 		return
@@ -28,11 +28,11 @@ func SetLabelsFunc(fn func(Files, media.Src, string) (classify.Labels, error)) {
 // GenerateLabels finds matching labels for the specified image.
 // Caller must pass the appropriate metadata source string (e.g., entity.SrcOllama, entity.SrcOpenAI)
 // so that downstream indexing can record where the labels originated.
-func GenerateLabels(images Files, mediaSrc media.Src, labelSrc string) (classify.Labels, error) {
+func GenerateLabels(images Files, mediaSrc media.Src, labelSrc entity.Src) (classify.Labels, error) {
 	return labelsFunc(images, mediaSrc, labelSrc)
 }
 
-func labelsInternal(images Files, mediaSrc media.Src, labelSrc string) (result classify.Labels, err error) {
+func labelsInternal(images Files, mediaSrc media.Src, labelSrc entity.Src) (result classify.Labels, err error) {
 	// Return if no thumbnail filenames were given.
 	if len(images) == 0 {
 		return result, errors.New("at least one image required")
@@ -127,7 +127,7 @@ func labelsInternal(images Files, mediaSrc media.Src, labelSrc string) (result c
 					return result, err
 				}
 
-				result = mergeLabels(result, labels)
+				result = mergeLabels(result, labels, labelSrc)
 			}
 		} else {
 			return result, errors.New("invalid labels model configuration")
@@ -141,15 +141,19 @@ func labelsInternal(images Files, mediaSrc media.Src, labelSrc string) (result c
 	return result, nil
 }
 
-// mergeLabels combines existing labels with newly detected labels and returns the result.
-func mergeLabels(result, labels classify.Labels) classify.Labels {
+// mergeLabels combines existing labels with newly detected labels, applies a custom source, and returns the result.
+func mergeLabels(result, labels classify.Labels, labelSrc entity.Src) classify.Labels {
 	if len(labels) == 0 {
 		return result
 	}
 
 	for j := range labels {
 		found := false
 
+		if labelSrc != entity.SrcAuto {
+			labels[j].Source = labelSrc
+		}
+
 		for k := range result {
 			if labels[j].Name == result[k].Name {
 				found = true

internal/ai/vision/labels_test.go

@@ -49,6 +49,17 @@ func TestGenerateLabels(t *testing.T) {
 		assert.InDelta(t, 60, result[0].Uncertainty, 10)
 		assert.InDelta(t, float32(0.4), result[0].Confidence(), 0.1)
 	})
+	t.Run("CustomSourceLocal", func(t *testing.T) {
+		labels, err := GenerateLabels(Files{examplesPath + "/cat_224.jpeg"}, media.SrcLocal, entity.SrcManual)
+		if err != nil {
+			t.Fatalf("GenerateLabels error: %v", err)
+		}
+		for _, label := range labels {
+			if label.Source != entity.SrcManual {
+				t.Fatalf("expected custom source %q, got %q", entity.SrcManual, label.Source)
+			}
+		}
+	})
 	t.Run("InvalidFile", func(t *testing.T) {
 		_, err := GenerateLabels(Files{examplesPath + "/notexisting.jpg"}, media.SrcLocal, entity.SrcAuto)
 		assert.Error(t, err)

internal/photoprism/index_vision_test.go

@@ -83,7 +83,7 @@ func TestIndexLabelsSource(t *testing.T) {
 
 	t.Run("AutoUsesModelSource", func(t *testing.T) {
 		var captured string
-		vision.SetLabelsFunc(func(files vision.Files, mediaSrc media.Src, src string) (classify.Labels, error) {
+		vision.SetLabelsFunc(func(files vision.Files, mediaSrc media.Src, src entity.Src) (classify.Labels, error) {
 			captured = src
 			return classify.Labels{{Name: "stub", Source: src, Uncertainty: 0}}, nil
 		})
@@ -96,7 +96,7 @@ func TestIndexLabelsSource(t *testing.T) {
 
 	t.Run("CustomSource", func(t *testing.T) {
 		var captured string
-		vision.SetLabelsFunc(func(files vision.Files, mediaSrc media.Src, src string) (classify.Labels, error) {
+		vision.SetLabelsFunc(func(files vision.Files, mediaSrc media.Src, src entity.Src) (classify.Labels, error) {
 			captured = src
 			return classify.Labels{{Name: "stub", Source: src, Uncertainty: 0}}, nil
 		})

internal/photoprism/mediafile_vision_test.go

@@ -86,7 +86,7 @@ func TestMediaFile_GenerateLabels(t *testing.T) {
 
 	t.Run("AutoUsesModelSource", func(t *testing.T) {
 		var captured string
-		vision.SetLabelsFunc(func(files vision.Files, mediaSrc media.Src, src string) (classify.Labels, error) {
+		vision.SetLabelsFunc(func(files vision.Files, mediaSrc media.Src, src entity.Src) (classify.Labels, error) {
 			captured = src
 			return classify.Labels{{Name: "stub", Source: src}}, nil
 		})
@@ -98,7 +98,7 @@ func TestMediaFile_GenerateLabels(t *testing.T) {
 
 	t.Run("CustomSourceOverrides", func(t *testing.T) {
 		var captured string
-		vision.SetLabelsFunc(func(files vision.Files, mediaSrc media.Src, src string) (classify.Labels, error) {
+		vision.SetLabelsFunc(func(files vision.Files, mediaSrc media.Src, src entity.Src) (classify.Labels, error) {
 			captured = src
 			return classify.Labels{{Name: "stub", Source: src}}, nil
 		})