photoprism/internal/commands/vision_reset.go

package commands

import (
	"fmt"
	"slices"
	"strings"

	"github.com/manifoldco/promptui"
	"github.com/urfave/cli/v2"

	"github.com/photoprism/photoprism/internal/ai/vision"
	"github.com/photoprism/photoprism/internal/config"
	"github.com/photoprism/photoprism/internal/workers"
	"github.com/photoprism/photoprism/pkg/txt"
)

// VisionResetCommand configures the command name, flags, and action.
var VisionResetCommand = &cli.Command{
	Name:      "reset",
	Usage:     "Resets data generated by computer vision models for pictures that match the specified search filters",
	ArgsUsage: "[filter]...",
	Flags: []cli.Flag{
		&cli.StringFlag{
			Name:    "models",
			Aliases: []string{"m"},
			Usage:   "computer vision `MODELS` to reset, e.g. caption or labels",
			Value:   "",
		},
		PicturesCountFlag(),
		VisionSourceFlag(vision.DefaultSrc),
		&cli.BoolFlag{
			Name:    "yes",
			Aliases: []string{"y"},
			Usage:   "runs the command non-interactively",
		},
	},
	Action: visionResetAction,
}

func visionResetAction(ctx *cli.Context) error {
	return CallWithDependencies(ctx, func(conf *config.Config) error {
		models := vision.ParseModelTypes(ctx.String("models"))
		resetCaptions := slices.Contains(models, vision.ModelTypeCaption)
		resetLabels := slices.Contains(models, vision.ModelTypeLabels)

		if !resetCaptions && !resetLabels {
			log.Warn("vision: no resettable models were specified, nothing to reset")
			return nil
		}

		selectedModels := make([]string, 0, 2)

		if resetCaptions {
			selectedModels = append(selectedModels, vision.ModelTypeCaption)
		}
		if resetLabels {
			selectedModels = append(selectedModels, vision.ModelTypeLabels)
		}

		confirmed := RunNonInteractively(ctx.Bool("yes"))

		if !confirmed && len(selectedModels) > 0 {
			label := fmt.Sprintf("Reset generated %s for matching pictures?", txt.JoinAnd(selectedModels))
			prompt := promptui.Prompt{Label: label, IsConfirm: true}
			if _, err := prompt.Run(); err != nil {
				return nil
			}
		}

		worker := workers.NewVision(conf)
		filter := strings.TrimSpace(strings.Join(ctx.Args().Slice(), " "))
		source, err := sanitizeVisionSource(ctx.String("source"))

		if err != nil {
			return cli.Exit(err.Error(), 1)
		}

		return worker.Reset(
			filter,
			ctx.Int("count"),
			selectedModels,
			string(source),
		)
	})
}