stash/internal/manager/task_generate_interactive_heatmap_speed.go
WithoutPants 4a054ab081
Support file-less scenes. Add scene split, merge and reassign file (#3006)
* Reassign scene file functionality
* Implement scene create
* Add scene create UI
* Add sceneMerge backend support
* Add merge scene to UI
* Populate split create with scene details
* Add merge button to duplicate checker
* Handle file-less scenes in marker preview generate
* Make unique file name for file-less scene exports
* Add o-counter to scene update input
* Hide rescan for file-less scenes
* Generate heatmap if no speed set on file
* Fix count in scene/image queries
2022-11-14 16:35:09 +11:00

71 lines
1.9 KiB
Go

package manager
import (
"context"
"fmt"
"github.com/stashapp/stash/pkg/file/video"
"github.com/stashapp/stash/pkg/fsutil"
"github.com/stashapp/stash/pkg/logger"
"github.com/stashapp/stash/pkg/models"
)
type GenerateInteractiveHeatmapSpeedTask struct {
Scene models.Scene
Overwrite bool
fileNamingAlgorithm models.HashAlgorithm
TxnManager Repository
}
func (t *GenerateInteractiveHeatmapSpeedTask) GetDescription() string {
return fmt.Sprintf("Generating heatmap and speed for %s", t.Scene.Path)
}
func (t *GenerateInteractiveHeatmapSpeedTask) Start(ctx context.Context) {
if !t.shouldGenerate() {
return
}
videoChecksum := t.Scene.GetHash(t.fileNamingAlgorithm)
funscriptPath := video.GetFunscriptPath(t.Scene.Path)
heatmapPath := instance.Paths.Scene.GetInteractiveHeatmapPath(videoChecksum)
generator := NewInteractiveHeatmapSpeedGenerator(funscriptPath, heatmapPath)
err := generator.Generate()
if err != nil {
logger.Errorf("error generating heatmap: %s", err.Error())
return
}
median := generator.InteractiveSpeed
if err := t.TxnManager.WithTxn(ctx, func(ctx context.Context) error {
primaryFile := t.Scene.Files.Primary()
primaryFile.InteractiveSpeed = &median
qb := t.TxnManager.File
return qb.Update(ctx, primaryFile)
}); err != nil {
logger.Error(err.Error())
}
}
func (t *GenerateInteractiveHeatmapSpeedTask) shouldGenerate() bool {
primaryFile := t.Scene.Files.Primary()
if primaryFile == nil || !primaryFile.Interactive {
return false
}
sceneHash := t.Scene.GetHash(t.fileNamingAlgorithm)
return !t.doesHeatmapExist(sceneHash) || primaryFile.InteractiveSpeed == nil || t.Overwrite
}
func (t *GenerateInteractiveHeatmapSpeedTask) doesHeatmapExist(sceneChecksum string) bool {
if sceneChecksum == "" {
return false
}
imageExists, _ := fsutil.FileExists(instance.Paths.Scene.GetInteractiveHeatmapPath(sceneChecksum))
return imageExists
}