diff --git a/manager/generator_sprite.go b/manager/generator_sprite.go new file mode 100644 index 000000000..cac63b0a3 --- /dev/null +++ b/manager/generator_sprite.go @@ -0,0 +1,135 @@ +package manager + +import ( + "fmt" + "github.com/bmatcuk/doublestar" + "github.com/disintegration/imaging" + "github.com/stashapp/stash/ffmpeg" + "github.com/stashapp/stash/logger" + "github.com/stashapp/stash/utils" + "image" + "image/color" + "io/ioutil" + "math" + "path/filepath" + "strings" +) + +type SpriteGenerator struct { + Info *GeneratorInfo + + ImageOutputPath string + VTTOutputPath string + Rows int + Columns int +} + +func NewSpriteGenerator(videoFile ffmpeg.VideoFile, imageOutputPath string, vttOutputPath string, rows int, cols int) (*SpriteGenerator, error) { + exists, err := utils.FileExists(videoFile.Path) + if !exists { + return nil, err + } + generator, err := newGeneratorInfo(videoFile) + if err != nil { + return nil, err + } + generator.ChunkCount = rows * cols + if err := generator.configure(); err != nil { + return nil, err + } + + return &SpriteGenerator{ + Info: generator, + ImageOutputPath: imageOutputPath, + VTTOutputPath: vttOutputPath, + Rows: rows, + Columns: cols, + }, nil +} + +func (g *SpriteGenerator) Generate() error { + instance.Paths.Generated.EmptyTmpDir() + encoder := ffmpeg.NewEncoder(instance.Paths.FixedPaths.FFMPEG) + + if err := g.generateSpriteImage(&encoder); err != nil { + return err + } + if err := g.generateSpriteVTT(&encoder); err != nil { + return err + } + return nil +} + +func (g *SpriteGenerator) generateSpriteImage(encoder *ffmpeg.Encoder) error { + logger.Infof("[generator] generating sprite image for %s", g.Info.VideoFile.Path) + + // Create `this.chunkCount` thumbnails in the tmp directory + stepSize := int(g.Info.VideoFile.Duration / float64(g.Info.ChunkCount)) + for i := 0; i < g.Info.ChunkCount; i++ { + time := i * stepSize + num := fmt.Sprintf("%.3d", i) + filename := "thumbnail" + num + ".jpg" + + options := ffmpeg.ScreenshotOptions{ + OutputPath: instance.Paths.Generated.GetTmpPath(filename), + Time: float64(time), + Width: 160, + } + encoder.Screenshot(g.Info.VideoFile, options) + } + + // Combine all of the thumbnails into a sprite image + globPath := filepath.Join(instance.Paths.Generated.Tmp, "thumbnail*.jpg") + imagePaths, _ := doublestar.Glob(globPath) + var images []image.Image + for _, imagePath := range imagePaths { + img, err := imaging.Open(imagePath) + if err != nil { + return err + } + images = append(images, img) + } + + width := images[0].Bounds().Size().X + height := images[0].Bounds().Size().Y + canvasWidth := width * g.Columns + canvasHeight := height * g.Rows + montage := imaging.New(canvasWidth, canvasHeight, color.NRGBA{}) + for index := 0; index < len(images); index++ { + x := width * (index % g.Columns) + y := height * int(math.Floor(float64(index) / float64(g.Rows))) + img := images[index] + montage = imaging.Paste(montage, img, image.Pt(x, y)) + } + + return imaging.Save(montage, g.ImageOutputPath) +} + +func (g *SpriteGenerator) generateSpriteVTT(encoder *ffmpeg.Encoder) error { + logger.Infof("[generator] generating sprite vtt for %s", g.Info.VideoFile.Path) + + spriteImage, err := imaging.Open(g.ImageOutputPath) + if err != nil { + return err + } + spriteImageName := filepath.Base(g.ImageOutputPath) + width := spriteImage.Bounds().Size().X / g.Columns + height := spriteImage.Bounds().Size().Y / g.Rows + + stepSize := float64(g.Info.NthFrame) / g.Info.FrameRate + + vttLines := []string{"WEBVTT", ""} + for index := 0; index < g.Info.ChunkCount; index++ { + x := width * (index % g.Columns) + y := height * int(math.Floor(float64(index) / float64(g.Rows))) + startTime := utils.GetVTTTime(float64(index) * stepSize) + endTime := utils.GetVTTTime(float64(index + 1) * stepSize) + + vttLines = append(vttLines, startTime + " --> " + endTime) + vttLines = append(vttLines, fmt.Sprintf("%s#xywh=%d,%d,%d,%d", spriteImageName, x, y, width, height)) + vttLines = append(vttLines, "") + } + vtt := strings.Join(vttLines, "\n") + + return ioutil.WriteFile(g.VTTOutputPath, []byte(vtt), 0755) +} diff --git a/manager/manager.go b/manager/manager.go index 218c7a17f..18d4af5bb 100644 --- a/manager/manager.go +++ b/manager/manager.go @@ -109,9 +109,8 @@ func (s *singleton) Generate(sprites bool, previews bool, markers bool, transcod wg.Add(delta) if sprites { - go func() { - wg.Done() // TODO - }() + task := GenerateSpriteTask{Scene: scene} + go task.Start(&wg) } if previews { diff --git a/manager/task_generate_sprite.go b/manager/task_generate_sprite.go new file mode 100644 index 000000000..c0a8e2557 --- /dev/null +++ b/manager/task_generate_sprite.go @@ -0,0 +1,50 @@ +package manager + +import ( + "github.com/stashapp/stash/ffmpeg" + "github.com/stashapp/stash/logger" + "github.com/stashapp/stash/models" + "github.com/stashapp/stash/utils" + "sync" +) + +type GenerateSpriteTask struct { + Scene models.Scene +} + +func (t *GenerateSpriteTask) Start(wg *sync.WaitGroup) { + if t.doesSpriteExist(t.Scene.Checksum) { + wg.Done() + return + } + + videoFile, err := ffmpeg.NewVideoFile(instance.Paths.FixedPaths.FFProbe, t.Scene.Path) + if err != nil { + logger.Errorf("error reading video file: %s", err.Error()) + wg.Done() + return + } + + imagePath := instance.Paths.Scene.GetSpriteImageFilePath(t.Scene.Checksum) + vttPath := instance.Paths.Scene.GetSpriteVttFilePath(t.Scene.Checksum) + generator, err := NewSpriteGenerator(*videoFile, imagePath, vttPath, 9, 9) + if err != nil { + logger.Errorf("error creating sprite generator: %s", err.Error()) + wg.Done() + return + } + + if err := generator.Generate(); err != nil { + logger.Errorf("error generating sprite: %s", err.Error()) + wg.Done() + return + } + + wg.Done() +} + +func (t *GenerateSpriteTask) doesSpriteExist(sceneChecksum string) bool { + imageExists, _ := utils.FileExists(instance.Paths.Scene.GetSpriteImageFilePath(sceneChecksum)) + vttExists, _ := utils.FileExists(instance.Paths.Scene.GetSpriteVttFilePath(sceneChecksum)) + return imageExists && vttExists +} \ No newline at end of file diff --git a/utils/file.go b/utils/file.go index 5d6963e4b..0ffa5f52f 100644 --- a/utils/file.go +++ b/utils/file.go @@ -42,7 +42,6 @@ func RemoveDir(path string) error { return os.RemoveAll(path) } -// TODO test func EmptyDir(path string) error { d, err := os.Open(path) if err != nil {