From d5e98afcef65f95198c8ce2f2421acc725d0c136 Mon Sep 17 00:00:00 2001 From: Alexandre Boulanger <44292157+aboulang2002@users.noreply.github.com> Date: Mon, 30 Sep 2019 00:40:32 -0400 Subject: [PATCH] RL4J: Add VideoRecorder (#8106) * Added VideoRecorder Signed-off-by: unknown * Added missing header Signed-off-by: unknown * Changed HistoryProcessor to use VideoRecorder Signed-off-by: unknown --- .../rl4j/learning/HistoryProcessor.java | 42 +-- .../rl4j/util/VideoRecorder.java | 326 ++++++++++++++++++ 2 files changed, 344 insertions(+), 24 deletions(-) create mode 100644 rl4j/rl4j-core/src/main/java/org/deeplearning4j/rl4j/util/VideoRecorder.java diff --git a/rl4j/rl4j-core/src/main/java/org/deeplearning4j/rl4j/learning/HistoryProcessor.java b/rl4j/rl4j-core/src/main/java/org/deeplearning4j/rl4j/learning/HistoryProcessor.java index 7b1481cbe..243422b13 100644 --- a/rl4j/rl4j-core/src/main/java/org/deeplearning4j/rl4j/learning/HistoryProcessor.java +++ b/rl4j/rl4j-core/src/main/java/org/deeplearning4j/rl4j/learning/HistoryProcessor.java @@ -21,6 +21,7 @@ import lombok.extern.slf4j.Slf4j; import org.apache.commons.collections4.queue.CircularFifoQueue; import org.bytedeco.javacv.*; import org.datavec.image.loader.NativeImageLoader; +import org.deeplearning4j.rl4j.util.VideoRecorder; import org.nd4j.linalg.api.buffer.DataType; import org.nd4j.linalg.api.ndarray.INDArray; import org.nd4j.linalg.factory.Nd4j; @@ -46,7 +47,7 @@ public class HistoryProcessor implements IHistoryProcessor { final private Configuration conf; final private OpenCVFrameConverter openCVFrameConverter = new OpenCVFrameConverter.ToMat(); private CircularFifoQueue history; - private FFmpegFrameRecorder fmpegFrameRecorder = null; + private VideoRecorder videoRecorder; public HistoryProcessor(Configuration conf) { this.conf = conf; @@ -60,46 +61,39 @@ public class HistoryProcessor implements IHistoryProcessor { } public void startMonitor(String filename, int[] shape) { - stopMonitor(); - fmpegFrameRecorder = new FFmpegFrameRecorder(filename, shape[1], shape[0]); - fmpegFrameRecorder.setVideoCodec(AV_CODEC_ID_H264); - fmpegFrameRecorder.setFrameRate(30.0); - fmpegFrameRecorder.setVideoQuality(30); + if(videoRecorder == null) { + videoRecorder = VideoRecorder.builder(shape[0], shape[1]) + .frameInputType(VideoRecorder.FrameInputTypes.Float) + .build(); + } + try { - log.info("Started monitoring: " + filename); - fmpegFrameRecorder.start(); - } catch (FrameRecorder.Exception e) { + videoRecorder.startRecording(filename); + } catch (Exception e) { e.printStackTrace(); } } public void stopMonitor() { - if (fmpegFrameRecorder != null) { + if(videoRecorder != null) { try { - fmpegFrameRecorder.stop(); - fmpegFrameRecorder.release(); - log.info("Stopped monitoring"); - } catch (FrameRecorder.Exception e) { + videoRecorder.stopRecording(); + } catch (Exception e) { e.printStackTrace(); } } - fmpegFrameRecorder = null; } public boolean isMonitoring() { - return fmpegFrameRecorder != null; + return videoRecorder != null && videoRecorder.isRecording(); } public void record(INDArray raw) { - if (fmpegFrameRecorder != null) { - long[] shape = raw.shape(); - Mat ocvmat = new Mat((int)shape[0], (int)shape[1], CV_32FC(3), raw.data().pointer()); - Mat cvmat = new Mat(shape[0], shape[1], CV_8UC(3)); - ocvmat.convertTo(cvmat, CV_8UC(3), 255.0, 0.0); - Frame frame = openCVFrameConverter.convert(cvmat); + if(isMonitoring()) { + VideoRecorder.VideoFrame frame = videoRecorder.createFrame(raw.data().pointer()); try { - fmpegFrameRecorder.record(frame); - } catch (FrameRecorder.Exception e) { + videoRecorder.record(frame); + } catch (Exception e) { e.printStackTrace(); } } diff --git a/rl4j/rl4j-core/src/main/java/org/deeplearning4j/rl4j/util/VideoRecorder.java b/rl4j/rl4j-core/src/main/java/org/deeplearning4j/rl4j/util/VideoRecorder.java new file mode 100644 index 000000000..5a91b71e4 --- /dev/null +++ b/rl4j/rl4j-core/src/main/java/org/deeplearning4j/rl4j/util/VideoRecorder.java @@ -0,0 +1,326 @@ +/******************************************************************************* + * Copyright (c) 2015-2018 Skymind, Inc. + * + * This program and the accompanying materials are made available under the + * terms of the Apache License, Version 2.0 which is available at + * https://www.apache.org/licenses/LICENSE-2.0. + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations + * under the License. + * + * SPDX-License-Identifier: Apache-2.0 + ******************************************************************************/ + +package org.deeplearning4j.rl4j.util; + +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; +import org.bytedeco.javacpp.BytePointer; +import org.bytedeco.javacpp.Pointer; +import org.bytedeco.javacv.FFmpegFrameRecorder; +import org.bytedeco.javacv.Frame; +import org.bytedeco.javacv.OpenCVFrameConverter; +import org.bytedeco.opencv.global.opencv_core; +import org.bytedeco.opencv.global.opencv_imgproc; +import org.bytedeco.opencv.opencv_core.Mat; +import org.bytedeco.opencv.opencv_core.Rect; +import org.bytedeco.opencv.opencv_core.Size; +import org.opencv.imgproc.Imgproc; + +import static org.bytedeco.ffmpeg.global.avcodec.*; +import static org.bytedeco.opencv.global.opencv_core.*; + +/** + * VideoRecorder is used to create a video from a sequence of individual frames. If using 3 channels + * images, it expects B-G-R order. A RGB order can be used by calling isRGBOrder(true).
+ * Example:
+ * 
+ * {@code
+ *        VideoRecorder recorder = VideoRecorder.builder(160, 100)
+ *             .numChannels(3)
+ *             .isRGBOrder(true)
+ *             .build();
+ *         recorder.startRecording("myVideo.mp4");
+ *         while(...) {
+ *             byte[] data = new byte[160*100*3];
+ *             // Todo: Fill data
+ *             VideoRecorder.VideoFrame frame = recorder.createFrame(data);
+ *             // Todo: Apply cropping or resizing to frame
+ *             recorder.record(frame);
+ *         }
+ *         recorder.stopRecording();
+ * }
+ *
+ * + * @author Alexandre Boulanger + */ +@Slf4j +public class VideoRecorder implements AutoCloseable { + + public enum FrameInputTypes { BGR, RGB, Float } + + private final int height; + private final int width; + private final int imageType; + private final OpenCVFrameConverter openCVFrameConverter = new OpenCVFrameConverter.ToMat(); + private final int codec; + private final double framerate; + private final int videoQuality; + private final FrameInputTypes frameInputType; + + private FFmpegFrameRecorder fmpegFrameRecorder = null; + + /** + * @return True if the instance is recording a video + */ + public boolean isRecording() { + return fmpegFrameRecorder != null; + } + + private VideoRecorder(Builder builder) { + this.height = builder.height; + this.width = builder.width; + imageType = CV_8UC(builder.numChannels); + codec = builder.codec; + framerate = builder.frameRate; + videoQuality = builder.videoQuality; + frameInputType = builder.frameInputType; + } + + /** + * Initiate the recording of a video + * @param filename Name of the video file to create + * @throws Exception + */ + public void startRecording(String filename) throws Exception { + stopRecording(); + + fmpegFrameRecorder = new FFmpegFrameRecorder(filename, width, height); + fmpegFrameRecorder.setVideoCodec(codec); + fmpegFrameRecorder.setFrameRate(framerate); + fmpegFrameRecorder.setVideoQuality(videoQuality); + fmpegFrameRecorder.start(); + } + + /** + * Terminates the recording of the video + * @throws Exception + */ + public void stopRecording() throws Exception { + if (fmpegFrameRecorder != null) { + fmpegFrameRecorder.stop(); + fmpegFrameRecorder.release(); + } + fmpegFrameRecorder = null; + } + + /** + * Add a frame to the video + * @param frame the VideoFrame to add to the video + * @throws Exception + */ + public void record(VideoFrame frame) throws Exception { + Size size = frame.getMat().size(); + if(size.height() != height || size.width() != width) { + throw new IllegalArgumentException(String.format("Wrong frame size. Got (%dh x %dw) expected (%dh x %dw)", size.height(), size.width(), height, width)); + } + Frame cvFrame = openCVFrameConverter.convert(frame.getMat()); + fmpegFrameRecorder.record(cvFrame); + } + + /** + * Create a VideoFrame from a byte array. + * @param data A byte array. Expect the index to be of the form [(Y*Width + X) * NumChannels + channel] + * @return An instance of VideoFrame + */ + public VideoFrame createFrame(byte[] data) { + return createFrame(new BytePointer(data)); + } + + /** + * Create a VideoFrame from a byte array with different height and width than the video + * the frame will need to be cropped or resized before being added to the video) + * + * @param data A byte array Expect the index to be of the form [(Y*customWidth + X) * NumChannels + channel] + * @param customHeight The actual height of the data + * @param customWidth The actual width of the data + * @return A VideoFrame instance + */ + public VideoFrame createFrame(byte[] data, int customHeight, int customWidth) { + return createFrame(new BytePointer(data), customHeight, customWidth); + } + + /** + * Create a VideoFrame from a Pointer (to use for example with a INDarray). + * @param data A Pointer (for example myINDArray.data().pointer()) + * @return An instance of VideoFrame + */ + public VideoFrame createFrame(Pointer data) { + return new VideoFrame(height, width, imageType, frameInputType, data); + } + + /** + * Create a VideoFrame from a Pointer with different height and width than the video + * the frame will need to be cropped or resized before being added to the video) + * @param data + * @param customHeight The actual height of the data + * @param customWidth The actual width of the data + * @return A VideoFrame instance + */ + public VideoFrame createFrame(Pointer data, int customHeight, int customWidth) { + return new VideoFrame(customHeight, customWidth, imageType, frameInputType, data); + } + + /** + * Terminate the recording and close the video file + * @throws Exception + */ + public void close() throws Exception { + stopRecording(); + } + + /** + * + * @param height The height of the video + * @param width Thw width of the video + * @return A VideoRecorder builder + */ + public static Builder builder(int height, int width) { + return new Builder(height, width); + } + + /** + * An individual frame for the video + */ + public static class VideoFrame { + + private final int height; + private final int width; + private final int imageType; + @Getter + private Mat mat; + + private VideoFrame(int height, int width, int imageType, FrameInputTypes frameInputType, Pointer data) { + this.height = height; + this.width = width; + this.imageType = imageType; + + switch(frameInputType) { + case RGB: + Mat src = new Mat(height, width, imageType, data); + mat = new Mat(height, width, imageType); + opencv_imgproc.cvtColor(src, mat, Imgproc.COLOR_RGB2BGR); + break; + + case BGR: + mat = new Mat(height, width, imageType, data); + break; + + case Float: + Mat tmpMat = new Mat(height, width, CV_32FC(3), data); + mat = new Mat(height, width, imageType); + tmpMat.convertTo(mat, CV_8UC(3), 255.0, 0.0); + } + } + + /** + * Crop the video to a specified size + * @param newHeight The new height of the frame + * @param newWidth The new width of the frame + * @param heightOffset The starting height offset in the uncropped frame + * @param widthOffset The starting weight offset in the uncropped frame + */ + public void crop(int newHeight, int newWidth, int heightOffset, int widthOffset) { + mat = mat.apply(new Rect(widthOffset, heightOffset, newWidth, newHeight)); + } + + /** + * Resize the frame to a specified size + * @param newHeight The new height of the frame + * @param newWidth The new width of the frame + */ + public void resize(int newHeight, int newWidth) { + mat = new Mat(newHeight, newWidth, imageType); + } + } + + /** + * A builder class for the VideoRecorder + */ + public static class Builder { + private final int height; + private final int width; + private int numChannels = 3; + private FrameInputTypes frameInputType = FrameInputTypes.BGR; + private int codec = AV_CODEC_ID_H264; + private double frameRate = 30.0; + private int videoQuality = 30; + + /** + * @param height The height of the video + * @param width The width of the video + */ + public Builder(int height, int width) { + this.height = height; + this.width = width; + } + + /** + * Specify the number of channels. Default is 3 + * @param numChannels + */ + public Builder numChannels(int numChannels) { + this.numChannels = numChannels; + return this; + } + + /** + * Tell the VideoRecorder what data it will receive (default is BGR) + * @param frameInputType (See {@link FrameInputTypes}} + */ + public Builder frameInputType(FrameInputTypes frameInputType) { + this.frameInputType = frameInputType; + return this; + } + + /** + * The codec to use for the video. Default is AV_CODEC_ID_H264 + * @param codec Code (see {@link org.bytedeco.ffmpeg.global.avcodec codec codes}) + */ + public Builder codec(int codec) { + this.codec = codec; + return this; + } + + /** + * The frame rate of the video. Default is 30.0 + * @param frameRate The frame rate + * @return + */ + public Builder frameRate(double frameRate) { + this.frameRate = frameRate; + return this; + } + + /** + * The video quality. Default is 30 + * @param videoQuality + * @return + */ + public Builder videoQuality(int videoQuality) { + this.videoQuality = videoQuality; + return this; + } + + /** + * Build an instance of the configured VideoRecorder + * @return A VideoRecorder instance + */ + public VideoRecorder build() { + return new VideoRecorder(this); + } + } +}