mirror of https://github.com/stashapp/stash.git
122 lines
3.0 KiB
Go
122 lines
3.0 KiB
Go
package manager
|
|
|
|
import (
|
|
"fmt"
|
|
"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"
|
|
"github.com/stashapp/stash/pkg/utils"
|
|
)
|
|
|
|
type PhashGenerator struct {
|
|
Info *GeneratorInfo
|
|
|
|
VideoChecksum string
|
|
Columns int
|
|
Rows int
|
|
}
|
|
|
|
func NewPhashGenerator(videoFile ffmpeg.VideoFile, checksum string) (*PhashGenerator, error) {
|
|
exists, err := utils.FileExists(videoFile.Path)
|
|
if !exists {
|
|
return nil, err
|
|
}
|
|
|
|
generator, err := newGeneratorInfo(videoFile)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &PhashGenerator{
|
|
Info: generator,
|
|
VideoChecksum: checksum,
|
|
Columns: 5,
|
|
Rows: 5,
|
|
}, nil
|
|
}
|
|
|
|
func (g *PhashGenerator) Generate() (*uint64, error) {
|
|
encoder := ffmpeg.NewEncoder(instance.FFMPEGPath)
|
|
|
|
sprite, err := g.generateSprite(&encoder)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
hash, err := goimagehash.PerceptionHash(sprite)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
hashValue := hash.GetHash()
|
|
return &hashValue, nil
|
|
}
|
|
|
|
func (g *PhashGenerator) generateSprite(encoder *ffmpeg.Encoder) (image.Image, error) {
|
|
logger.Infof("[generator] generating phash sprite for %s", g.Info.VideoFile.Path)
|
|
|
|
// Generate sprite image offset by 5% on each end to avoid intro/outros
|
|
chunkCount := g.Columns * g.Rows
|
|
offset := 0.05 * g.Info.VideoFile.Duration
|
|
stepSize := (0.9 * g.Info.VideoFile.Duration) / float64(chunkCount)
|
|
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
|
|
}
|
|
}
|
|
|
|
// 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)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
images = append(images, img)
|
|
}
|
|
|
|
if len(images) == 0 {
|
|
return nil, fmt.Errorf("images slice is empty, failed to generate phash sprite for %s", g.Info.VideoFile.Path)
|
|
}
|
|
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))
|
|
}
|
|
|
|
for _, imagePath := range imagePaths {
|
|
os.Remove(imagePath)
|
|
}
|
|
|
|
return montage, nil
|
|
}
|