animate/animate.go

// This package enables creating simple animations by drawing onto an image, using ffmpeg which could be in the system path.

package animate

import (
	"cmp"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"
)

// Animation implements the [draw.Image] interface; you can draw to it and then invoke [EmitFrame].
//
// That said, some implementations might be optimized for known image types, in which case, passing a pointer to the
// [NRGBA] field might be more efficient.
type Animation struct {
	image.NRGBA

	// the first error where we encountered during writing
	err error

	// the pipe to ffmpeg's stdin, nil if the animation is closed
	w io.WriteCloser

	// the ffmpeg process, only applicable if w above is not nil
	c *exec.Cmd

	// the name of the file we're writing to,
	// and where we'll move it to if nothing explodes
	tmp, out string

	// framerate
	rate float64

	// the current frame number
	frame int
}

// EmitFrame writes the image as it presently appears as a frame to the animation.
//
// If this is successful, then the frame counter is advanced. The frame buffer isn't cleared.
func (a *Animation) EmitFrame() error {
	// if we've already encountered an error, return it
	if a.err != nil {
		return a.err
	}

	// if the pipe is closed, then we can't write, obviously.
	if a.w == nil {
		return os.ErrClosed
	}

	// try to write the frame to the pipe,
	// if if that fails, record the error and close the animation.
	if _, err := a.w.Write(a.Pix); err != nil {
		a.err = err
		return a.Close()
	}

	a.frame++

	// success!
	return nil
}

// Frame returns the current frame number, with the first frame being 0.
func (a *Animation) Frame() int {
	return a.frame
}

// Elapsed returns the timestamp for the current frame, with first frame having a timestamp of 0.
func (a *Animation) Elapsed() time.Duration {
	return time.Duration(float64(a.frame) * float64(time.Second) / a.rate)
}

// Close finishes writing the animation, and cleans up any resources.
func (a *Animation) Close() error {
	// the pipe is open, shut everything down.
	if a.w != nil {
		// no longer needs to be finalized
		runtime.SetFinalizer(a, nil)

		// close the pipe and wait for ffmpeg to finish,
		// setting a.err to the first error encountered.
		a.err = cmp.Or(a.err, a.w.Close(), a.c.Wait())
		a.w, a.c = nil, nil

		// if there are no errors so far, attempt to move the file to the destination.
		if a.err == nil {
			a.err = os.Rename(a.tmp, a.out)
		}

		// if there was an error, delete the temporary file we wrote.
		if a.err != nil {
			os.Remove(a.tmp)
		}
	}

	// return any error we may have encountered.
	return a.err
}

func paletteToFile(palette color.Palette) (*os.File, error) {
	r, w, err := os.Pipe()
	if err != nil {
		return nil, err
	}

	go func() {
		defer w.Close()

		img := image.NewPaletted(image.Rect(0, 0, 16, 16), palette)

		for i := 0; i < 256; i++ {
			img.Pix[i] = uint8(min(i, len(palette)-1))
		}

		// I'm not going to bother dealing with possible write errors,
		// those would be because ffmpeg terminated unexpectedly,
		// and has nothing to do with us.
		png.Encode(w, img)
	}()

	return r, nil
}

func finalize(a *Animation) {
	err := a.Close()
	log.Printf("animation was finalized without being closed; err=%v", err)
}

func new(r image.Rectangle, rate float64, palette color.Palette, out string, fmtArgs []string) (a *Animation, err error) {
	tmp, err := os.CreateTemp(filepath.Dir(out), "."+strings.TrimSuffix(filepath.Base(out), filepath.Ext(out))+"*.tmp")

	if err != nil {
		return nil, err
	}

	defer func() {
		// if things go wrong, delete the file we wrote.
		if err != nil {
			os.Remove(tmp.Name())
		}
	}()

	// we don't need the file open, we just want a safe place for ffmpeg to write to.
	if err := tmp.Close(); err != nil {
		return nil, err
	}

	var args []string

	args = append(args,
		"-hide_banner",
		"-loglevel", "error",
		"-f", "rawvideo",
		"-pixel_format", "rgba",
		"-video_size", fmt.Sprintf("%dx%d", r.Dx(), r.Dy()),
		"-framerate", fmt.Sprint(rate),
		"-i", "pipe:0",
	)

	var extraFiles []*os.File

	if palette != nil {
		f, err := paletteToFile(palette)
		if err != nil {
			return nil, err
		}

		args = append(args,
			"-i", fmt.Sprintf("pipe:%d", len(extraFiles)+3),
			"-filter_complex", "[0:v][1:v]paletteuse=dither=bayer",
		)

		extraFiles = append(extraFiles, f)
	}

	args = append(args, fmtArgs...)

	args = append(args, "-y", "file:"+tmp.Name())

	cmd := exec.Command("ffmpeg", args...)

	cmd.ExtraFiles = extraFiles
	cmd.Stderr = os.Stderr

	w, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}

	err = cmd.Start()

	// close our copies of the file descriptors. The ffmpeg process, if it even started, has inherited them.
	// and if it didn't, we'd needed to close them ourselves.
	for _, f := range extraFiles {
		f.Close()
	}

	if err != nil {
		return nil, err
	}

	animation := &Animation{
		NRGBA: image.NRGBA{
			Pix:    make([]uint8, r.Dx()*r.Dy()*4),
			Stride: r.Dx() * 4,
			Rect:   r,
		},
		w:     w,
		c:     cmd,
		tmp:   tmp.Name(),
		out:   out,
		rate:  rate,
		frame: 0,
	}

	// failing to close the animation will leave temporary files on the filesystem, possibly leave the ffmpeg process hanging,
	// so let's try not to let that happen.
	runtime.SetFinalizer(animation, finalize)

	return animation, nil
}

// New begins a new animation at the given size and frame rate, to be written to the specified output file.
//
// The output filename must have one of the following extensions:
// - .gif (lossy, uses a default fixed palette)
// - .png / .apng (lossless)
// - .mp4 (lossless RGB, doesn't support transparency)
// - .mkv (lossless RGB, doesn't support transparency)
// - .webp (lossy YUV colorspace conversion, alpha supported)
// - .webm (lossy YUV colorspace conversion, alpha supported)
//
// Note that the output file won't be created immediately, the animation will be written to a temporary file and then renamed
// by the [Close] method if no errors are encountered.
func New(r image.Rectangle, rate float64, out string) (*Animation, error) {
	switch ext := filepath.Ext(out); ext {
	case ".gif":
		return new(r, rate, defaultPalette, out, []string{
			"-loop", "0",
			"-f", "gif",
		})
	case ".png", ".apng":
		return new(r, rate, nil, out, []string{
			"-vcodec", "apng",
			"-plays", "0",
			"-f", "apng",
		})
	case ".mp4":
		return new(r, rate, nil, out, []string{
			"-vcodec", "libx264rgb",
			"-preset", "ultrafast",
			"-qp", "0",
			"-f", "mp4",
		})
	case ".mkv":
		return new(r, rate, nil, out, []string{
			"-vcodec", "libx264rgb",
			"-preset", "ultrafast",
			"-qp", "0",
			"-f", "matroska",
		})
	case ".webp":
		return new(r, rate, nil, out, []string{
			"-vcodec", "libwebp_anim",
			"-lossless", "1",
			"-loop", "0",
			"-f", "webp",
		})
	case ".webm":
		return new(r, rate, nil, out, []string{
			"-vcodec", "vp9",
			"-lossless", "1",
			"-f", "webm",
		})
	default:
		return nil, fmt.Errorf("unsupported extension: %q", ext)
	}
}