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

Signed-off-by: Michael Mayer <michael@photoprism.app>
This commit is contained in:
Michael Mayer 2025-10-30 12:40:03 +01:00
parent 46b219770a
commit 0a1b5071fd
5 changed files with 29 additions and 14 deletions

View file

@ -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 {

View file

@ -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,8 +141,8 @@ 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
}
@ -150,6 +150,10 @@ func mergeLabels(result, labels classify.Labels) classify.Labels {
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

View file

@ -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)

View file

@ -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
})

View file

@ -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
})