Now splittable into segments with variable intervals, and more Cli options

2 files changed, 50 insertions, 31 deletions

diff --git a/cmd/thumbnail-gen/main.go b/cmd/thumbnail-gen/main.go
index 8783685..4837770 100644
--- a/cmd/thumbnail-gen/main.go
+++ b/cmd/thumbnail-gen/main.go
@@ -3,6 +3,8 @@ package main
 import (
 	"flag"
 	"fmt"
+	"os"
+	"path/filepath"
 
 	thumbnailgen "github.com/pablu23/thumbnail-gen"
 )
@@ -11,37 +13,27 @@ var (
 	pathFlag          = flag.String("path", "", "path to video")
 	intervalFlag      = flag.Int("interval", 20, "Interval in seconds between thumbnails")
 	maxThumbnailsFlag = flag.Int("max", 0, "Max Thumbnails, default as much as possible with interval")
+	blackFilterFlag   = flag.Bool("filter", true, "Try to filter out black frames, might be an expensive operation")
 )
 
 func main() {
 	flag.Parse()
 
 	if *pathFlag != "" {
-    _, err := thumbnailgen.GetFramerate(*pathFlag)
+		// thumbnails, num, err := thumbnailgen.GetThumbnail(*pathFlag, *intervalFlag, *maxThumbnailsFlag, *blackFilterFlag)
+		thumbnails, num, err := thumbnailgen.GetThumbnailSegments(*pathFlag, 12, *blackFilterFlag)
 		if err != nil {
-			fmt.Printf("Framerate: %s\n", err)
+			panic(err)
 		}
 
-		_, err = thumbnailgen.GetVideoLength(*pathFlag)
-		if err != nil {
-			fmt.Printf("Video Length: %s\n", err)
-		}
+		fmt.Printf("Extracted %d thumbnails\n", num)
 
-    _, err = thumbnailgen.GetFilter(*pathFlag)
-    if err != nil {
-			fmt.Printf("Filter: %s\n", err)
+		name := filepath.Base(*pathFlag)
+		for i, thumbnail := range thumbnails[:num] {
+			err := os.WriteFile(fmt.Sprintf("%s-%d.png", name, i), thumbnail, 0600)
+			if err != nil {
+				panic(err)
+			}
 		}
-		// thumbnails, err := thumbnailgen.GetThumbnail(*pathFlag, *intervalFlag, *maxThumbnailsFlag)
-		// if err != nil {
-		// 	panic(err)
-		// }
-		//
-		// name := filepath.Base(*pathFlag)
-		// for i, thumbnail := range thumbnails {
-		// 	err := os.WriteFile(fmt.Sprintf("%s-%d.png", name, i), thumbnail, 0600)
-		// 	if err != nil {
-		// 		panic(err)
-		// 	}
-		// }
 	}
 }
diff --git a/generator.go b/generator.go
index a3c6e85..6f5cbff 100644
--- a/generator.go
+++ b/generator.go
@@ -16,7 +16,7 @@ type TimeFilter struct {
 	End   float64
 }
 
-func GetThumbnail(path string, intervalSeconds int, maxThumbnails int, enableFilter bool) ([][]byte, int, error) {
+func GetThumbnailSegments(path string, segments int, enableFilter bool) ([][]byte, int, error) {
 	var filters []TimeFilter
 	if enableFilter {
 		f, err := GetFilter(path)
@@ -26,18 +26,36 @@ func GetThumbnail(path string, intervalSeconds int, maxThumbnails int, enableFil
 		filters = f
 	}
 
-	fps, err := GetFramerate(path)
+	length, err := GetVideoLength(path)
 	if err != nil {
 		return nil, 0, err
 	}
 
+	interval := length / float64(segments)
+
+	return getThumbnailUnderlying(path, 0, filters, length, int(interval), enableFilter)
+}
+
+func GetThumbnail(path string, intervalSeconds int, maxThumbnails int, enableFilter bool) ([][]byte, int, error) {
+	var filters []TimeFilter
+	if enableFilter {
+		f, err := GetFilter(path)
+		if err != nil {
+			return nil, 0, err
+		}
+		filters = f
+	}
+
 	length, err := GetVideoLength(path)
 	if err != nil {
 		return nil, 0, err
 	}
 
+	return getThumbnailUnderlying(path, maxThumbnails, filters, length, intervalSeconds, enableFilter)
+}
+
+func getThumbnailUnderlying(path string, maxThumbnails int, filters []TimeFilter, length float64, intervalSeconds int, enableFilter bool) ([][]byte, int, error) {
 	buf := bytes.NewBuffer(nil)
-	currFrame := 0
 	framesExtracted := 0
 
 	var out [][]byte
@@ -47,15 +65,21 @@ func GetThumbnail(path string, intervalSeconds int, maxThumbnails int, enableFil
 		out = make([][]byte, 0)
 	}
 
+	// fmt.Printf("FPS: %f, Lenght: %f, Interval: %d, Filters: %t\n", fps, length, intervalSeconds, enableFilter)
+
+	var time float64 = 1
 	for {
-		time := float64(currFrame) / fps
+		fmt.Printf("Time: %f\n", time)
 		if (maxThumbnails > 0 && framesExtracted >= maxThumbnails) || time >= length {
 			break
 		}
 
-		currFrame += int(fps) * intervalSeconds
-		if enableFilter && FrameLiesWithinFilter(time, filters) {
-			continue
+		if enableFilter {
+			if ok, next := FrameLiesWithinFilter(time, filters); ok {
+				fmt.Printf("Black Frame, skipping to: %f\n", next)
+				time = next
+				continue
+			}
 		}
 
 		err := GetImage(buf, path, int(time), "png")
@@ -71,20 +95,23 @@ func GetThumbnail(path string, intervalSeconds int, maxThumbnails int, enableFil
 			out = append(out, b)
 		}
 
+		time += float64(intervalSeconds)
 		framesExtracted += 1
 		buf.Reset()
 	}
 	return out, framesExtracted, nil
 }
 
-func FrameLiesWithinFilter(time float64, filters []TimeFilter) bool {
+// Return true, if time is inside filter, also returns next time, when not in filter anymore.
+// If False it returns input time
+func FrameLiesWithinFilter(time float64, filters []TimeFilter) (bool, float64) {
 	for _, filter := range filters {
 		if time >= filter.Start && time <= filter.End {
-			return true
+			return true, filter.End
 		}
 
 	}
-	return false
+	return false, time
 }
 
 func GetImage(buf *bytes.Buffer, path string, timestamp int, format string) error {