141 lines
4.3 KiB
C++
Executable file
141 lines
4.3 KiB
C++
Executable file
/*
|
||
* Copyright (C) 2024–2025 Amlogic, Inc. All rights reserved.
|
||
*
|
||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||
* you may not use this file except in compliance with the License.
|
||
* You may obtain a copy of the License at
|
||
*
|
||
* http://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.
|
||
*/
|
||
|
||
#include <iostream>
|
||
#include <string>
|
||
#include <vector>
|
||
#include <chrono>
|
||
#include <tuple>
|
||
#include <iomanip>
|
||
#include <fstream>
|
||
#include <opencv2/opencv.hpp>
|
||
#include "postprocess.h"
|
||
#include "model_loader.h"
|
||
|
||
const std::string DEFAULT_OUTPUT_PATH = "./result.jpg";
|
||
const int MODEL_INPUT_WIDTH = 256;
|
||
const int MODEL_INPUT_HEIGHT = 256;
|
||
const float SCORE_THRESHOLD = 0.5f;
|
||
|
||
int main(int argc, char **argv)
|
||
{
|
||
std::string model_path;
|
||
std::string image_path;
|
||
|
||
if (argc != 3)
|
||
{
|
||
printf("%s <model_path> <image_path>\n", argv[0]);
|
||
return -1;
|
||
}
|
||
|
||
if (argc > 1)
|
||
model_path = argv[1];
|
||
if (argc > 2)
|
||
image_path = argv[2];
|
||
|
||
std::cout << "Blazepose Detect Demo" << std::endl;
|
||
std::cout << "Model: " << model_path << std::endl;
|
||
std::cout << "Image: " << image_path << std::endl;
|
||
std::cout << "Output: " << DEFAULT_OUTPUT_PATH << std::endl;
|
||
|
||
// 1. Load Image
|
||
cv::Mat img = cv::imread(image_path);
|
||
if (img.empty())
|
||
{
|
||
std::cerr << "Failed to load image from " << image_path << std::endl;
|
||
return -1;
|
||
}
|
||
// Load detections
|
||
// n * 13 detections
|
||
// image_path -> txt_path
|
||
std::vector<std::vector<float>> detections;
|
||
|
||
std::string txt_path = image_path.substr(0, image_path.find_last_of('.'));
|
||
txt_path += ".txt";
|
||
std::ifstream ifs(txt_path);
|
||
for (std::string line; std::getline(ifs, line);)
|
||
{
|
||
std::istringstream iss(line);
|
||
std::vector<float> det;
|
||
float val;
|
||
while (iss >> val)
|
||
det.push_back(val);
|
||
if (!det.empty())
|
||
detections.push_back(det);
|
||
}
|
||
|
||
// 2. Initialize Network
|
||
void *context = init_network(model_path.c_str());
|
||
if (!context)
|
||
{
|
||
std::cerr << "Failed to initialize network." << std::endl;
|
||
return -1;
|
||
}
|
||
|
||
// 3. Preprocess
|
||
auto start_time = std::chrono::high_resolution_clock::now();
|
||
|
||
auto [preprocessed, affine] = preprocess(img, detections, std::make_tuple(MODEL_INPUT_HEIGHT, MODEL_INPUT_WIDTH));
|
||
|
||
// Quantize to int16 (model expects quantized input)
|
||
cv::Mat quantized_img = quantize_input(preprocessed, 0.000030518509447574615f);
|
||
|
||
// 4. Set input and run inference
|
||
nn_input inData;
|
||
memset(&inData, 0, sizeof(nn_input));
|
||
inData.input_type = BINARY_RAW_DATA;
|
||
inData.input = quantized_img.data;
|
||
inData.input_index = 0;
|
||
inData.size = quantized_img.total() * quantized_img.elemSize();
|
||
|
||
if (aml_module_input_set(context, &inData) != 0)
|
||
{
|
||
std::cerr << "Failed to set input." << std::endl;
|
||
uninit_network(context);
|
||
return -1;
|
||
}
|
||
|
||
aml_output_config_t outconfig;
|
||
memset(&outconfig, 0, sizeof(aml_output_config_t));
|
||
outconfig.typeSize = sizeof(aml_output_config_t);
|
||
outconfig.format = AML_OUTDATA_FLOAT32;
|
||
|
||
nn_output *outdata = (nn_output *)aml_module_output_get(context, outconfig);
|
||
if (!outdata)
|
||
{
|
||
std::cerr << "Failed to run network." << std::endl;
|
||
uninit_network(context);
|
||
return -1;
|
||
}
|
||
|
||
// 5. Postprocess
|
||
std::vector<BlazePoseLandmark> landmarks = postprocess(outdata, affine);
|
||
auto end_time = std::chrono::high_resolution_clock::now();
|
||
std::chrono::duration<double, std::milli> inference_time = end_time - start_time;
|
||
|
||
std::cout << "Inference time: " << inference_time.count() << " ms" << std::endl;
|
||
std::cout << "Landmarks: " << landmarks.size() << std::endl;
|
||
|
||
// 6. Draw and Save
|
||
cv::Mat result_img = draw_landmarks(img, landmarks, SCORE_THRESHOLD);
|
||
cv::imwrite(DEFAULT_OUTPUT_PATH, result_img);
|
||
std::cout << "Result saved to " << DEFAULT_OUTPUT_PATH << std::endl;
|
||
|
||
// 7. Cleanup
|
||
uninit_network(context);
|
||
|
||
return 0;
|
||
}
|