1025 lines
34 KiB
C++
1025 lines
34 KiB
C++
#ifndef CAFFE2_VIDEO_VIDEO_INPUT_OP_H_
|
|
#define CAFFE2_VIDEO_VIDEO_INPUT_OP_H_
|
|
|
|
#include <exception>
|
|
#include <istream>
|
|
#include <ostream>
|
|
#include <random>
|
|
#include <string>
|
|
|
|
#include <c10/core/thread_pool.h>
|
|
#include <c10/util/irange.h>
|
|
#include <caffe2/core/db.h>
|
|
#include <caffe2/core/logging.h>
|
|
#include <caffe2/operators/prefetch_op.h>
|
|
#include <caffe2/utils/math.h>
|
|
#include <caffe2/video/video_decoder.h>
|
|
#include <caffe2/video/video_io.h>
|
|
|
|
namespace caffe2 {
|
|
|
|
template <class Context>
|
|
class VideoInputOp final : public PrefetchOperator<Context> {
|
|
public:
|
|
using OperatorBase::OutputSize;
|
|
using PrefetchOperator<Context>::context_;
|
|
using PrefetchOperator<Context>::prefetch_thread_;
|
|
explicit VideoInputOp(const OperatorDef& operator_def, Workspace* ws);
|
|
~VideoInputOp() {
|
|
PrefetchOperator<Context>::Finalize();
|
|
}
|
|
|
|
// override methods
|
|
bool Prefetch() override;
|
|
bool CopyPrefetched() override;
|
|
|
|
private:
|
|
void CheckParamsAndPrint();
|
|
|
|
bool GetClipsAndLabelsFromDBValue(
|
|
const std::string& value,
|
|
int& height,
|
|
int& width,
|
|
std::vector<unsigned char*>& buffer_rgb,
|
|
int* label_data,
|
|
int64_t* video_id_data,
|
|
int* start_frame_data,
|
|
std::mt19937* randgen);
|
|
|
|
void DecodeAndTransform(
|
|
const std::string& value,
|
|
float* clip_rgb_data,
|
|
float* clip_of_data,
|
|
int* label_data,
|
|
int64_t* video_id_data,
|
|
int* start_frame_data,
|
|
std::mt19937* randgen,
|
|
std::bernoulli_distribution* mirror_this_clip);
|
|
|
|
void GetLabelsFromProto(const TensorProto& label_proto, int* label_data);
|
|
|
|
bool GetImageAndLabelsFromDBValue(
|
|
const std::string& value,
|
|
int& height,
|
|
int& width,
|
|
std::vector<unsigned char*>& buffer_rgb,
|
|
int* label_data);
|
|
|
|
const db::DBReader* reader_;
|
|
CPUContext cpu_context_;
|
|
Tensor prefetched_clip_rgb_;
|
|
Tensor prefetched_clip_of_;
|
|
Tensor prefetched_label_;
|
|
Tensor prefetched_video_id_;
|
|
Tensor prefetched_start_frame_;
|
|
Tensor prefetched_clip_rgb_on_device_{Context::GetDeviceType()};
|
|
Tensor prefetched_clip_of_on_device_{Context::GetDeviceType()};
|
|
Tensor prefetched_label_on_device_{Context::GetDeviceType()};
|
|
Tensor prefetched_video_id_on_device_{Context::GetDeviceType()};
|
|
Tensor prefetched_start_frame_on_device_{Context::GetDeviceType()};
|
|
|
|
int batch_size_;
|
|
int clip_per_video_;
|
|
std::vector<int> clip_start_positions_;
|
|
std::vector<float> mean_rgb_;
|
|
std::vector<float> inv_std_rgb_;
|
|
std::vector<float> mean_of_;
|
|
std::vector<float> inv_std_of_;
|
|
int channels_rgb_;
|
|
int channels_of_;
|
|
int crop_size_;
|
|
int scale_h_;
|
|
int scale_w_;
|
|
int short_edge_;
|
|
std::vector<int> jitter_scales_;
|
|
int length_rgb_;
|
|
int sampling_rate_rgb_;
|
|
int random_sampling_rate_;
|
|
int num_of_required_frame_;
|
|
int length_of_;
|
|
int sampling_rate_of_;
|
|
int frame_gap_of_;
|
|
bool random_mirror_;
|
|
int num_of_class_;
|
|
bool use_local_file_;
|
|
bool random_crop_;
|
|
int crop_per_clip_;
|
|
int flow_data_type_;
|
|
int flow_alg_type_;
|
|
int decode_type_;
|
|
int video_res_type_;
|
|
bool do_flow_aggregation_;
|
|
bool image_as_input_;
|
|
bool get_rgb_;
|
|
bool get_optical_flow_;
|
|
bool get_video_id_;
|
|
bool get_start_frame_;
|
|
bool do_multi_label_;
|
|
|
|
// thread pool for parse + decode
|
|
int num_decode_threads_;
|
|
std::shared_ptr<TaskThreadPool> thread_pool_;
|
|
};
|
|
|
|
template <class Context>
|
|
void VideoInputOp<Context>::CheckParamsAndPrint() {
|
|
// check whether the input parameters are valid or not
|
|
CAFFE_ENFORCE_GT(batch_size_, 0, "Batch size should be positive.");
|
|
CAFFE_ENFORCE_GT(
|
|
clip_per_video_, 0, "Number of clips per video should be positive.");
|
|
CAFFE_ENFORCE_GT(crop_size_, 0, "Must provide the cropping value.");
|
|
|
|
if (!image_as_input_) {
|
|
CAFFE_ENFORCE_GT(
|
|
num_of_required_frame_,
|
|
0,
|
|
"Required number of frames must be positive.");
|
|
}
|
|
|
|
if (image_as_input_) {
|
|
CAFFE_ENFORCE_EQ(
|
|
video_res_type_,
|
|
VideoResType::USE_WIDTH_HEIGHT,
|
|
"Currently only USE_WIDTH_HEIGHT option is supported with images");
|
|
}
|
|
|
|
if (video_res_type_ == VideoResType::USE_SHORT_EDGE) {
|
|
CAFFE_ENFORCE_GT(short_edge_, 0, "Must provide the short edge value.");
|
|
CAFFE_ENFORCE_GE(
|
|
short_edge_,
|
|
crop_size_,
|
|
"The short edge must be no smaller than the crop value.");
|
|
} else if (video_res_type_ == VideoResType::USE_WIDTH_HEIGHT) {
|
|
CAFFE_ENFORCE_GT(scale_h_, 0, "Must provide the scale height value.");
|
|
CAFFE_ENFORCE_GT(scale_w_, 0, "Must provide the scale width value.");
|
|
CAFFE_ENFORCE_GE(
|
|
scale_h_,
|
|
crop_size_,
|
|
"The scaled height must be no smaller than the crop value.");
|
|
CAFFE_ENFORCE_GE(
|
|
scale_w_,
|
|
crop_size_,
|
|
"The scaled width must be no smaller than the crop value.");
|
|
}
|
|
|
|
if (jitter_scales_.size() > 0) {
|
|
CAFFE_ENFORCE_GE(
|
|
video_res_type_,
|
|
VideoResType::USE_SHORT_EDGE,
|
|
"Scale jittering is used with short_edge scaling only");
|
|
}
|
|
|
|
if (get_rgb_) {
|
|
CAFFE_ENFORCE_GT(length_rgb_, 0, "Must provide rgb clip length.");
|
|
CAFFE_ENFORCE_GT(
|
|
sampling_rate_rgb_, 0, "4 frames for mc2; 2 frames for res3d.");
|
|
CAFFE_ENFORCE_EQ(
|
|
channels_rgb_, mean_rgb_.size(), "Number rgb channels is wrong!");
|
|
CAFFE_ENFORCE_EQ(
|
|
channels_rgb_, inv_std_rgb_.size(), "Number rgb channels is wrong!");
|
|
}
|
|
|
|
if (get_optical_flow_) {
|
|
CAFFE_ENFORCE_GT(length_of_, 0, "Must provide optical flow clip length.");
|
|
CAFFE_ENFORCE_GT(
|
|
sampling_rate_of_, 0, "4 frames for mc2; 2 frames for res3d.");
|
|
CAFFE_ENFORCE_EQ(
|
|
channels_of_,
|
|
mean_of_.size(),
|
|
"Number of optical flow channels is wrong!");
|
|
CAFFE_ENFORCE_EQ(
|
|
channels_of_,
|
|
inv_std_of_.size(),
|
|
"Number of optical flow channels is wrong!");
|
|
}
|
|
|
|
if (clip_per_video_ > 1) {
|
|
CAFFE_ENFORCE_EQ(
|
|
decode_type_,
|
|
DecodeType::DO_UNIFORM_SMP,
|
|
"Only uniformly sampling is supported when sampling multiple clips!");
|
|
}
|
|
|
|
if (do_multi_label_) {
|
|
CAFFE_ENFORCE_GT(
|
|
num_of_class_,
|
|
0,
|
|
"Number of classes must be set when using multiple labels.");
|
|
}
|
|
|
|
// print out the parameter settings
|
|
LOG(INFO) << "Creating a clip input op with the following setting: ";
|
|
LOG(INFO) << " Input Type: " << (image_as_input_ ? "Image" : "Video");
|
|
LOG(INFO) << " Using " << num_decode_threads_ << " CPU threads;";
|
|
LOG(INFO) << " Outputting in batches of " << batch_size_ << " videos;";
|
|
LOG(INFO) << " Each video has " << clip_per_video_ << " clips;";
|
|
LOG(INFO) << " Scaling image to " << scale_h_ << "x" << scale_w_;
|
|
LOG(INFO) << " Cropping video frame to " << crop_size_
|
|
<< (random_mirror_ ? " with " : " without ") << "random mirroring;";
|
|
LOG(INFO) << " Using " << (random_crop_ ? "random" : "center") << " crop";
|
|
LOG(INFO) << " Using " << crop_per_clip_ << " spatial crop(s)";
|
|
|
|
if (get_rgb_) {
|
|
LOG(INFO) << " Using a clip of " << length_rgb_ << " rgb frames "
|
|
<< "with " << channels_rgb_ << " channels "
|
|
<< "and a sampling rate of 1:" << sampling_rate_rgb_;
|
|
if (random_sampling_rate_) {
|
|
LOG(INFO) << "random sampling with max:" << random_sampling_rate_;
|
|
}
|
|
for (const auto i : c10::irange(channels_rgb_)) {
|
|
LOG(INFO) << " RGB " << i << "-th channel mean: " << mean_rgb_[i]
|
|
<< " std: " << 1.f / inv_std_rgb_[i];
|
|
}
|
|
}
|
|
|
|
if (get_optical_flow_) {
|
|
LOG(INFO) << " Using a clip of " << length_of_ << " optical flow frames "
|
|
<< "with " << channels_of_ << " channels "
|
|
<< "and a sampling rate of 1:" << sampling_rate_of_
|
|
<< " flow_data_type_: " << flow_data_type_
|
|
<< " flow_alg_type_: " << flow_alg_type_;
|
|
for (const auto i : c10::irange(channels_of_)) {
|
|
LOG(INFO) << " Optical flow" << i
|
|
<< "-th channel mean: " << mean_of_[i]
|
|
<< " std: " << 1.f / inv_std_of_[i];
|
|
}
|
|
}
|
|
|
|
if (video_res_type_ == VideoResType::ORIGINAL_RES) {
|
|
LOG(INFO) << " Use original resolution";
|
|
} else if (video_res_type_ == VideoResType::USE_SHORT_EDGE) {
|
|
LOG(INFO) << " Resize and keep aspect ratio";
|
|
} else if (video_res_type_ == VideoResType::USE_WIDTH_HEIGHT) {
|
|
LOG(INFO) << " Resize and ignore aspect ratio";
|
|
} else {
|
|
LOG(ERROR) << " Unknown video resolution type";
|
|
}
|
|
|
|
if (video_res_type_ == VideoResType::USE_SHORT_EDGE) {
|
|
if (jitter_scales_.size() > 0) {
|
|
LOG(INFO) << "Using scale jittering:";
|
|
for (const auto idx : c10::irange(jitter_scales_.size())) {
|
|
LOG(INFO) << "scale " << idx << ": " << jitter_scales_[idx];
|
|
}
|
|
} else {
|
|
LOG(INFO) << "No scale jittering is used.";
|
|
}
|
|
}
|
|
|
|
if (decode_type_ == DecodeType::DO_TMP_JITTER) {
|
|
LOG(INFO) << " Do temporal jittering";
|
|
} else if (decode_type_ == DecodeType::USE_START_FRM) {
|
|
LOG(INFO) << " Use start_frm for decoding";
|
|
} else if (decode_type_ == DecodeType::DO_UNIFORM_SMP) {
|
|
LOG(INFO) << " Do uniformly sampling";
|
|
} else {
|
|
LOG(ERROR) << " Unknown video decoding type";
|
|
}
|
|
if (get_start_frame_) {
|
|
CAFFE_ENFORCE_EQ(
|
|
decode_type_,
|
|
DecodeType::USE_START_FRM,
|
|
"Only decoding with starting frame is supported w/ get start_frame!");
|
|
CAFFE_ENFORCE_EQ(
|
|
clip_per_video_, 1, "get start frame support only clip per video = 1");
|
|
}
|
|
}
|
|
|
|
template <class Context>
|
|
VideoInputOp<Context>::VideoInputOp(
|
|
const OperatorDef& operator_def,
|
|
Workspace* ws)
|
|
: PrefetchOperator<Context>(operator_def, ws),
|
|
reader_(nullptr),
|
|
batch_size_(
|
|
OperatorBase::template GetSingleArgument<int>("batch_size", 0)),
|
|
clip_per_video_(
|
|
OperatorBase::template GetSingleArgument<int>("clip_per_video", 1)),
|
|
clip_start_positions_(OperatorBase::template GetRepeatedArgument<int>(
|
|
"clip_start_positions",
|
|
{})),
|
|
channels_rgb_(
|
|
OperatorBase::template GetSingleArgument<int>("channels_rgb", 3)),
|
|
channels_of_(
|
|
OperatorBase::template GetSingleArgument<int>("channels_of", 2)),
|
|
crop_size_(OperatorBase::template GetSingleArgument<int>("crop_size", 0)),
|
|
scale_h_(OperatorBase::template GetSingleArgument<int>("scale_h", 0)),
|
|
scale_w_(OperatorBase::template GetSingleArgument<int>("scale_w", 0)),
|
|
short_edge_(
|
|
OperatorBase::template GetSingleArgument<int>("short_edge", 0)),
|
|
jitter_scales_(
|
|
OperatorBase::template GetRepeatedArgument<int>("jitter_scales", {})),
|
|
length_rgb_(
|
|
OperatorBase::template GetSingleArgument<int>("length_rgb", 0)),
|
|
sampling_rate_rgb_(OperatorBase::template GetSingleArgument<int>(
|
|
"sampling_rate_rgb",
|
|
1)),
|
|
random_sampling_rate_(OperatorBase::template GetSingleArgument<int>(
|
|
"random_sampling_rate",
|
|
0)),
|
|
length_of_(OperatorBase::template GetSingleArgument<int>("length_of", 0)),
|
|
sampling_rate_of_(
|
|
OperatorBase::template GetSingleArgument<int>("sampling_rate_of", 1)),
|
|
frame_gap_of_(
|
|
OperatorBase::template GetSingleArgument<int>("frame_gap_of", 1)),
|
|
random_mirror_(OperatorBase::template GetSingleArgument<bool>(
|
|
"random_mirror",
|
|
true)),
|
|
num_of_class_(
|
|
OperatorBase::template GetSingleArgument<int>("num_of_class", 0)),
|
|
use_local_file_(OperatorBase::template GetSingleArgument<bool>(
|
|
"use_local_file",
|
|
false)),
|
|
random_crop_(
|
|
OperatorBase::template GetSingleArgument<bool>("random_crop", true)),
|
|
crop_per_clip_(
|
|
OperatorBase::template GetSingleArgument<int>("crop_per_clip", 1)),
|
|
flow_data_type_(
|
|
OperatorBase::template GetSingleArgument<int>("flow_data_type", 0)),
|
|
flow_alg_type_(
|
|
OperatorBase::template GetSingleArgument<int>("flow_alg_type", 0)),
|
|
decode_type_(
|
|
OperatorBase::template GetSingleArgument<int>("decode_type", 0)),
|
|
video_res_type_(
|
|
OperatorBase::template GetSingleArgument<int>("video_res_type", 0)),
|
|
do_flow_aggregation_(OperatorBase::template GetSingleArgument<bool>(
|
|
"do_flow_aggregation",
|
|
true)),
|
|
image_as_input_(OperatorBase::template GetSingleArgument<bool>(
|
|
"image_as_input",
|
|
false)),
|
|
get_rgb_(OperatorBase::template GetSingleArgument<bool>("get_rgb", true)),
|
|
get_optical_flow_(OperatorBase::template GetSingleArgument<bool>(
|
|
"get_optical_flow",
|
|
false)),
|
|
get_video_id_(OperatorBase::template GetSingleArgument<bool>(
|
|
"get_video_id",
|
|
false)),
|
|
get_start_frame_(OperatorBase::template GetSingleArgument<bool>(
|
|
"get_start_frame",
|
|
false)),
|
|
do_multi_label_(OperatorBase::template GetSingleArgument<bool>(
|
|
"do_multi_label",
|
|
false)),
|
|
num_decode_threads_(OperatorBase::template GetSingleArgument<int>(
|
|
"num_decode_threads",
|
|
4)),
|
|
thread_pool_(std::make_shared<TaskThreadPool>(num_decode_threads_)) {
|
|
try {
|
|
num_of_required_frame_ = 0;
|
|
|
|
// mean and std for normalizing different optical flow data type;
|
|
// Example statistics generated from SOA are shown below, and you may
|
|
// want to change them if you are running on a different dataset;
|
|
|
|
// 7 channels: (flow_x, flow_y, flow_magitude, gray, Red, Green, Blue)
|
|
const std::vector<float> InputDataMean = {
|
|
0.0046635, 0.0046261, 0.963986, 102.976, 110.201, 100.64, 95.9966};
|
|
const std::vector<float> InputDataStd = {
|
|
0.972347, 0.755146, 1.43588, 55.3691, 58.1489, 56.4701, 55.3324};
|
|
|
|
// if we need RGB as an input
|
|
if (get_rgb_ && !image_as_input_) {
|
|
// how many frames we need for RGB
|
|
num_of_required_frame_ = std::max(
|
|
num_of_required_frame_, (length_rgb_ - 1) * sampling_rate_rgb_ + 1);
|
|
|
|
if (random_sampling_rate_) {
|
|
num_of_required_frame_ = std::max(
|
|
num_of_required_frame_,
|
|
(length_rgb_ - 1) * random_sampling_rate_ + 1);
|
|
}
|
|
|
|
channels_rgb_ = 3;
|
|
for (const auto i : c10::irange(4, 7)) {
|
|
mean_rgb_.push_back(InputDataMean[i]);
|
|
inv_std_rgb_.push_back(1.f / InputDataStd[i]);
|
|
}
|
|
}
|
|
|
|
if (image_as_input_) {
|
|
channels_rgb_ = 3;
|
|
length_rgb_ = 1;
|
|
clip_per_video_ = 1;
|
|
get_optical_flow_ = false;
|
|
get_rgb_ = true;
|
|
sampling_rate_rgb_ = 1;
|
|
for (const auto i : c10::irange(4, 7)) {
|
|
mean_rgb_.push_back(InputDataMean[i]);
|
|
inv_std_rgb_.push_back(1.f / InputDataStd[i]);
|
|
}
|
|
}
|
|
|
|
// if we need optical flow as an input
|
|
if (get_optical_flow_) {
|
|
// how many frames we need for optical flow
|
|
num_of_required_frame_ = std::max(
|
|
num_of_required_frame_,
|
|
(length_of_ - 1) * sampling_rate_of_ + frame_gap_of_ + 1);
|
|
|
|
// set the parameters for different input data types
|
|
switch (flow_data_type_) {
|
|
case FlowDataType::Flow2C:
|
|
channels_of_ = 2;
|
|
for (const auto i : c10::irange(channels_of_)) {
|
|
mean_of_.push_back(InputDataMean[i]);
|
|
inv_std_of_.push_back(1.f / InputDataStd[i]);
|
|
}
|
|
break;
|
|
|
|
case FlowDataType::Flow3C:
|
|
channels_of_ = 3;
|
|
for (const auto i : c10::irange(channels_of_)) {
|
|
mean_of_.push_back(InputDataMean[i]);
|
|
inv_std_of_.push_back(1.f / InputDataStd[i]);
|
|
}
|
|
break;
|
|
|
|
// early fusion with gray
|
|
case FlowDataType::FlowWithGray:
|
|
channels_of_ = 3;
|
|
for (const auto i : c10::irange(2)) {
|
|
mean_of_.push_back(InputDataMean[i]);
|
|
inv_std_of_.push_back(1.f / InputDataStd[i]);
|
|
}
|
|
mean_of_.push_back(InputDataMean[3]);
|
|
inv_std_of_.push_back(1.f / InputDataStd[3]);
|
|
break;
|
|
|
|
// early fusion with RGB
|
|
case FlowDataType::FlowWithRGB:
|
|
channels_of_ = 5;
|
|
for (const auto i : c10::irange(2)) {
|
|
mean_of_.push_back(InputDataMean[i]);
|
|
inv_std_of_.push_back(1.f / InputDataStd[i]);
|
|
}
|
|
for (const auto i : c10::irange(4, 7)) {
|
|
mean_of_.push_back(InputDataMean[i]);
|
|
inv_std_of_.push_back(1.f / InputDataStd[i]);
|
|
}
|
|
break;
|
|
|
|
default:
|
|
LOG(ERROR) << "Unknown optical flow type " << flow_data_type_;
|
|
break;
|
|
}
|
|
}
|
|
|
|
CheckParamsAndPrint();
|
|
// Always need a dbreader, even when using local video files
|
|
CAFFE_ENFORCE_GT(
|
|
operator_def.input_size(), 0, "Need to have a DBReader blob input");
|
|
|
|
vector<int64_t> data_shape(5);
|
|
vector<int64_t> label_shape(2);
|
|
|
|
// In case clip_start_positions are given, set the clip_per_video arg
|
|
if (clip_start_positions_.size() > 0) {
|
|
clip_per_video_ = clip_start_positions_.size();
|
|
}
|
|
|
|
// for RGB data
|
|
data_shape[0] = batch_size_ * clip_per_video_ * crop_per_clip_;
|
|
data_shape[1] = channels_rgb_;
|
|
data_shape[2] = length_rgb_;
|
|
data_shape[3] = crop_size_;
|
|
data_shape[4] = crop_size_;
|
|
ReinitializeTensor(
|
|
&prefetched_clip_rgb_, data_shape, at::dtype<float>().device(CPU));
|
|
|
|
// for optical flow data
|
|
data_shape[1] = channels_of_;
|
|
data_shape[2] = length_of_;
|
|
ReinitializeTensor(
|
|
&prefetched_clip_of_, data_shape, at::dtype<float>().device(CPU));
|
|
|
|
// If do_multi_label is used, output label is a binary vector
|
|
// of length num_of_class indicating which labels present
|
|
if (do_multi_label_) {
|
|
label_shape[0] = batch_size_ * clip_per_video_ * crop_per_clip_;
|
|
label_shape[1] = num_of_class_;
|
|
ReinitializeTensor(
|
|
&prefetched_label_, label_shape, at::dtype<int>().device(CPU));
|
|
} else {
|
|
ReinitializeTensor(
|
|
&prefetched_label_,
|
|
vector<int64_t>(1, batch_size_ * clip_per_video_ * crop_per_clip_),
|
|
at::dtype<int>().device(CPU));
|
|
}
|
|
|
|
ReinitializeTensor(
|
|
&prefetched_video_id_,
|
|
vector<int64_t>(1, batch_size_ * clip_per_video_ * crop_per_clip_),
|
|
at::dtype<int>().device(CPU));
|
|
ReinitializeTensor(
|
|
&prefetched_start_frame_,
|
|
vector<int64_t>(1, batch_size_ * clip_per_video_ * crop_per_clip_),
|
|
at::dtype<int>().device(CPU));
|
|
|
|
} catch (const std::exception& exc) {
|
|
std::cerr << "While calling VideoInputOp initialization\n";
|
|
std::cerr << exc.what();
|
|
}
|
|
}
|
|
|
|
template <class Context>
|
|
void VideoInputOp<Context>::GetLabelsFromProto(
|
|
const TensorProto& label_proto,
|
|
int* label_data) {
|
|
int num_clips = clip_per_video_ * crop_per_clip_;
|
|
if (!do_multi_label_) {
|
|
for (const auto i : c10::irange(num_clips)) {
|
|
label_data[i] = label_proto.int32_data(0);
|
|
}
|
|
} else {
|
|
// For multiple label case, output label is a binary vector
|
|
// where presented concepts are marked 1
|
|
memset(label_data, 0, sizeof(int) * num_of_class_ * num_clips);
|
|
for (const auto i : c10::irange(num_clips)) {
|
|
for (const auto j : c10::irange(label_proto.int32_data_size())) {
|
|
CAFFE_ENFORCE_LT(
|
|
label_proto.int32_data(j),
|
|
num_of_class_,
|
|
"Label should be less than the number of classes.");
|
|
label_data[i * num_of_class_ + label_proto.int32_data(j)] = 1;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <class Context>
|
|
bool VideoInputOp<Context>::GetImageAndLabelsFromDBValue(
|
|
const std::string& value,
|
|
int& height,
|
|
int& width,
|
|
std::vector<unsigned char*>& buffer_rgb,
|
|
int* label_data) {
|
|
TensorProtos protos;
|
|
CAFFE_ENFORCE(protos.ParseFromString(value));
|
|
const TensorProto& image_proto = protos.protos(0);
|
|
const TensorProto& label_proto = protos.protos(1);
|
|
|
|
GetLabelsFromProto(label_proto, label_data);
|
|
|
|
cv::Mat src;
|
|
if (image_proto.data_type() == TensorProto::STRING) {
|
|
// encoded image string.
|
|
TORCH_DCHECK_EQ(image_proto.string_data_size(), 1);
|
|
const string& encoded_image_str = image_proto.string_data(0);
|
|
int encoded_size = encoded_image_str.size();
|
|
// We use a cv::Mat to wrap the encoded str so we do not need a copy.
|
|
src = cv::imdecode(
|
|
cv::Mat(
|
|
1,
|
|
&encoded_size,
|
|
CV_8UC1,
|
|
const_cast<char*>(encoded_image_str.data())),
|
|
cv::IMREAD_COLOR);
|
|
if (src.rows == 0 || src.cols == 0) {
|
|
throw std::runtime_error("Both rows and cols are 0 for image");
|
|
}
|
|
} else if (image_proto.data_type() == TensorProto::BYTE) {
|
|
// raw image content.
|
|
int src_c = (image_proto.dims_size() == 3) ? image_proto.dims(2) : 1;
|
|
CAFFE_ENFORCE(src_c == 3 || src_c == 1);
|
|
|
|
src.create(
|
|
image_proto.dims(0),
|
|
image_proto.dims(1),
|
|
(src_c == 3) ? CV_8UC3 : CV_8UC1);
|
|
memcpy(
|
|
src.ptr<uchar>(0),
|
|
image_proto.byte_data().data(),
|
|
image_proto.byte_data().size());
|
|
} else {
|
|
throw std::runtime_error(
|
|
"Unknown image data type: " +
|
|
caffe2::to_string(image_proto.data_type()));
|
|
}
|
|
CAFFE_ENFORCE(src.isContinuous());
|
|
|
|
cv::Mat scaled_img;
|
|
cv::resize(
|
|
src, scaled_img, cv::Size(scale_w_, scale_h_), 0, 0, cv::INTER_AREA);
|
|
|
|
cv::Mat img;
|
|
if (channels_rgb_ == src.channels()) {
|
|
img = scaled_img;
|
|
} else {
|
|
cv::cvtColor(
|
|
scaled_img,
|
|
img,
|
|
(channels_rgb_ == 1) ? cv::COLOR_BGR2GRAY : cv::COLOR_GRAY2BGR);
|
|
}
|
|
|
|
cv::Mat rgb_img;
|
|
|
|
if (channels_rgb_ == 1) {
|
|
cv::cvtColor(img, rgb_img, cv::COLOR_BGR2RGB);
|
|
} else {
|
|
rgb_img = img;
|
|
}
|
|
CAFFE_ENFORCE(rgb_img.isContinuous());
|
|
|
|
unsigned char* data = new unsigned char[scale_h_ * scale_w_ * channels_rgb_];
|
|
memcpy(
|
|
data,
|
|
rgb_img.data,
|
|
scale_h_ * scale_w_ * channels_rgb_ * sizeof(unsigned char));
|
|
buffer_rgb.push_back(data);
|
|
width = scale_w_;
|
|
height = scale_h_;
|
|
return true;
|
|
}
|
|
|
|
template <class Context>
|
|
bool VideoInputOp<Context>::GetClipsAndLabelsFromDBValue(
|
|
const std::string& value,
|
|
int& height,
|
|
int& width,
|
|
std::vector<unsigned char*>& buffer_rgb,
|
|
int* label_data,
|
|
int64_t* video_id_data,
|
|
int* start_frame_data,
|
|
std::mt19937* randgen) {
|
|
TensorProtos protos;
|
|
int curr_proto_idx = 0;
|
|
CAFFE_ENFORCE(protos.ParseFromString(value));
|
|
const TensorProto& video_proto = protos.protos(curr_proto_idx++);
|
|
const TensorProto& label_proto = protos.protos(curr_proto_idx++);
|
|
|
|
int start_frm = 0;
|
|
int num_clips = clip_per_video_ * crop_per_clip_;
|
|
// start_frm is only valid when sampling 1 clip per video without
|
|
// temporal jitterring
|
|
if (decode_type_ == DecodeType::USE_START_FRM) {
|
|
CAFFE_ENFORCE_GE(
|
|
protos.protos_size(),
|
|
curr_proto_idx + 1,
|
|
"Start frm proto not provided");
|
|
const TensorProto& start_frm_proto = protos.protos(curr_proto_idx++);
|
|
start_frm = start_frm_proto.int32_data(0);
|
|
if (get_start_frame_) {
|
|
for (const auto i : c10::irange(num_clips)) {
|
|
start_frame_data[i] = start_frm;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (get_video_id_) {
|
|
CAFFE_ENFORCE_GE(
|
|
protos.protos_size(), curr_proto_idx + 1, "Video Id not provided");
|
|
const TensorProto& video_id_proto = protos.protos(curr_proto_idx);
|
|
for (const auto i : c10::irange(num_clips)) {
|
|
video_id_data[i] = video_id_proto.int64_data(0);
|
|
}
|
|
}
|
|
|
|
// assign labels
|
|
GetLabelsFromProto(label_proto, label_data);
|
|
|
|
if (use_local_file_) {
|
|
CAFFE_ENFORCE_EQ(
|
|
video_proto.data_type(),
|
|
TensorProto::STRING,
|
|
"Database with a file_list is expected to be string data");
|
|
}
|
|
|
|
// initializing the decoding params
|
|
Params params;
|
|
params.maximumOutputFrames_ = MAX_DECODING_FRAMES;
|
|
params.video_res_type_ = video_res_type_;
|
|
params.crop_size_ = crop_size_;
|
|
params.short_edge_ = short_edge_;
|
|
params.outputWidth_ = scale_w_;
|
|
params.outputHeight_ = scale_h_;
|
|
params.decode_type_ = decode_type_;
|
|
params.num_of_required_frame_ = num_of_required_frame_;
|
|
|
|
if (jitter_scales_.size() > 0) {
|
|
int select_idx =
|
|
std::uniform_int_distribution<>(0, jitter_scales_.size() - 1)(*randgen);
|
|
params.short_edge_ = jitter_scales_[select_idx];
|
|
}
|
|
|
|
char* video_buffer = nullptr; // for decoding from buffer
|
|
std::string video_filename; // for decoding from file
|
|
int encoded_size = 0;
|
|
if (video_proto.data_type() == TensorProto::STRING) {
|
|
const string& encoded_video_str = video_proto.string_data(0);
|
|
if (!use_local_file_) {
|
|
encoded_size = encoded_video_str.size();
|
|
video_buffer = const_cast<char*>(encoded_video_str.data());
|
|
} else {
|
|
video_filename = encoded_video_str;
|
|
}
|
|
} else if (video_proto.data_type() == TensorProto::BYTE) {
|
|
if (!use_local_file_) {
|
|
encoded_size = video_proto.byte_data().size();
|
|
video_buffer = const_cast<char*>(video_proto.byte_data().data());
|
|
} else {
|
|
// TODO: does this works?
|
|
video_filename = video_proto.string_data(0);
|
|
}
|
|
} else {
|
|
CAFFE_ENFORCE(false, "Unknown video data type.");
|
|
}
|
|
|
|
DecodeMultipleClipsFromVideo(
|
|
video_buffer,
|
|
video_filename,
|
|
encoded_size,
|
|
params,
|
|
start_frm,
|
|
clip_per_video_,
|
|
clip_start_positions_,
|
|
use_local_file_,
|
|
height,
|
|
width,
|
|
buffer_rgb);
|
|
return true;
|
|
}
|
|
|
|
template <class Context>
|
|
void VideoInputOp<Context>::DecodeAndTransform(
|
|
const std::string& value,
|
|
float* clip_rgb_data,
|
|
float* clip_of_data,
|
|
int* label_data,
|
|
int64_t* video_id_data,
|
|
int* start_frame_data,
|
|
std::mt19937* randgen,
|
|
std::bernoulli_distribution* mirror_this_clip) {
|
|
try {
|
|
std::vector<unsigned char*> buffer_rgb;
|
|
// get the video resolution after decoding
|
|
int height = 0;
|
|
int width = 0;
|
|
|
|
if (image_as_input_) {
|
|
CHECK(GetImageAndLabelsFromDBValue(
|
|
value, height, width, buffer_rgb, label_data));
|
|
} else {
|
|
// Decode the video from memory or read from a local file
|
|
CHECK(GetClipsAndLabelsFromDBValue(
|
|
value,
|
|
height,
|
|
width,
|
|
buffer_rgb,
|
|
label_data,
|
|
video_id_data,
|
|
start_frame_data,
|
|
randgen));
|
|
}
|
|
int clip_offset_rgb = channels_rgb_ * length_rgb_ * crop_size_ * crop_size_;
|
|
int clip_offset_of = channels_of_ * length_of_ * crop_size_ * crop_size_;
|
|
for (int i = 0; i < std::min(clip_per_video_, int(buffer_rgb.size()));
|
|
i++) {
|
|
for (const auto j : c10::irange(crop_per_clip_)) {
|
|
// get the rectangle for cropping
|
|
int h_off = 0;
|
|
int w_off = 0;
|
|
if (crop_per_clip_ > 1) {
|
|
CAFFE_ENFORCE(
|
|
random_crop_ == false,
|
|
"Only using multiple crops w/o random cropping");
|
|
}
|
|
if (random_crop_) {
|
|
// using random crop for training
|
|
h_off =
|
|
std::uniform_int_distribution<>(0, height - crop_size_)(*randgen);
|
|
w_off =
|
|
std::uniform_int_distribution<>(0, width - crop_size_)(*randgen);
|
|
} else {
|
|
// using multiple spatial crops
|
|
if (crop_per_clip_ > 1) { // normally 3 crops
|
|
if (height < width) {
|
|
h_off = (height - crop_size_) / 2;
|
|
w_off = j * (width - crop_size_) / (crop_per_clip_ - 1);
|
|
} else {
|
|
h_off = j * (height - crop_size_) / (crop_per_clip_ - 1);
|
|
w_off = (width - crop_size_) / 2;
|
|
}
|
|
// LOG(INFO) << "crop " << j << "-th " << h_off << " & " << w_off;
|
|
} else { // using center crop for testing
|
|
h_off = (height - crop_size_) / 2;
|
|
w_off = (width - crop_size_) / 2;
|
|
}
|
|
}
|
|
cv::Rect rect(w_off, h_off, crop_size_, crop_size_);
|
|
|
|
int this_clip_sampling_rate;
|
|
if (random_sampling_rate_) {
|
|
this_clip_sampling_rate = std::uniform_int_distribution<>(
|
|
1, random_sampling_rate_)(*randgen);
|
|
}
|
|
|
|
// randomly mirror the image or not
|
|
bool mirror_me = random_mirror_ && (*mirror_this_clip)(*randgen);
|
|
|
|
if (get_rgb_ && clip_rgb_data) {
|
|
ClipTransformRGB(
|
|
buffer_rgb[i],
|
|
crop_size_,
|
|
length_rgb_,
|
|
channels_rgb_,
|
|
(random_sampling_rate_ == 0 ? sampling_rate_rgb_
|
|
: this_clip_sampling_rate),
|
|
height,
|
|
width,
|
|
h_off,
|
|
w_off,
|
|
mirror_me,
|
|
mean_rgb_,
|
|
inv_std_rgb_,
|
|
clip_rgb_data + ((i * crop_per_clip_ + j) * clip_offset_rgb));
|
|
}
|
|
|
|
if (get_optical_flow_ && clip_of_data) {
|
|
ClipTransformOpticalFlow(
|
|
buffer_rgb[i],
|
|
crop_size_,
|
|
length_of_,
|
|
channels_of_,
|
|
sampling_rate_of_,
|
|
height,
|
|
width,
|
|
rect,
|
|
channels_rgb_,
|
|
mirror_me,
|
|
flow_alg_type_,
|
|
flow_data_type_,
|
|
frame_gap_of_,
|
|
do_flow_aggregation_,
|
|
mean_of_,
|
|
inv_std_of_,
|
|
clip_of_data + ((i * crop_per_clip_ + j) * clip_offset_of));
|
|
}
|
|
}
|
|
}
|
|
if (buffer_rgb.size() > 0) {
|
|
for (const auto i : c10::irange(buffer_rgb.size())) {
|
|
unsigned char* buff = buffer_rgb[i];
|
|
delete[] buff;
|
|
}
|
|
}
|
|
buffer_rgb.clear();
|
|
} catch (const std::exception& exc) {
|
|
std::cerr << "While calling DecodeAndTransform()\n";
|
|
std::cerr << exc.what();
|
|
}
|
|
}
|
|
|
|
template <class Context>
|
|
bool VideoInputOp<Context>::Prefetch() {
|
|
try {
|
|
// We will get the reader pointer from input.
|
|
// If we use local clips, db will store the list
|
|
reader_ = &OperatorBase::Input<db::DBReader>(0);
|
|
|
|
// Call mutable_data() once to allocate the underlying memory.
|
|
prefetched_clip_rgb_.mutable_data<float>();
|
|
prefetched_clip_of_.mutable_data<float>();
|
|
prefetched_label_.mutable_data<int>();
|
|
prefetched_video_id_.mutable_data<int64_t>();
|
|
prefetched_start_frame_.mutable_data<int>();
|
|
|
|
// Prefetching handled with a thread pool of "decode_threads" threads.
|
|
std::mt19937 meta_randgen(time(nullptr));
|
|
std::vector<std::mt19937> randgen_per_thread;
|
|
for (const auto i : c10::irange(num_decode_threads_)) {
|
|
randgen_per_thread.emplace_back(meta_randgen());
|
|
}
|
|
|
|
std::bernoulli_distribution mirror_this_clip(0.5);
|
|
for (const auto item_id : c10::irange(batch_size_)) {
|
|
std::mt19937* randgen =
|
|
&randgen_per_thread[item_id % num_decode_threads_];
|
|
|
|
int frame_size = crop_size_ * crop_size_;
|
|
// get the clip data pointer for the item_id -th example
|
|
float* clip_rgb_data = prefetched_clip_rgb_.mutable_data<float>() +
|
|
frame_size * length_rgb_ * channels_rgb_ * item_id * clip_per_video_ *
|
|
crop_per_clip_;
|
|
|
|
// get the optical flow data for the current clip
|
|
float* clip_of_data = prefetched_clip_of_.mutable_data<float>() +
|
|
frame_size * length_of_ * channels_of_ * item_id * clip_per_video_ *
|
|
crop_per_clip_;
|
|
|
|
// get the label data pointer for the item_id -th example
|
|
int* label_data = prefetched_label_.mutable_data<int>() +
|
|
(do_multi_label_ ? num_of_class_ : 1) * item_id * clip_per_video_ *
|
|
crop_per_clip_;
|
|
|
|
// get the video id data pointer for the item_id -th example
|
|
int64_t* video_id_data = prefetched_video_id_.mutable_data<int64_t>() +
|
|
item_id * clip_per_video_ * crop_per_clip_;
|
|
|
|
int* start_frame_data = prefetched_start_frame_.mutable_data<int>() +
|
|
item_id * clip_per_video_ * crop_per_clip_;
|
|
|
|
std::string key, value;
|
|
// read data
|
|
reader_->Read(&key, &value);
|
|
|
|
thread_pool_->run(std::bind(
|
|
&VideoInputOp<Context>::DecodeAndTransform,
|
|
this,
|
|
std::string(value),
|
|
clip_rgb_data,
|
|
clip_of_data,
|
|
label_data,
|
|
video_id_data,
|
|
start_frame_data,
|
|
randgen,
|
|
&mirror_this_clip));
|
|
} // for over the batch
|
|
thread_pool_->waitWorkComplete();
|
|
|
|
// If the context is not CPUContext, we will need to do a copy in the
|
|
// prefetch function as well.
|
|
if (!std::is_same<Context, CPUContext>::value) {
|
|
if (get_rgb_) {
|
|
prefetched_clip_rgb_on_device_.CopyFrom(
|
|
prefetched_clip_rgb_, &context_);
|
|
}
|
|
if (get_optical_flow_) {
|
|
prefetched_clip_of_on_device_.CopyFrom(prefetched_clip_of_, &context_);
|
|
}
|
|
prefetched_label_on_device_.CopyFrom(prefetched_label_, &context_);
|
|
if (get_video_id_) {
|
|
prefetched_video_id_on_device_.CopyFrom(
|
|
prefetched_video_id_, &context_);
|
|
}
|
|
if (get_start_frame_) {
|
|
prefetched_start_frame_on_device_.CopyFrom(
|
|
prefetched_start_frame_, &context_);
|
|
}
|
|
}
|
|
} catch (const std::exception& exc) {
|
|
std::cerr << "While calling Prefetch()\n";
|
|
std::cerr << exc.what();
|
|
}
|
|
return true;
|
|
}
|
|
|
|
template <class Context>
|
|
bool VideoInputOp<Context>::CopyPrefetched() {
|
|
try {
|
|
int index = 0;
|
|
auto type = Context::GetDeviceType();
|
|
if (get_rgb_) {
|
|
auto* clip_rgb_output = OperatorBase::Output<Tensor>(index++, type);
|
|
if (std::is_same<Context, CPUContext>::value) {
|
|
clip_rgb_output->CopyFrom(prefetched_clip_rgb_, &context_);
|
|
} else {
|
|
clip_rgb_output->CopyFrom(prefetched_clip_rgb_on_device_, &context_);
|
|
}
|
|
}
|
|
|
|
if (get_optical_flow_) {
|
|
auto* clip_of_output = OperatorBase::Output<Tensor>(index++, type);
|
|
if (std::is_same<Context, CPUContext>::value) {
|
|
clip_of_output->CopyFrom(prefetched_clip_of_, &context_);
|
|
} else {
|
|
clip_of_output->CopyFrom(prefetched_clip_of_on_device_, &context_);
|
|
}
|
|
}
|
|
|
|
auto* label_output = OperatorBase::Output<Tensor>(index++, type);
|
|
if (std::is_same<Context, CPUContext>::value) {
|
|
label_output->CopyFrom(prefetched_label_, &context_);
|
|
} else {
|
|
label_output->CopyFrom(prefetched_label_on_device_, &context_);
|
|
}
|
|
|
|
if (get_video_id_) {
|
|
auto* video_id_output = OperatorBase::Output<Tensor>(index++, type);
|
|
if (std::is_same<Context, CPUContext>::value) {
|
|
video_id_output->CopyFrom(prefetched_video_id_, &context_);
|
|
} else {
|
|
video_id_output->CopyFrom(prefetched_video_id_on_device_, &context_);
|
|
}
|
|
}
|
|
if (get_start_frame_) {
|
|
auto* start_frame_output = OperatorBase::Output<Tensor>(index, type);
|
|
if (std::is_same<Context, CPUContext>::value) {
|
|
start_frame_output->CopyFrom(prefetched_start_frame_, &context_);
|
|
} else {
|
|
start_frame_output->CopyFrom(
|
|
prefetched_start_frame_on_device_, &context_);
|
|
}
|
|
}
|
|
} catch (const std::exception& exc) {
|
|
std::cerr << "While calling CopyPrefetched()\n";
|
|
std::cerr << exc.what();
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
} // namespace caffe2
|
|
|
|
#endif // CAFFE2_VIDEO_VIDEO_INPUT_OP_H_
|