mirror of
https://github.com/stashapp/stash.git
synced 2025-12-17 20:34:37 +03:00
Restructure ffmpeg (#2392)
* Refactor transcode generation * Move phash generation into separate package * Refactor image thumbnail generation * Move JSONTime to separate package * Ffmpeg refactoring * Refactor live transcoding * Refactor scene marker preview generation * Refactor preview generation * Refactor screenshot generation * Refactor sprite generation * Change ffmpeg.IsStreamable to return error * Move frame rate calculation into ffmpeg * Refactor file locking * Refactor title set during scan * Add missing lockmanager instance * Return error instead of logging in MatchContainer
This commit is contained in:
WithoutPants
330
pkg/scene/generate/sprite.go
Normal file
330
pkg/scene/generate/sprite.go
Normal file
@@ -0,0 +1,330 @@
|
||||
package generate
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/disintegration/imaging"
|
||||
"github.com/stashapp/stash/pkg/ffmpeg"
|
||||
"github.com/stashapp/stash/pkg/ffmpeg/transcoder"
|
||||
"github.com/stashapp/stash/pkg/fsutil"
|
||||
"github.com/stashapp/stash/pkg/utils"
|
||||
)
|
||||
|
||||
const (
|
||||
spriteScreenshotWidth = 160
|
||||
|
||||
spriteRows = 9
|
||||
spriteCols = 9
|
||||
spriteChunks = spriteRows * spriteCols
|
||||
)
|
||||
|
||||
func (g Generator) SpriteScreenshot(ctx context.Context, input string, seconds float64) (image.Image, error) {
|
||||
lockCtx := g.LockManager.ReadLock(ctx, input)
|
||||
defer lockCtx.Cancel()
|
||||
|
||||
ssOptions := transcoder.ScreenshotOptions{
|
||||
OutputPath: "-",
|
||||
OutputType: transcoder.ScreenshotOutputTypeBMP,
|
||||
Width: spriteScreenshotWidth,
|
||||
}
|
||||
|
||||
args := transcoder.ScreenshotTime(input, seconds, ssOptions)
|
||||
|
||||
return g.generateImage(lockCtx, args)
|
||||
}
|
||||
|
||||
func (g Generator) SpriteScreenshotSlow(ctx context.Context, input string, frame int) (image.Image, error) {
|
||||
lockCtx := g.LockManager.ReadLock(ctx, input)
|
||||
defer lockCtx.Cancel()
|
||||
|
||||
ssOptions := transcoder.ScreenshotOptions{
|
||||
OutputPath: "-",
|
||||
OutputType: transcoder.ScreenshotOutputTypeBMP,
|
||||
Width: spriteScreenshotWidth,
|
||||
}
|
||||
|
||||
args := transcoder.ScreenshotFrame(input, frame, ssOptions)
|
||||
|
||||
return g.generateImage(lockCtx, args)
|
||||
}
|
||||
|
||||
func (g Generator) generateImage(lockCtx *fsutil.LockContext, args ffmpeg.Args) (image.Image, error) {
|
||||
out, err := g.generateOutput(lockCtx, args)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
img, _, err := image.Decode(bytes.NewReader(out))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("decoding image from ffmpeg: %w", err)
|
||||
}
|
||||
|
||||
return img, nil
|
||||
}
|
||||
|
||||
func (g Generator) CombineSpriteImages(images []image.Image) image.Image {
|
||||
// Combine all of the thumbnails into a sprite image
|
||||
width := images[0].Bounds().Size().X
|
||||
height := images[0].Bounds().Size().Y
|
||||
canvasWidth := width * spriteCols
|
||||
canvasHeight := height * spriteRows
|
||||
montage := imaging.New(canvasWidth, canvasHeight, color.NRGBA{})
|
||||
for index := 0; index < len(images); index++ {
|
||||
x := width * (index % spriteCols)
|
||||
y := height * int(math.Floor(float64(index)/float64(spriteRows)))
|
||||
img := images[index]
|
||||
montage = imaging.Paste(montage, img, image.Pt(x, y))
|
||||
}
|
||||
|
||||
return montage
|
||||
}
|
||||
|
||||
func (g Generator) SpriteVTT(ctx context.Context, output string, spritePath string, stepSize float64) error {
|
||||
lockCtx := g.LockManager.ReadLock(ctx, spritePath)
|
||||
defer lockCtx.Cancel()
|
||||
|
||||
return g.generateFile(lockCtx, g.ScenePaths, vttPattern, output, g.spriteVTT(spritePath, stepSize))
|
||||
}
|
||||
|
||||
func (g Generator) spriteVTT(spritePath string, stepSize float64) generateFn {
|
||||
return func(lockCtx *fsutil.LockContext, tmpFn string) error {
|
||||
spriteImage, err := os.Open(spritePath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer spriteImage.Close()
|
||||
spriteImageName := filepath.Base(spritePath)
|
||||
image, _, err := image.DecodeConfig(spriteImage)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
width := image.Width / spriteCols
|
||||
height := image.Height / spriteRows
|
||||
|
||||
vttLines := []string{"WEBVTT", ""}
|
||||
for index := 0; index < spriteChunks; index++ {
|
||||
x := width * (index % spriteCols)
|
||||
y := height * int(math.Floor(float64(index)/float64(spriteRows)))
|
||||
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 os.WriteFile(tmpFn, []byte(vtt), 0644)
|
||||
}
|
||||
}
|
||||
|
||||
// TODO - move all sprite generation code here
|
||||
// WIP
|
||||
// func (g Generator) Sprite(ctx context.Context, videoFile *ffmpeg.VideoFile, hash string) error {
|
||||
// input := videoFile.Path
|
||||
// if err := g.generateSpriteImage(ctx, videoFile, hash); err != nil {
|
||||
// return fmt.Errorf("generating sprite image for %s: %w", input, err)
|
||||
// }
|
||||
|
||||
// output := g.ScenePaths.GetSpriteVttFilePath(hash)
|
||||
// if !g.Overwrite {
|
||||
// if exists, _ := fsutil.FileExists(output); exists {
|
||||
// return nil
|
||||
// }
|
||||
// }
|
||||
|
||||
// if err := g.generateFile(ctx, g.ScenePaths, vttPattern, output, g.spriteVtt(input, screenshotOptions{
|
||||
// Time: at,
|
||||
// Quality: screenshotQuality,
|
||||
// // default Width is video width
|
||||
// })); err != nil {
|
||||
// return err
|
||||
// }
|
||||
|
||||
// logger.Debug("created screenshot: ", output)
|
||||
|
||||
// return nil
|
||||
// }
|
||||
|
||||
// func (g Generator) generateSpriteImage(ctx context.Context, videoFile *ffmpeg.VideoFile, hash string) error {
|
||||
// output := g.ScenePaths.GetSpriteImageFilePath(hash)
|
||||
// if !g.Overwrite {
|
||||
// if exists, _ := fsutil.FileExists(output); exists {
|
||||
// return nil
|
||||
// }
|
||||
// }
|
||||
|
||||
// var images []image.Image
|
||||
// var err error
|
||||
// if options.VideoDuration > 0 {
|
||||
// images, err = g.generateSprites(ctx, input, options.VideoDuration)
|
||||
// } else {
|
||||
// images, err = g.generateSpritesSlow(ctx, input, options.FrameCount)
|
||||
// }
|
||||
|
||||
// if len(images) == 0 {
|
||||
// return errors.New("images slice is empty")
|
||||
// }
|
||||
|
||||
// montage, err := g.combineSpriteImages(images)
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
|
||||
// if err := imaging.Save(montage, output); err != nil {
|
||||
// return err
|
||||
// }
|
||||
|
||||
// logger.Debug("created sprite image: ", output)
|
||||
|
||||
// return nil
|
||||
// }
|
||||
|
||||
// func useSlowSeek(videoFile *ffmpeg.VideoFile) (bool, error) {
|
||||
// // For files with small duration / low frame count try to seek using frame number intead of seconds
|
||||
// // some files can have FrameCount == 0, only use SlowSeek if duration < 5
|
||||
// if videoFile.Duration < 5 || (videoFile.FrameCount > 0 && videoFile.FrameCount <= int64(spriteChunks)) {
|
||||
// if videoFile.Duration <= 0 {
|
||||
// return false, fmt.Errorf("duration(%.3f)/frame count(%d) invalid", videoFile.Duration, videoFile.FrameCount)
|
||||
// }
|
||||
|
||||
// logger.Warnf("[generator] video %s too short (%.3fs, %d frames), using frame seeking", videoFile.Path, videoFile.Duration, videoFile.FrameCount)
|
||||
// return true, nil
|
||||
// }
|
||||
// }
|
||||
|
||||
// func (g Generator) combineSpriteImages(images []image.Image) (image.Image, error) {
|
||||
// // Combine all of the thumbnails into a sprite image
|
||||
// width := images[0].Bounds().Size().X
|
||||
// height := images[0].Bounds().Size().Y
|
||||
// canvasWidth := width * spriteCols
|
||||
// canvasHeight := height * spriteRows
|
||||
// montage := imaging.New(canvasWidth, canvasHeight, color.NRGBA{})
|
||||
// for index := 0; index < len(images); index++ {
|
||||
// x := width * (index % spriteCols)
|
||||
// y := height * int(math.Floor(float64(index)/float64(spriteRows)))
|
||||
// img := images[index]
|
||||
// montage = imaging.Paste(montage, img, image.Pt(x, y))
|
||||
// }
|
||||
|
||||
// return montage, nil
|
||||
// }
|
||||
|
||||
// func (g Generator) generateSprites(ctx context.Context, input string, videoDuration float64) ([]image.Image, error) {
|
||||
// logger.Infof("[generator] generating sprite image for %s", input)
|
||||
// // generate `ChunkCount` thumbnails
|
||||
// stepSize := videoDuration / float64(spriteChunks)
|
||||
|
||||
// var images []image.Image
|
||||
// for i := 0; i < spriteChunks; i++ {
|
||||
// time := float64(i) * stepSize
|
||||
|
||||
// img, err := g.spriteScreenshot(ctx, input, time)
|
||||
// if err != nil {
|
||||
// return nil, err
|
||||
// }
|
||||
// images = append(images, img)
|
||||
// }
|
||||
|
||||
// return images, nil
|
||||
// }
|
||||
|
||||
// func (g Generator) generateSpritesSlow(ctx context.Context, input string, frameCount int) ([]image.Image, error) {
|
||||
// logger.Infof("[generator] generating sprite image for %s (%d frames)", input, frameCount)
|
||||
|
||||
// stepFrame := float64(frameCount-1) / float64(spriteChunks)
|
||||
|
||||
// var images []image.Image
|
||||
// for i := 0; i < spriteChunks; i++ {
|
||||
// // generate exactly `ChunkCount` thumbnails, using duplicate frames if needed
|
||||
// frame := math.Round(float64(i) * stepFrame)
|
||||
// if frame >= math.MaxInt || frame <= math.MinInt {
|
||||
// return nil, errors.New("invalid frame number conversion")
|
||||
// }
|
||||
|
||||
// img, err := g.spriteScreenshotSlow(ctx, input, int(frame))
|
||||
// if err != nil {
|
||||
// return nil, err
|
||||
// }
|
||||
// images = append(images, img)
|
||||
// }
|
||||
|
||||
// return images, nil
|
||||
// }
|
||||
|
||||
// func (g Generator) spriteScreenshot(ctx context.Context, input string, seconds float64) (image.Image, error) {
|
||||
// ssOptions := transcoder.ScreenshotOptions{
|
||||
// OutputPath: "-",
|
||||
// OutputType: transcoder.ScreenshotOutputTypeBMP,
|
||||
// Width: spriteScreenshotWidth,
|
||||
// }
|
||||
|
||||
// args := transcoder.ScreenshotTime(input, seconds, ssOptions)
|
||||
|
||||
// return g.generateImage(ctx, args)
|
||||
// }
|
||||
|
||||
// func (g Generator) spriteScreenshotSlow(ctx context.Context, input string, frame int) (image.Image, error) {
|
||||
// ssOptions := transcoder.ScreenshotOptions{
|
||||
// OutputPath: "-",
|
||||
// OutputType: transcoder.ScreenshotOutputTypeBMP,
|
||||
// Width: spriteScreenshotWidth,
|
||||
// }
|
||||
|
||||
// args := transcoder.ScreenshotFrame(input, frame, ssOptions)
|
||||
|
||||
// return g.generateImage(ctx, args)
|
||||
// }
|
||||
|
||||
// func (g Generator) spriteVTT(videoFile ffmpeg.VideoFile, spriteImagePath string, slowSeek bool) generateFn {
|
||||
// return func(ctx context.Context, tmpFn string) error {
|
||||
// logger.Infof("[generator] generating sprite vtt for %s", input)
|
||||
|
||||
// spriteImage, err := os.Open(spriteImagePath)
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
// defer spriteImage.Close()
|
||||
// spriteImageName := filepath.Base(spriteImagePath)
|
||||
// image, _, err := image.DecodeConfig(spriteImage)
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
// width := image.Width / spriteCols
|
||||
// height := image.Height / spriteRows
|
||||
|
||||
// var stepSize float64
|
||||
// if !slowSeek {
|
||||
// nthFrame = g.NumberOfFrames / g.ChunkCount
|
||||
// stepSize = float64(g.Info.NthFrame) / g.Info.FrameRate
|
||||
// } else {
|
||||
// // for files with a low framecount (<ChunkCount) g.Info.NthFrame can be zero
|
||||
// // so recalculate from scratch
|
||||
// stepSize = float64(videoFile.FrameCount-1) / float64(spriteChunks)
|
||||
// stepSize /= g.Info.FrameRate
|
||||
// }
|
||||
|
||||
// vttLines := []string{"WEBVTT", ""}
|
||||
// for index := 0; index < spriteChunks; index++ {
|
||||
// x := width * (index % spriteCols)
|
||||
// y := height * int(math.Floor(float64(index)/float64(spriteRows)))
|
||||
// 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 os.WriteFile(tmpFn, []byte(vtt), 0644)
|
||||
// }
|
||||
// }
|
||||
Reference in New Issue
Block a user