1#include "blace_ai.h"
 2#include <opencv2/opencv.hpp>
 3
 4// include the models you want to use
 5#include "raft_v1_default_v2_ALL_export_version_v17.h"
 6
 7/**
 8x
 9*/
10
11std::shared_ptr<blace::RawMemoryObject> memory_from_file(std::string file) {
12  // read image into memory
13  cv::Mat image = cv::imread(file, cv::IMREAD_COLOR);
14
15  // construct a hash from the filename
16  blace::ml_core::BlaceHash random_hash(file);
17
18  // construct the memory object. We set copy_memory to true, since image will
19  // be out-of-scope upon method return and therefore we need to take ownership
20  // of the data
21  blace::RawMemoryObject raw_mem(
22      (void *)image.data, blace::ml_core::DataTypeEnum::BLACE_BYTE,
23      blace::ml_core::ColorFormatEnum::BGR,
24      std::vector<int64_t>{1, image.rows, image.cols, 3}, blace::ml_core::BHWC,
25      blace::ml_core::ZERO_TO_255, random_hash, true);
26
27  return std::make_shared<blace::RawMemoryObject>(raw_mem);
28}
29
30int main() {
31  ::blace::workload_management::BlaceWorld blace;
32
33  // register model at server
34  blace::util::registerModel(raft_v1_default_v2_ALL_export_version_v17,
35                             blace::util::getPathToExe());
36
37  // load image into op
38  auto exe_path = blace::util::getPathToExe();
39  std::filesystem::path frame_0 = exe_path / "raft_frame_0.png";
40  std::filesystem::path frame_1 = exe_path / "raft_frame_1.png";
41
42  auto frame_0_mem = memory_from_file(frame_0.string());
43  auto frame_1_mem = memory_from_file(frame_1.string());
44
45  auto frame_0_op = CONSTRUCT_OP(blace::ops::FromRawMemoryOp(frame_0_mem));
46  auto frame_1_op = CONSTRUCT_OP(blace::ops::FromRawMemoryOp(frame_1_mem));
47
48  // construct model inference arguments
49  blace::ml_core::InferenceArgsCollection infer_args;
50  infer_args.inference_args.device = blace::util::get_accelerator().value();
51
52  // construct inference operation
53  auto infer_op = CONSTRUCT_OP(
54      blace::ops::InferenceOp(raft_v1_default_v2_ALL_export_version_v17_IDENT,
55                              {frame_0_op, frame_1_op}, infer_args, 0));
56
57  // normalize optical flow to zero-one range for plotting. The model returns
58  // relative offsets in -1 to 1 pixel space, so the raw values are to small to
59  // plot
60  infer_op = CONSTRUCT_OP(blace::ops::NormalizeToZeroOneOP(infer_op));
61
62  // convert uv color to rgb (by U->R, V->G, 1->B)
63  infer_op = CONSTRUCT_OP(blace::ops::ToColorOp(infer_op, blace::ml_core::RGB));
64
65  // we prepare the result for later copy to cv::Mat. The values set here are
66  // based on implicit knowledge of cv::Mat internal data storage.
67  auto normalized_matte = CONSTRUCT_OP(blace::ops::PrepareForHostCopyOP(
68      infer_op, blace::ml_core::BLACE_BYTE, blace::ml_core::RGB,
69      blace::ml_core::HWC, blace::ml_core::ZERO_TO_255, blace::ml_core::CPU));
70
71  // construct evaluator and evaluate to raw memory object
72  blace::computation_graph::GraphEvaluator evaluator(normalized_matte);
73  auto raw_mem = evaluator.evaluateToRawMemory().value();
74
75  // get the sizes
76  int w = raw_mem.get_memory_sizes()[1];
77  int h = raw_mem.get_memory_sizes()[0];
78
79  // initialize an empty cv::Mat
80  cv::Mat cv_mat(h, w, CV_8UC3);
81  cv_mat.setTo(cv::Scalar(0, 0, 0));
82
83  // and copy the memory
84  std::memcpy(cv_mat.data, raw_mem.get_data_ptr(), raw_mem.get_memory_size());
85
86  // save to disk and return
87  auto out_file = exe_path / "optical_flow.png";
88  cv::imwrite(out_file.string(), cv_mat);
89
90  return 0;
91}