Merge pull request #7 from Krzysztofz01/frame-seeking-feature-implementation

Frame seeking feature implementation

README.md

@@ -16,7 +16,7 @@ go get github.com/AlexEidt/Vidio
 
 The `Video` struct stores data about a video file you give it. The code below shows an example of sequentially reading the frames of the given video.
 
-Calling the `Read()` function will fill in the `Video` struct `framebuffer` with the next frame data as 8-bit RGBA data, stored in a flattened byte array in row-major order where each pixel is represented by four consecutive bytes representing the R, G, B and A components of that pixel. Note that the A (alpha) component will always be 255.
+Calling the `Read()` function will fill in the `Video` struct `framebuffer` with the next frame data as 8-bit RGBA data, stored in a flattened byte array in row-major order where each pixel is represented by four consecutive bytes representing the R, G, B and A components of that pixel. Note that the A (alpha) component will always be 255. When iteration over the entire video file is not required, we can lookup a specific frame by calling `ReadFrame(n int)`. By calling `ReadFrames(n ...int)`, we can immediately access multiple frames as a slice of RGBA images and skip the `framebuffer`.
 
 ```go
 vidio.NewVideo(filename string) (*vidio.Video, error)
@@ -38,6 +38,8 @@ MetaData() map[string]string
 SetFrameBuffer(buffer []byte) error
 
 Read() bool
+ReadFrame(n int) error
+ReadFrames(n ...int) ([]*image.RGBA, error)
 Close()
 ```
 
@@ -184,6 +186,35 @@ for video.Read() {
 }
 ```
 
+Write the last frame of `video.mp4` as `jpg` image (without iterating over all video frames).
+
+```go
+video, _ := video.NewVideo("video.mp4")
+
+img := image.NewRGBA(image.Rect(0, 0, video.Width(), video.Height()))
+video.SetFrameBuffer(img.Pix)
+
+video.ReadFrame(video.Frames() - 1)
+
+f, _ := os.Create(fmt.Sprintf("%d.jpg", video.Frames() - 1))
+jpeg.Encode(f, img, nil)
+f.Close()
+```
+
+Write the first and last frames of `video.mp4` as `jpg` images (without iterating over all video frames).
+
+```go
+video, _ := vidio.NewVideo("video.mp4")
+
+frames, _ := video.ReadFrames(0, video.Frames() - 1)
+
+for index, frame := range frames {
+	f, _ := os.Create(fmt.Sprintf("%d.jpg", index))
+	jpeg.Encode(f, frame, nil)
+	f.Close()
+}
+```
+
 # Acknowledgements
 
 * Special thanks to [Zulko](http://zulko.github.io/) and his [blog post](http://zulko.github.io/blog/2013/09/27/read-and-write-video-frames-in-python-using-ffmpeg/) about using FFmpeg to process video.

utils.go

@@ -214,3 +214,44 @@ func getDevicesWindows() ([]string, error) {
 	devices := parseDevices(builder.String())
 	return devices, nil
 }
+
+// Error representing a strings.Builder failure in the buildSelectExpression func.
+var errExpressionBuilder = fmt.Errorf("vidio: failed to write tokens to the frame select expresion")
+
+// Helper function used to generate a "-vf select" expression that specifies which video frames should be exported.
+func buildSelectExpression(n ...int) (string, error) {
+	sb := strings.Builder{}
+	if _, err := sb.WriteString("select='"); err != nil {
+		return "", errExpressionBuilder
+	}
+
+	for index, frame := range n {
+		if index != 0 {
+			if _, err := sb.WriteRune('+'); err != nil {
+				return "", errExpressionBuilder
+			}
+		}
+
+		if _, err := sb.WriteString("eq(n\\,"); err != nil {
+
+			return "", errExpressionBuilder
+		}
+
+		if _, err := sb.WriteString(strconv.Itoa(frame)); err != nil {
+
+			return "", errExpressionBuilder
+		}
+
+		if _, err := sb.WriteRune(')'); err != nil {
+
+			return "", errExpressionBuilder
+		}
+	}
+
+	if _, err := sb.WriteRune('\''); err != nil {
+
+		return "", errExpressionBuilder
+	}
+
+	return sb.String(), nil
+}

video.go

@@ -2,6 +2,7 @@ package vidio
 
 import (
 	"fmt"
+	"image"
 	"io"
 	"os"
 	"os/exec"
@@ -239,6 +240,145 @@ func (video *Video) Read() bool {
 	return true
 }
 
+// Reads the N-th frame from the video and stores it in the framebuffer. If the index is out of range or
+// the operation failes, the function will return an error. The frames are indexed from 0.
+func (video *Video) ReadFrame(n int) error {
+	if n >= video.frames {
+		return fmt.Errorf("vidio: provided frame index %d is not in frame count range", n)
+	}
+
+	if video.framebuffer == nil {
+		video.framebuffer = make([]byte, video.width*video.height*video.depth)
+	}
+
+	selectExpression, err := buildSelectExpression(n)
+	if err != nil {
+		return fmt.Errorf("vidio: failed to parse the specified frame index: %w", err)
+	}
+
+	cmd := exec.Command(
+		"ffmpeg",
+		"-i", video.filename,
+		"-f", "image2pipe",
+		"-loglevel", "quiet",
+		"-pix_fmt", "rgba",
+		"-vcodec", "rawvideo",
+		"-map", fmt.Sprintf("0:v:%d", video.stream),
+		"-vf", selectExpression,
+		"-vsync", "0",
+		"-",
+	)
+
+	stdoutPipe, err := cmd.StdoutPipe()
+	if err != nil {
+		return fmt.Errorf("vidio: failed to access the ffmpeg stdout pipe: %w", err)
+	}
+
+	if err := cmd.Start(); err != nil {
+		return fmt.Errorf("vidio: failed to start the ffmpeg cmd: %w", err)
+	}
+
+	interruptChan := make(chan os.Signal, 1)
+	signal.Notify(interruptChan, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		<-interruptChan
+		if stdoutPipe != nil {
+			stdoutPipe.Close()
+		}
+		if cmd != nil {
+			cmd.Process.Kill()
+		}
+		os.Exit(1)
+	}()
+
+	if _, err := io.ReadFull(stdoutPipe, video.framebuffer); err != nil {
+		return fmt.Errorf("vidio: failed to read the ffmpeg cmd result to the image buffer: %w", err)
+	}
+
+	if err := stdoutPipe.Close(); err != nil {
+		return fmt.Errorf("vidio: failed to close the ffmpeg stdout pipe: %w", err)
+	}
+
+	if err := cmd.Wait(); err != nil {
+		return fmt.Errorf("vidio: failed to free resources after the ffmpeg cmd: %w", err)
+	}
+
+	return nil
+}
+
+// Read the N-amount of frames with the given indexes and return them as a slice of RGBA image pointers. If one of
+// the indexes is out of range, the function will return an error. The frames are indexes from 0.
+func (video *Video) ReadFrames(n ...int) ([]*image.RGBA, error) {
+	if len(n) == 0 {
+		return nil, fmt.Errorf("vidio: no frames indexes specified")
+	}
+
+	for _, nValue := range n {
+		if nValue >= video.frames {
+			return nil, fmt.Errorf("vidio: provided frame index %d is not in frame count range", nValue)
+		}
+	}
+
+	selectExpression, err := buildSelectExpression(n...)
+	if err != nil {
+		return nil, fmt.Errorf("vidio: failed to parse the specified frame index: %w", err)
+	}
+
+	cmd := exec.Command(
+		"ffmpeg",
+		"-i", video.filename,
+		"-f", "image2pipe",
+		"-loglevel", "quiet",
+		"-pix_fmt", "rgba",
+		"-vcodec", "rawvideo",
+		"-map", fmt.Sprintf("0:v:%d", video.stream),
+		"-vf", selectExpression,
+		"-vsync", "0",
+		"-",
+	)
+
+	stdoutPipe, err := cmd.StdoutPipe()
+	if err != nil {
+		return nil, fmt.Errorf("vidio: failed to access the ffmpeg stdout pipe: %w", err)
+	}
+
+	if err := cmd.Start(); err != nil {
+		return nil, fmt.Errorf("vidio: failed to start the ffmpeg cmd: %w", err)
+	}
+
+	interruptChan := make(chan os.Signal, 1)
+	signal.Notify(interruptChan, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		<-interruptChan
+		if stdoutPipe != nil {
+			stdoutPipe.Close()
+		}
+		if cmd != nil {
+			cmd.Process.Kill()
+		}
+		os.Exit(1)
+	}()
+
+	frames := make([]*image.RGBA, len(n))
+	for frameIndex := range frames {
+		frames[frameIndex] = image.NewRGBA(image.Rect(0, 0, video.width, video.height))
+
+		if _, err := io.ReadFull(stdoutPipe, frames[frameIndex].Pix); err != nil {
+			return nil, fmt.Errorf("vidio: failed to read the ffmpeg cmd result to the image buffer: %w", err)
+		}
+	}
+
+	if err := stdoutPipe.Close(); err != nil {
+		return nil, fmt.Errorf("vidio: failed to close the ffmpeg stdout pipe: %w", err)
+	}
+
+	if err := cmd.Wait(); err != nil {
+		return nil, fmt.Errorf("vidio: failed to free resources after the ffmpeg cmd: %w", err)
+	}
+
+	return frames, nil
+}
+
 // Closes the pipe and stops the ffmpeg process.
 func (video *Video) Close() {
 	if video.pipe != nil {

vidio_test.go

@@ -2,6 +2,8 @@ package vidio
 
 import (
 	"fmt"
+	"image"
+	"image/png"
 	"os"
 	"testing"
 )
@@ -283,3 +285,142 @@ func TestImageWrite(t *testing.T) {
 
 	fmt.Println("Image Writing Test Passed")
 }
+
+func TestReadFrameShouldReturnErrorOnOutOfRangeFrame(t *testing.T) {
+	path := "test/koala.mp4"
+
+	video, err := NewVideo(path)
+	if err != nil {
+		t.Errorf("Failed to create the video: %s", err)
+	}
+
+	err = video.ReadFrame(video.Frames() + 1)
+	if err == nil {
+		t.Error("Error was expected to no be nil")
+	}
+}
+
+func TestReadFrameShouldReturnCorrectFrame(t *testing.T) {
+	path := "test/koala.mp4"
+
+	expectedFrameFile, err := os.Open("test/koala-frame5.png")
+	if err != nil {
+		t.Errorf("Failed to arrange the test: %s", err)
+	}
+
+	defer expectedFrameFile.Close()
+
+	expectedFrame, err := png.Decode(expectedFrameFile)
+	if err != nil {
+		t.Errorf("Failed to arrange the test: %s", err)
+	}
+
+	video, err := NewVideo(path)
+	if err != nil {
+		t.Errorf("Failed to create the video: %s", err)
+	}
+
+	actualFrame := image.NewRGBA(expectedFrame.Bounds())
+	if err := video.SetFrameBuffer(actualFrame.Pix); err != nil {
+		t.Errorf("Failed to set the frame buffer: %s", err)
+	}
+
+	if err := video.ReadFrame(5); err != nil {
+		t.Errorf("Failed to read the given frame: %s", err)
+	}
+
+	for xIndex := 0; xIndex < expectedFrame.Bounds().Dx(); xIndex += 1 {
+		for yIndex := 0; yIndex < expectedFrame.Bounds().Dy(); yIndex += 1 {
+			eR, eG, eB, eA := expectedFrame.At(xIndex, yIndex).RGBA()
+			aR, aG, aB, aA := actualFrame.At(xIndex, yIndex).RGBA()
+
+			if eR != aR || eG != aG || eB != aB || eA != aA {
+				t.Error("The expected and actual frames were expected to be equal")
+			}
+		}
+	}
+}
+
+func TestReadFramesShouldReturnErrorOnNoFramesSpecified(t *testing.T) {
+	path := "test/koala.mp4"
+
+	video, err := NewVideo(path)
+	if err != nil {
+		t.Errorf("Failed to create the video: %s", err)
+	}
+
+	_, err = video.ReadFrames()
+	if err == nil {
+		t.Error("Error was expected to no be nil")
+	}
+}
+
+func TestReadFramesShouldReturnErrorOnOutOfRangeFrame(t *testing.T) {
+	path := "test/koala.mp4"
+
+	video, err := NewVideo(path)
+	if err != nil {
+		t.Errorf("Failed to create the video: %s", err)
+	}
+
+	_, err = video.ReadFrames(0, video.Frames()-1, video.Frames()+1)
+	if err == nil {
+		t.Error("Error was expected to no be nil")
+	}
+}
+
+func TestReadFramesShouldReturnCorrectFrames(t *testing.T) {
+	path := "test/koala.mp4"
+
+	expectedFrames := make([]image.Image, 0, 2)
+
+	expectedFrameFile, err := os.Open("test/koala-frame5.png")
+	if err != nil {
+		t.Errorf("Failed to arrange the test: %s", err)
+	}
+
+	expectedFrame, err := png.Decode(expectedFrameFile)
+	if err != nil {
+		t.Errorf("Failed to arrange the test: %s", err)
+	}
+
+	expectedFrameFile.Close()
+	expectedFrames = append(expectedFrames, expectedFrame)
+
+	expectedFrameFile, err = os.Open("test/koala-frame15.png")
+	if err != nil {
+		t.Errorf("Failed to arrange the test: %s", err)
+	}
+
+	expectedFrame, err = png.Decode(expectedFrameFile)
+	if err != nil {
+		t.Errorf("Failed to arrange the test: %s", err)
+	}
+
+	expectedFrameFile.Close()
+	expectedFrames = append(expectedFrames, expectedFrame)
+
+	video, err := NewVideo(path)
+	if err != nil {
+		t.Errorf("Failed to create the video: %s", err)
+	}
+
+	frames, err := video.ReadFrames(5, 15)
+	if err != nil {
+		t.Errorf("Failed to read frames: %s", err)
+	}
+
+	for index, actualFrame := range frames {
+		expectedFrame := expectedFrames[index]
+		for xIndex := 0; xIndex < expectedFrame.Bounds().Dx(); xIndex += 1 {
+			for yIndex := 0; yIndex < expectedFrame.Bounds().Dy(); yIndex += 1 {
+				eR, eG, eB, eA := expectedFrame.At(xIndex, yIndex).RGBA()
+				aR, aG, aB, aA := actualFrame.At(xIndex, yIndex).RGBA()
+
+				if eR != aR || eG != aG || eB != aB || eA != aA {
+					t.Error("The expected and actual frames were expected to be equal")
+				}
+			}
+		}
+	}
+}