89 lines
2.5 KiB
C++
89 lines
2.5 KiB
C++
|
|
#include <map>
|
|
#include <string>
|
|
|
|
#include "mmdeploy/video_recognizer.hpp"
|
|
#include "opencv2/videoio.hpp"
|
|
#include "utils/argparse.h"
|
|
|
|
void SampleFrames(const char* video_path, std::map<int, cv::Mat>& buffer,
|
|
std::vector<mmdeploy::Mat>& clips, int clip_len, int frame_interval = 1,
|
|
int num_clips = 1) {
|
|
cv::VideoCapture cap = cv::VideoCapture(video_path);
|
|
if (!cap.isOpened()) {
|
|
fprintf(stderr, "failed to load video: %s\n", video_path);
|
|
exit(1);
|
|
}
|
|
|
|
int num_frames = cap.get(cv::CAP_PROP_FRAME_COUNT);
|
|
printf("num_frames %d\n", num_frames);
|
|
|
|
int ori_clip_len = clip_len * frame_interval;
|
|
float avg_interval = (num_frames - ori_clip_len + 1.f) / num_clips;
|
|
std::vector<int> frame_inds;
|
|
for (int i = 0; i < num_clips; i++) {
|
|
int clip_offset = i * avg_interval + avg_interval / 2.0;
|
|
for (int j = 0; j < clip_len; j++) {
|
|
int ind = (j * frame_interval + clip_offset) % num_frames;
|
|
if (num_frames <= ori_clip_len - 1) {
|
|
ind = j % num_frames;
|
|
}
|
|
frame_inds.push_back(ind);
|
|
}
|
|
}
|
|
|
|
std::vector<int> unique_inds(frame_inds.begin(), frame_inds.end());
|
|
std::sort(unique_inds.begin(), unique_inds.end());
|
|
auto last = std::unique(unique_inds.begin(), unique_inds.end());
|
|
unique_inds.erase(last, unique_inds.end());
|
|
|
|
int ind = 0;
|
|
for (int i = 0; i < unique_inds.size(); i++) {
|
|
int tid = unique_inds[i];
|
|
cv::Mat frame;
|
|
while (ind < tid) {
|
|
cap.read(frame);
|
|
ind++;
|
|
}
|
|
cap.read(frame);
|
|
buffer[tid] = frame;
|
|
ind++;
|
|
}
|
|
|
|
clips.resize(frame_inds.size());
|
|
for (int i = 0; i < frame_inds.size(); i++) {
|
|
auto& img = buffer[frame_inds[i]];
|
|
clips[i] = img;
|
|
}
|
|
}
|
|
|
|
DEFINE_ARG_string(model, "Model path");
|
|
DEFINE_ARG_string(video, "Input video path");
|
|
DEFINE_string(device, "cpu", R"(Device name, e.g. "cpu", "cuda")");
|
|
|
|
int main(int argc, char* argv[]) {
|
|
if (!utils::ParseArguments(argc, argv)) {
|
|
return -1;
|
|
}
|
|
|
|
int clip_len = 1;
|
|
int frame_interval = 1;
|
|
int num_clips = 25;
|
|
|
|
std::map<int, cv::Mat> buffer;
|
|
std::vector<mmdeploy::Mat> clips;
|
|
mmdeploy::VideoSampleInfo clip_info = {clip_len, num_clips};
|
|
SampleFrames(ARGS_video.c_str(), buffer, clips, clip_len, frame_interval, num_clips);
|
|
|
|
mmdeploy::Model model(ARGS_model);
|
|
mmdeploy::VideoRecognizer recognizer(model, mmdeploy::Device{FLAGS_device});
|
|
|
|
auto res = recognizer.Apply(clips, clip_info);
|
|
|
|
for (const auto& cls : res) {
|
|
fprintf(stderr, "label: %d, score: %.4f\n", cls.label_id, cls.score);
|
|
}
|
|
|
|
return 0;
|
|
}
|