internal/workers/vision_reset.go - develop - githubmirror/photoprism

internal/workers/vision_reset.go
152 строки4 KB
Michael Mayer
CLI: Ensure vision reset command and worker use same default src #5233
01 окт 2025, 09:52
01 окт 2025, 09:5254d96f6
package workers

import (
	"fmt"
	"path"
	"runtime/debug"
	"slices"
	"time"

	"github.com/dustin/go-humanize/english"

	"github.com/photoprism/photoprism/internal/ai/vision"
	"github.com/photoprism/photoprism/internal/entity"
	"github.com/photoprism/photoprism/internal/entity/query"
	"github.com/photoprism/photoprism/internal/entity/search"
	"github.com/photoprism/photoprism/internal/entity/sortby"
	"github.com/photoprism/photoprism/internal/form"
	"github.com/photoprism/photoprism/internal/mutex"
	"github.com/photoprism/photoprism/internal/photoprism"
	"github.com/photoprism/photoprism/pkg/clean"
	"github.com/photoprism/photoprism/pkg/txt"
)

// Reset removes data generated by the specified model types for photos matching the search filter.
func (w *Vision) Reset(filter string, count int, models []string, provider string) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("vision: %s (worker panic)\nstack: %s", r, debug.Stack())
			log.Error(err)
		}
	}()

	if err = mutex.VisionWorker.Start(); err != nil {
		return err
	}
	defer mutex.VisionWorker.Stop()

	resetLabels := slices.Contains(models, vision.ModelTypeLabels)
	resetCaptions := slices.Contains(models, vision.ModelTypeCaption)

	if n := len(models); n == 0 {
		log.Warnf("vision: no models were specified")
		return nil
	} else {
		log.Infof("vision: resetting %s model data", txt.JoinAnd(models))
	}

	provider = clean.ShortTypeLower(provider)
	if provider == "" {
		provider = vision.DefaultSrc
	}

	start := time.Now()
	done := make(map[string]bool)

	if count < 1 || count > search.MaxResults {
		count = search.MaxResults
	}

	frm := form.SearchPhotos{
		Query:   filter,
		Primary: true,
		Merged:  false,
		Count:   count,
		Offset:  0,
		Order:   sortby.Added,
	}

	photos, _, queryErr := search.Photos(frm)
	if queryErr != nil {
		return queryErr
	}

	if len(photos) == 0 {
		log.Info("vision: no pictures to reset")
		return nil
	}

	var (
		updated            int
		removedLabels      int64
		resetCaptionsCount int
	)

	for _, photo := range photos {
		if done[photo.PhotoUID] {
			continue
		}
		done[photo.PhotoUID] = true

		m, loadErr := query.PhotoByUID(photo.PhotoUID)
		if loadErr != nil {
			log.Errorf("vision: failed to load %s (%s)", photo.PhotoUID, loadErr)
			continue
		}

		changed := false

		if resetCaptions {
			if changedCaption, capErr := m.ResetCaption(provider); capErr != nil {
				log.Warnf("vision: %s (reset caption)", clean.Error(capErr))
			} else if changedCaption {
				resetCaptionsCount++
				changed = true
			}
		}

		if resetLabels {
			if removed, lblErr := m.ResetLabels(provider); lblErr != nil {
				log.Warnf("vision: %s (reset labels)", clean.Error(lblErr))
			} else if removed > 0 {
				removedLabels += removed
				changed = true
				log.Debugf("vision: removed %d labels from %s", removed, clean.Log(path.Join(m.PhotoPath, m.PhotoName)))
			}
		}

		if !changed {
			continue
		}

		updated++

		if err := m.GenerateAndSaveTitle(); err != nil {
			log.Warnf("vision: %s (generate title)", clean.Error(err))
		}
	}

	if removedLabels > 0 {
		entity.FlushPhotoLabelCache()
	}

	if updated > 0 {
		if moments := photoprism.NewMoments(w.conf); moments == nil {
			log.Errorf("vision: failed to update moments")
		} else if err = moments.Start(); err != nil {
			log.Warnf("moments: %s in optimization worker", err)
		}

		if err = entity.UpdateCounts(); err != nil {
			log.Warnf("vision: %s in optimization worker", err)
		}

		if err = query.UpdateCovers(); err != nil {
			log.Warnf("vision: %s in optimization worker", err)
		}
	}

	log.Infof("vision: reset %s [%s]", english.Plural(updated, "picture", "pictures"), time.Since(start))

	return nil
}