RL4J: Add VideoRecorder (#8106)
* Added VideoRecorder Signed-off-by: unknown <aboulang2002@yahoo.com> * Added missing header Signed-off-by: unknown <aboulang2002@yahoo.com> * Changed HistoryProcessor to use VideoRecorder Signed-off-by: unknown <aboulang2002@yahoo.com>master
Alexandre Boulanger
Samuel Audet
parent
f98f8be7b6
commit
d5e98afcef
|
|
@ -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<INDArray> 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();
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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).<br>
|
||||
* Example:<br>
|
||||
* <pre>
|
||||
* {@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();
|
||||
* }
|
||||
* </pre>
|
||||
*
|
||||
* @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);
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue