diff --git a/README.md b/README.md index ec55543..36a014d 100644 --- a/README.md +++ b/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. diff --git a/test/koala-frame15.png b/test/koala-frame15.png new file mode 100644 index 0000000..5ecd15b Binary files /dev/null and b/test/koala-frame15.png differ diff --git a/test/koala-frame5.png b/test/koala-frame5.png new file mode 100644 index 0000000..92c853e Binary files /dev/null and b/test/koala-frame5.png differ diff --git a/utils.go b/utils.go index c8423f2..90def84 100644 --- a/utils.go +++ b/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 +} diff --git a/video.go b/video.go index 5ec7de4..685cc26 100644 --- a/video.go +++ b/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 { diff --git a/vidio_test.go b/vidio_test.go index 03998d3..caf5567 100644 --- a/vidio_test.go +++ b/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") + } + } + } + } +}