2024-10-18 18:25:09 -05:00

63 lines
1.8 KiB
C++

/* Computer vision demo. Performs object tracking on pumpkins to determine velocity, launch angle, and predicted range.
*
*/
#include <iostream>
#include <DarkHelp.hpp>
#include <unistd.h>
#include <opencv4/opencv2/core/mat.hpp>
#include <opencv4/opencv2/videoio.hpp>
#include <opencv4/opencv2/highgui.hpp>
// Default configuration values. Using YOLOv7 tiny model. Make sure to download the weights.
#define DARKNET_CFG "cfg/yolov7-tiny.cfg"
#define DARKNET_WEIGHTS "cfg/yolov7-tiny.weights"
#define DARKNET_DATA "cfg/coco.names"
#define CAMERA_ID 0
#define SHOW_GUI true
using namespace std;
using namespace cv;
using namespace DarkHelp;
int main() {
// Configure the neural network
NN nn(DARKNET_CFG, DARKNET_WEIGHTS, DARKNET_DATA);
// Setup the capture stream from the webcam
VideoCapture cap(CAMERA_ID);
if (not cap.isOpened())
{
throw std::runtime_error("failed to open the webcam");
}
cap.set(CAP_PROP_FRAME_WIDTH, 640.0);
cap.set(CAP_PROP_FRAME_HEIGHT, 480.0);
cap.set(CAP_PROP_FPS, 30.0);
while (cap.isOpened()) {
Mat frame;
cap >> frame;
if (frame.empty()) break;
const PredictionResults results = nn.predict(frame);
// print results to console
cout << "RESULTS" << endl;
for (int i = 0; i < results.size(); i++) {
const PredictionResult result = results[i];
cout << result.name << " | (" << result.original_point.x << ", " << result.original_point.y << ")" << endl;
}
cout << endl;
// Sleep between frames; if the UI is enabled, show the latest (annotated) frame
if constexpr (SHOW_GUI) {
frame = nn.annotate();
imshow("pumpkin-tracker", frame);
if (const auto key = waitKey(15); key == 27) break;
}
else {
usleep(15000);
}
}
}