Files
thumbnail-gen/generator.go

317 lines
7.0 KiB
Go

package thumbnailgen
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"os/exec"
"regexp"
"slices"
"strconv"
"strings"
"time"
)
type TimeFilter struct {
Start float64
End float64
}
type Options struct {
Interval int
Segments int
UseSegments bool
Format string
Scale string
EnableFilter bool
MaxThumbnails int
}
func NewDefaultOptions() Options {
return Options{
Interval: 5,
Segments: 0,
UseSegments: false,
Format: "webp",
Scale: "-1:-1",
EnableFilter: true,
MaxThumbnails: 0,
}
}
func (options *Options) Apply(opts ...func(*Options)) *Options {
for _, opt := range opts {
opt(options)
}
return options
}
type videoStreamProbe struct {
Streams []struct {
CodecType string `json:"codec_type"`
Duration string `json:"duration"`
} `json:"streams"`
}
func getThumbnailDuration(path string) (float64, error) {
stderr := bytes.NewBuffer(nil)
stdout := bytes.NewBuffer(nil)
cmd := exec.Command(
"ffprobe",
"-v", "error",
"-select_streams", "v:0",
"-show_entries", "stream=codec_type,duration",
"-of", "json",
path,
)
cmd.Stdout = stdout
cmd.Stderr = stderr
if err := cmd.Run(); err != nil {
return 0, fmt.Errorf(
"probe video stream: %w: %s",
err,
strings.TrimSpace(stderr.String()),
)
}
var probe videoStreamProbe
if err := json.Unmarshal(stdout.Bytes(), &probe); err != nil {
return 0, fmt.Errorf("decode video-stream probe: %w", err)
}
if len(probe.Streams) == 0 || probe.Streams[0].CodecType != "video" {
return 0, errors.New("input has no video stream")
}
durationText := strings.TrimSpace(probe.Streams[0].Duration)
if durationText == "" || durationText == "N/A" {
return GetVideoLength(path) // container-duration fallback for animated WebP
}
duration, err := strconv.ParseFloat(durationText, 64)
if err != nil {
return 0, fmt.Errorf("parse video-stream duration %q: %w", durationText, err)
}
if duration <= 0 {
return 0, fmt.Errorf("invalid video-stream duration %v", duration)
}
return duration, nil
}
func (opts *Options) GetThumbnail(path string) ([][]byte, int, error) {
var filters []TimeFilter
if opts.EnableFilter {
f, err := GetFilter(path)
if err != nil {
return nil, 0, err
}
filters = f
}
length, err := getThumbnailDuration(path)
if err != nil {
return nil, 0, err
}
var interval int
if opts.UseSegments {
interval = int(length / float64(opts.Segments))
} else {
interval = opts.Interval
}
return GetThumbnailUnderlying(path, opts.MaxThumbnails, filters, length, interval, opts.EnableFilter, opts.Format, opts.Scale)
}
func GetThumbnailUnderlying(path string, maxThumbnails int, filters []TimeFilter, length float64, intervalSeconds int, enableFilter bool, format string, scale string) ([][]byte, int, error) {
buf := bytes.NewBuffer(nil)
framesExtracted := 0
var out [][]byte
if maxThumbnails > 0 {
out = make([][]byte, maxThumbnails)
} else {
out = make([][]byte, 0)
}
var time float64 = 1
for {
if (maxThumbnails > 0 && framesExtracted >= maxThumbnails) || time >= length {
break
}
if enableFilter {
if ok, next := FrameLiesWithinFilter(time, filters); ok {
time = next + 1
continue
}
}
err := GetImage(buf, path, int(time), format, scale)
if err != nil {
return nil, 0, err
}
b := bytes.Clone(buf.Bytes())[0:buf.Len()]
if maxThumbnails > 0 {
out[framesExtracted] = b
} else {
out = append(out, b)
}
time += float64(intervalSeconds)
framesExtracted += 1
buf.Reset()
}
return out, framesExtracted, nil
}
// 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, filter.End
}
}
return false, time
}
func GetImage(buf *bytes.Buffer, path string, timestamp int, format string, scale string) error {
var t time.Time
t = t.Add(time.Duration(timestamp) * time.Second)
stderr := bytes.NewBuffer(nil)
cmd := exec.Command(
"ffmpeg",
"-ss",
t.Format("15:04:05"),
"-i",
path,
"-vframes",
"1",
"-c:v",
format,
"-filter:v",
fmt.Sprintf("scale=%s", scale),
"-f",
"image2pipe",
"-",
)
cmd.Stdout = buf
cmd.Stderr = stderr
if err := cmd.Run(); err != nil {
return fmt.Errorf("generate thumbnail at %s: %w: %s", t.Format("15:04:05"), err, strings.TrimSpace(stderr.String()))
}
return nil
}
func GetVideoLength(path string) (float64, error) {
cmd := exec.Command("ffprobe", "-v", "quiet", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", path)
stderr := bytes.NewBuffer(nil)
buf := bytes.NewBuffer(nil)
cmd.Stdout = buf
cmd.Stderr = stderr
if err := cmd.Run(); err != nil {
return 0, fmt.Errorf(
"probe video duration: %w: %s",
err,
strings.TrimSpace(stderr.String()),
)
}
return strconv.ParseFloat(strings.ReplaceAll(buf.String(), "\n", ""), 64)
}
func escapeFilterValue(value string) string {
return strings.NewReplacer(
"\\", "\\\\",
"'", "\\'",
":", "\\:",
",", "\\,",
";", "\\;",
"[", "\\[",
"]", "\\]",
).Replace(value)
}
func GetFilter(path string) ([]TimeFilter, error) {
buf := bytes.NewBuffer(nil)
stderr := bytes.NewBuffer(nil)
filter := fmt.Sprintf("movie=filename='%s',blackdetect[out0]", escapeFilterValue(path))
cmd := exec.Command("ffprobe", "-f", "lavfi", "-i", filter, "-show_entries", "tags=lavfi.black_start,lavfi.black_end", "-of", "default=nw=1", "-v", "quiet")
cmd.Stderr = stderr
cmd.Stdout = buf
if err := cmd.Run(); err != nil {
return nil, fmt.Errorf("detect black frames: %w: %s", err, strings.TrimSpace(stderr.String()))
}
filterStr := buf.String()
filters := strings.Split(filterStr, "\n")
filters = slices.Compact(filters)
i := 0
reg := regexp.MustCompile(`([0-9]*\.?[0-9]+)`)
blackFilters := make([]TimeFilter, len(filters)/2)
for {
if i >= len(filters) {
break
}
start := reg.FindString(filters[i])
var end string
if i+1 >= len(filters) || filters[i+1] == "" {
end = start
} else {
end = reg.FindString(filters[i+1])
}
if start == "" || end == "" {
break
}
s, err := strconv.ParseFloat(start, 64)
if err != nil {
return nil, err
}
e, err := strconv.ParseFloat(end, 64)
if err != nil {
return nil, err
}
blackFilters[i/2] = TimeFilter{
Start: s,
End: e,
}
i += 2
}
return blackFilters, nil
}
func GetFramerate(path string) (float64, error) {
cmd := exec.Command("ffprobe", "-v", "0", "-of", "csv=p=0", "-select_streams", "v:0", "-show_entries", "stream=r_frame_rate", path)
buf := bytes.NewBuffer(nil)
cmd.Stdout = buf
err := cmd.Run()
if err != nil {
return 0, err
}
s := buf.String()
s = strings.ReplaceAll(s, "\n", "")
s = strings.TrimSpace(s)
ss := strings.Split(s, "/")
f1, err := strconv.Atoi(ss[0])
if err != nil {
return 0, err
}
f2, err := strconv.Atoi(ss[1])
if err != nil {
return 0, err
}
return float64(f1) / float64(f2), nil
}