Add in-memory screenshot generation for sprites and phash (#1316)

pkg/ffmpeg/encoder.go

@@ -1,7 +1,7 @@
 package ffmpeg
 
 import (
-	"fmt"
+	"bytes"
 	"io/ioutil"
 	"os"
 	"os/exec"
@@ -62,7 +62,7 @@ func KillRunningEncoders(path string) {
 
 	for _, process := range processes {
 		// assume it worked, don't check for error
-		fmt.Printf("Killing encoder process for file: %s", path)
+		logger.Infof("Killing encoder process for file: %s", path)
 		process.Kill()
 
 		// wait for the process to die before returning
@@ -82,7 +82,8 @@ func KillRunningEncoders(path string) {
 	}
 }
 
-func (e *Encoder) run(probeResult VideoFile, args []string) (string, error) {
+// FFmpeg runner with progress output, used for transcodes
+func (e *Encoder) runTranscode(probeResult VideoFile, args []string) (string, error) {
 	cmd := exec.Command(e.Path, args...)
 
 	stderr, err := cmd.StderrPipe()
@@ -137,3 +138,26 @@ func (e *Encoder) run(probeResult VideoFile, args []string) (string, error) {
 
 	return stdoutString, nil
 }
+
+func (e *Encoder) run(probeResult VideoFile, args []string) (string, error) {
+	cmd := exec.Command(e.Path, args...)
+
+	var stdout, stderr bytes.Buffer
+	cmd.Stdout = &stdout
+	cmd.Stderr = &stderr
+
+	if err := cmd.Start(); err != nil {
+		return "", err
+	}
+
+	registerRunningEncoder(probeResult.Path, cmd.Process)
+	err := waitAndDeregister(probeResult.Path, cmd)
+
+	if err != nil {
+		// error message should be in the stderr stream
+		logger.Errorf("ffmpeg error when running command <%s>: %s", strings.Join(cmd.Args, " "), stderr.String())
+		return stdout.String(), err
+	}
+
+	return stdout.String(), nil
+}

pkg/ffmpeg/encoder_sprite_screenshot.go

@@ -0,0 +1,38 @@
+package ffmpeg
+
+import (
+	"fmt"
+	"image"
+	"strings"
+)
+
+type SpriteScreenshotOptions struct {
+	Time  float64
+	Width int
+}
+
+func (e *Encoder) SpriteScreenshot(probeResult VideoFile, options SpriteScreenshotOptions) (image.Image, error) {
+	args := []string{
+		"-v", "error",
+		"-ss", fmt.Sprintf("%v", options.Time),
+		"-i", probeResult.Path,
+		"-vframes", "1",
+		"-vf", fmt.Sprintf("scale=%v:-1", options.Width),
+		"-c:v", "bmp",
+		"-f", "rawvideo",
+		"-",
+	}
+	data, err := e.run(probeResult, args)
+	if err != nil {
+		return nil, err
+	}
+
+	reader := strings.NewReader(data)
+
+	img, _, err := image.Decode(reader)
+	if err != nil {
+		return nil, err
+	}
+
+	return img, err
+}

pkg/ffmpeg/encoder_transcode.go

@@ -64,7 +64,7 @@ func (e *Encoder) Transcode(probeResult VideoFile, options TranscodeOptions) {
 		"-strict", "-2",
 		options.OutputPath,
 	}
-	_, _ = e.run(probeResult, args)
+	_, _ = e.runTranscode(probeResult, args)
 }
 
 //transcode the video, remove the audio
@@ -84,7 +84,7 @@ func (e *Encoder) TranscodeVideo(probeResult VideoFile, options TranscodeOptions
 		"-vf", "scale=" + scale,
 		options.OutputPath,
 	}
-	_, _ = e.run(probeResult, args)
+	_, _ = e.runTranscode(probeResult, args)
 }
 
 //copy the video stream as is, transcode audio
@@ -96,7 +96,7 @@ func (e *Encoder) TranscodeAudio(probeResult VideoFile, options TranscodeOptions
 		"-strict", "-2",
 		options.OutputPath,
 	}
-	_, _ = e.run(probeResult, args)
+	_, _ = e.runTranscode(probeResult, args)
 }
 
 //copy the video stream as is, drop audio
@@ -107,5 +107,5 @@ func (e *Encoder) CopyVideo(probeResult VideoFile, options TranscodeOptions) {
 		"-c:v", "copy",
 		options.OutputPath,
 	}
-	_, _ = e.run(probeResult, args)
+	_, _ = e.runTranscode(probeResult, args)
 }

pkg/manager/generator_phash.go

@@ -5,12 +5,9 @@ import (
 	"image"
 	"image/color"
 	"math"
-	"os"
-	"sort"
 
 	"github.com/corona10/goimagehash"
 	"github.com/disintegration/imaging"
-	"github.com/fvbommel/sortorder"
 
 	"github.com/stashapp/stash/pkg/ffmpeg"
 	"github.com/stashapp/stash/pkg/logger"
@@ -67,37 +64,22 @@ func (g *PhashGenerator) generateSprite(encoder *ffmpeg.Encoder) (image.Image, e
 	chunkCount := g.Columns * g.Rows
 	offset := 0.05 * g.Info.VideoFile.Duration
 	stepSize := (0.9 * g.Info.VideoFile.Duration) / float64(chunkCount)
+	var images []image.Image
 	for i := 0; i < chunkCount; i++ {
 		time := offset + (float64(i) * stepSize)
-		num := fmt.Sprintf("%.3d", i)
-		filename := "phash_" + g.VideoChecksum + "_" + num + ".bmp"
 
-		options := ffmpeg.ScreenshotOptions{
-			OutputPath: instance.Paths.Generated.GetTmpPath(filename),
-			Time:       time,
-			Width:      160,
-		}
-		if err := encoder.Screenshot(g.Info.VideoFile, options); err != nil {
-			return nil, err
+		options := ffmpeg.SpriteScreenshotOptions{
+			Time:  time,
+			Width: 160,
 		}
-	}
-
-	// Combine all of the thumbnails into a sprite image
-	pattern := fmt.Sprintf("phash_%s_.+\\.bmp$", g.VideoChecksum)
-	imagePaths, err := utils.MatchEntries(instance.Paths.Generated.Tmp, pattern)
-	if err != nil {
-		return nil, err
-	}
-	sort.Sort(sortorder.Natural(imagePaths))
-	var images []image.Image
-	for _, imagePath := range imagePaths {
-		img, err := imaging.Open(imagePath)
+		img, err := encoder.SpriteScreenshot(g.Info.VideoFile, options)
 		if err != nil {
 			return nil, err
 		}
 		images = append(images, img)
 	}
 
+	// Combine all of the thumbnails into a sprite image
 	if len(images) == 0 {
 		return nil, fmt.Errorf("images slice is empty, failed to generate phash sprite for %s", g.Info.VideoFile.Path)
 	}
@@ -113,9 +95,5 @@ func (g *PhashGenerator) generateSprite(encoder *ffmpeg.Encoder) (image.Image, e
 		montage = imaging.Paste(montage, img, image.Pt(x, y))
 	}
 
-	for _, imagePath := range imagePaths {
-		os.Remove(imagePath)
-	}
-
 	return montage, nil
 }

pkg/manager/generator_sprite.go

@@ -8,11 +8,9 @@ import (
 	"math"
 	"os"
 	"path/filepath"
-	"sort"
 	"strings"
 
 	"github.com/disintegration/imaging"
-	"github.com/fvbommel/sortorder"
 
 	"github.com/stashapp/stash/pkg/ffmpeg"
 	"github.com/stashapp/stash/pkg/logger"
@@ -75,29 +73,15 @@ func (g *SpriteGenerator) generateSpriteImage(encoder *ffmpeg.Encoder) error {
 
 	// Create `this.chunkCount` thumbnails in the tmp directory
 	stepSize := g.Info.VideoFile.Duration / float64(g.Info.ChunkCount)
+	var images []image.Image
 	for i := 0; i < g.Info.ChunkCount; i++ {
 		time := float64(i) * stepSize
-		num := fmt.Sprintf("%.3d", i)
-		filename := "thumbnail_" + g.VideoChecksum + "_" + num + ".jpg"
 
-		options := ffmpeg.ScreenshotOptions{
-			OutputPath: instance.Paths.Generated.GetTmpPath(filename),
-			Time:       time,
-			Width:      160,
+		options := ffmpeg.SpriteScreenshotOptions{
+			Time:  time,
+			Width: 160,
 		}
-		encoder.Screenshot(g.Info.VideoFile, options)
-	}
-
-	// Combine all of the thumbnails into a sprite image
-	pattern := fmt.Sprintf("thumbnail_%s_.+\\.jpg$", g.VideoChecksum)
-	imagePaths, err := utils.MatchEntries(instance.Paths.Generated.Tmp, pattern)
-	if err != nil {
-		return err
-	}
-	sort.Sort(sortorder.Natural(imagePaths))
-	var images []image.Image
-	for _, imagePath := range imagePaths {
-		img, err := imaging.Open(imagePath)
+		img, err := encoder.SpriteScreenshot(g.Info.VideoFile, options)
 		if err != nil {
 			return err
 		}
@@ -107,6 +91,7 @@ func (g *SpriteGenerator) generateSpriteImage(encoder *ffmpeg.Encoder) error {
 	if len(images) == 0 {
 		return fmt.Errorf("images slice is empty, failed to generate sprite images for %s", g.Info.VideoFile.Path)
 	}
+	// Combine all of the thumbnails into a sprite image
 	width := images[0].Bounds().Size().X
 	height := images[0].Bounds().Size().Y
 	canvasWidth := width * g.Columns

ui/v2.5/src/components/Changelog/versions/v070.md

@@ -15,6 +15,7 @@
 * Added scene queue.
 
 ### 🎨 Improvements
+* Improve sprite generation performance when using network storage.
 * Remove duplicate values when scraping lists of elements.
 * Improved performance of the auto-tagger.
 * Clean generation artifacts after generating each scene.