AI: Add "source" flag to the "photoprism vision run" command #3438 #5011

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

Author: lastzero

internal/commands/vision_run.go

@@ -7,24 +7,31 @@ import (
 
 	"github.com/photoprism/photoprism/internal/ai/vision"
 	"github.com/photoprism/photoprism/internal/config"
+	"github.com/photoprism/photoprism/internal/entity"
 	"github.com/photoprism/photoprism/internal/workers"
 )
 
 // VisionRunCommand configures the command name, flags, and action.
 var VisionRunCommand = &cli.Command{
 	Name:      "run",
-	Usage:     "Runs a computer vision model",
+	Usage:     "Runs one or more computer vision models on a set of pictures, as specified by the search filter",
 	ArgsUsage: "[filter]",
 	Flags: []cli.Flag{
 		&cli.StringFlag{
 			Name:    "models",
 			Aliases: []string{"m"},
-			Usage:   "model types (labels, nsfw, caption)",
+			Usage:   "computer vision `MODELS` to run, e.g. caption, labels, or nsfw",
+		},
+		&cli.StringFlag{
+			Name:    "source",
+			Aliases: []string{"s"},
+			Value:   entity.SrcImage,
+			Usage:   "custom data source `TYPE`, e.g. estimate, image, meta, or manual",
 		},
 		&cli.BoolFlag{
 			Name:    "force",
 			Aliases: []string{"f"},
-			Usage:   "force existing metadata to be updated",
+			Usage:   "force existing data to be updated if the source priority is equal to or higher than the current one",
 		},
 	},
 	Action: visionRunAction,
@@ -34,6 +41,11 @@ var VisionRunCommand = &cli.Command{
 func visionRunAction(ctx *cli.Context) error {
 	return CallWithDependencies(ctx, func(conf *config.Config) error {
 		worker := workers.NewVision(conf)
-		return worker.Start(strings.TrimSpace(ctx.Args().First()), vision.ParseTypes(ctx.String("models")), ctx.Bool("force"))
+		return worker.Start(
+			strings.TrimSpace(ctx.Args().First()),
+			vision.ParseTypes(ctx.String("models")),
+			ctx.String("source"),
+			ctx.Bool("force"),
+		)
 	})
 }

internal/workers/vision.go

@@ -40,7 +40,7 @@ func (w *Vision) originalsPath() string {
 }
 
 // Start runs the specified model types for the photos that match the search query.
-func (w *Vision) Start(q string, models []string, force bool) (err error) {
+func (w *Vision) Start(q string, models []string, customSrc string, force bool) (err error) {
 	defer func() {
 		if r := recover(); r != nil {
 			err = fmt.Errorf("vision: %s (worker panic)\nstack: %s", r, debug.Stack())
@@ -68,6 +68,15 @@ func (w *Vision) Start(q string, models []string, force bool) (err error) {
 		log.Infof("vision: running %s models", strings.Join(models, " and "))
 	}
 
+	// Source type for AI generated data.
+	var dataSrc string
+
+	if customSrc = clean.ShortTypeLower(customSrc); customSrc != "" {
+		dataSrc = customSrc
+	} else {
+		dataSrc = entity.SrcImage
+	}
+
 	// Check time when worker was last executed.
 	updateIndex := false
 
@@ -128,13 +137,15 @@ func (w *Vision) Start(q string, models []string, force bool) (err error) {
 
 			changed := false
 
+			// Generate labels.
 			if updateLabels && (len(m.Labels) == 0 || force) {
 				if labels := ind.Labels(file); len(labels) > 0 {
 					m.AddLabels(labels)
 					changed = true
 				}
 			}
 
+			// Detect NSFW content.
 			if updateNsfw && (!photo.PhotoPrivate || force) {
 				if isNsfw := ind.IsNsfw(file); photo.PhotoPrivate != isNsfw {
 					photo.PhotoPrivate = isNsfw
@@ -143,19 +154,15 @@ func (w *Vision) Start(q string, models []string, force bool) (err error) {
 				}
 			}
 
-			if updateCaptions && (m.PhotoCaption == "" || force) {
+			// Generate a caption if none exists or the force flag is used,
+			// and only if no caption was set or removed by a higher-priority source.
+			if updateCaptions && entity.SrcPriority[dataSrc] >= entity.SrcPriority[m.CaptionSrc] && (m.NoCaption() || force) {
 				if caption, captionErr := ind.Caption(file); captionErr != nil {
 					log.Warnf("vision: %s in %s (generate caption)", clean.Error(captionErr), photoName)
-				} else if caption.Text != "" {
-					if caption.Source == "" {
-						caption.Source = entity.SrcImage
-					}
-
-					if (entity.SrcPriority[caption.Source] >= entity.SrcPriority[m.CaptionSrc]) || !m.HasCaption() {
-						m.SetCaption(caption.Text, caption.Source)
-						changed = true
-						log.Infof("vision: changed caption of %s to %s", photoName, clean.Log(m.PhotoCaption))
-					}
+				} else if caption.Text = strings.TrimSpace(caption.Text); caption.Text != "" {
+					m.SetCaption(caption.Text, dataSrc)
+					changed = true
+					log.Infof("vision: changed caption of %s to %s", photoName, clean.Log(m.PhotoCaption))
 				}
 			}