diff --git a/internal/ai/vision/config.go b/internal/ai/vision/config.go index 031125d34..625dfb426 100644 --- a/internal/ai/vision/config.go +++ b/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 { diff --git a/internal/ai/vision/labels.go b/internal/ai/vision/labels.go index e27767e07..f80f90da8 100644 --- a/internal/ai/vision/labels.go +++ b/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,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 diff --git a/internal/ai/vision/labels_test.go b/internal/ai/vision/labels_test.go index 7ff10b727..a944c92ca 100644 --- a/internal/ai/vision/labels_test.go +++ b/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) diff --git a/internal/photoprism/index_vision_test.go b/internal/photoprism/index_vision_test.go index 0e8bcb746..4c92dd309 100644 --- a/internal/photoprism/index_vision_test.go +++ b/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 }) diff --git a/internal/photoprism/mediafile_vision_test.go b/internal/photoprism/mediafile_vision_test.go index ffe23be9a..23abe4990 100644 --- a/internal/photoprism/mediafile_vision_test.go +++ b/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 })