Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 5 commits into from
May 5, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 27 additions & 3 deletions pkg/ffmpeg/encoder.go
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
package ffmpeg

import (
"fmt"
"bytes"
"io/ioutil"
"os"
"os/exec"
Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand Down Expand Up @@ -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
}
38 changes: 38 additions & 0 deletions pkg/ffmpeg/encoder_sprite_screenshot.go
Original file line number Diff line number Diff line change
@@ -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
}
8 changes: 4 additions & 4 deletions pkg/ffmpeg/encoder_transcode.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
}
34 changes: 6 additions & 28 deletions pkg/manager/generator_phash.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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)
}
Expand All @@ -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
}
27 changes: 6 additions & 21 deletions pkg/manager/generator_sprite.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
}
Expand All @@ -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
Expand Down
1 change: 1 addition & 0 deletions ui/v2.5/src/components/Changelog/versions/v070.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down