From abbc2fa230bd4e63491de4e51c3a2dbc2081a467 Mon Sep 17 00:00:00 2001
From: Stash Dev <stashappdev@gmail.com>
Date: Thu, 28 Mar 2019 15:58:13 -0700
Subject: [PATCH] Another tweak to fix sprite generation

---
 pkg/ffmpeg/encoder.go    | 15 ++-------------
 pkg/ffmpeg/regex.go      | 34 ++++++++++++++++++++++++++++++++++
 pkg/manager/generator.go | 34 +++++++++++++++++++++-------------
 3 files changed, 57 insertions(+), 26 deletions(-)
 create mode 100644 pkg/ffmpeg/regex.go

diff --git a/pkg/ffmpeg/encoder.go b/pkg/ffmpeg/encoder.go
index b5ac39a8c..47c780f3a 100644
--- a/pkg/ffmpeg/encoder.go
+++ b/pkg/ffmpeg/encoder.go
@@ -4,12 +4,8 @@ import (
 	"github.com/stashapp/stash/pkg/logger"
 	"io/ioutil"
 	"os/exec"
-	"regexp"
-	"strconv"
 )
 
-var progressRegex = regexp.MustCompile(`time=(\d+):(\d+):(\d+.\d+)`)
-
 type Encoder struct {
 	Path string
 }
@@ -42,15 +38,8 @@ func (e *Encoder) run(probeResult VideoFile, args []string) (string, error) {
 		n, err := stderr.Read(buf)
 		if n > 0 {
 			data := string(buf[0:n])
-			regexResult := progressRegex.FindStringSubmatch(data)
-			if len(regexResult) == 4 && probeResult.Duration > 0 {
-				h, _ := strconv.ParseFloat(regexResult[1], 64)
-				m, _ := strconv.ParseFloat(regexResult[2], 64)
-				s, _ := strconv.ParseFloat(regexResult[3], 64)
-				hours := h * 3600
-				mins := m * 60
-				secs := s
-				time := hours + mins + secs
+			time := GetTimeFromRegex(data)
+			if time > 0 && probeResult.Duration > 0 {
 				progress := time / probeResult.Duration
 				logger.Infof("Progress %.2f", progress)
 			}
diff --git a/pkg/ffmpeg/regex.go b/pkg/ffmpeg/regex.go
new file mode 100644
index 000000000..3ffe0fbfc
--- /dev/null
+++ b/pkg/ffmpeg/regex.go
@@ -0,0 +1,34 @@
+package ffmpeg
+
+import (
+	"regexp"
+	"strconv"
+)
+
+var TimeRegex = regexp.MustCompile(`time=\s*(\d+):(\d+):(\d+.\d+)`)
+var FrameRegex = regexp.MustCompile(`frame=\s*([0-9]+)`)
+
+func GetTimeFromRegex(str string) float64 {
+	regexResult := TimeRegex.FindStringSubmatch(str)
+
+	// Bail early if we don't have the results we expect
+	if len(regexResult) != 4 { return 0 }
+
+	h, _ := strconv.ParseFloat(regexResult[1], 64)
+	m, _ := strconv.ParseFloat(regexResult[2], 64)
+	s, _ := strconv.ParseFloat(regexResult[3], 64)
+	hours := h * 3600
+	minutes := m * 60
+	seconds := s
+	return hours + minutes + seconds
+}
+
+func GetFrameFromRegex(str string) int {
+	regexResult := FrameRegex.FindStringSubmatch(str)
+
+	// Bail early if we don't have the results we expect
+	if len(regexResult) < 2 { return 0 }
+
+	result, _ := strconv.Atoi(regexResult[1])
+	return result
+}
\ No newline at end of file
diff --git a/pkg/manager/generator.go b/pkg/manager/generator.go
index ca1d0368a..4a7a1bf16 100644
--- a/pkg/manager/generator.go
+++ b/pkg/manager/generator.go
@@ -6,8 +6,8 @@ import (
 	"github.com/stashapp/stash/pkg/ffmpeg"
 	"github.com/stashapp/stash/pkg/logger"
 	"github.com/stashapp/stash/pkg/utils"
+	"math"
 	"os/exec"
-	"regexp"
 	"runtime"
 	"strconv"
 )
@@ -44,10 +44,15 @@ func (g *GeneratorInfo) configure() error {
 	} else {
 		framerate = g.VideoFile.FrameRate
 	}
-	g.FrameRate = framerate
 
 	numberOfFrames, _ := strconv.Atoi(videoStream.NbFrames)
-	if numberOfFrames == 0 {
+
+	if numberOfFrames == 0 && framerate > 0 && g.VideoFile.Duration > 0 { // TODO: test
+		numberOfFrames = int(framerate * g.VideoFile.Duration)
+	}
+
+	// If we are missing the frame count or frame rate then seek through the file and extract the info with regex
+	if numberOfFrames == 0 || framerate == 0 {
 		args := []string{
 			"-nostats",
 			"-i", g.VideoFile.Path,
@@ -65,25 +70,28 @@ func (g *GeneratorInfo) configure() error {
 		var stdErrBuffer bytes.Buffer
 		command.Stderr = &stdErrBuffer // Frames go to stderr rather than stdout
 		if err := command.Run(); err == nil {
-			re := regexp.MustCompile(`frame[=] ([0-9]+)`)
-			frames := re.FindStringSubmatch(stdErrBuffer.String())
-			if frames != nil && len(frames) > 1 {
-				numberOfFrames, _ = strconv.Atoi(frames[1])
+			stdErrString := stdErrBuffer.String()
+			if numberOfFrames == 0 {
+				numberOfFrames = ffmpeg.GetFrameFromRegex(stdErrString)
+			}
+			if framerate == 0 {
+				time := ffmpeg.GetTimeFromRegex(stdErrString)
+				framerate = math.Round((float64(numberOfFrames)/time)*100) / 100
 			}
 		}
 	}
-	if numberOfFrames == 0 { // TODO: test
-		numberOfFrames = int(framerate * g.VideoFile.Duration)
-	}
-	if numberOfFrames == 0 {
+
+	// Something seriously wrong with this file
+	if numberOfFrames == 0 || framerate == 0 {
 		logger.Errorf(
-			"number of frames is 0.  nb_frames <%s> framerate <%s> duration <%s>",
+			"number of frames or framerate is 0.  nb_frames <%s> framerate <%s> duration <%s>",
 			videoStream.NbFrames,
 			framerate,
 			g.VideoFile.Duration,
-			)
+		)
 	}
 
+	g.FrameRate = framerate
 	g.NumberOfFrames = numberOfFrames
 	g.NthFrame = g.NumberOfFrames / g.ChunkCount