Skip to content

Instantly share code, notes, and snippets.

@dataplayer12

dataplayer12/CMakeLists.txt

Last active August 10, 2020 12:11
Show Gist options
  • Save dataplayer12/6b1fe64b783d2b319f43bfa2f8bd9a5c to your computer and use it in GitHub Desktop.
Save dataplayer12/6b1fe64b783d2b319f43bfa2f8bd9a5c to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. dataplayer12 revised this gist Nov 13, 2019. No changes.
  2. dataplayer12 created this gist Nov 13, 2019.
    2 changes: 2 additions & 0 deletions CMakeLists.txt
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,2 @@
    add_executable(infervideo infervideo.cpp)
    target_link_libraries(infervideo PRIVATE retinanet ${OpenCV_LIBS} cuda ${CUDA_LIBRARIES})
    125 changes: 125 additions & 0 deletions infervideo.cpp
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,125 @@
    #include <iostream>
    #include <stdexcept>
    #include <fstream>
    #include <vector>
    #include <chrono>

    #include <opencv2/opencv.hpp>
    #include <opencv2/opencv.hpp>
    #include <opencv2/core/core.hpp>
    #include <opencv2/highgui/highgui.hpp>

    #include <cuda_runtime.h>

    #include "../../csrc/engine.h"

    using namespace std;
    using namespace cv;

    int main(int argc, char *argv[]) {
    if (argc != 4) {
    cerr << "Usage: " << argv[0] << " engine.plan input.mov output.mp4" << endl;
    return 1;
    }

    cout << "Loading engine..." << endl;
    auto engine = retinanet::Engine(argv[1]);
    VideoCapture src(argv[2]);

    if (!src.isOpened()){
    cerr << "Could not read " << argv[2] << endl;
    return 1;
    }

    auto fh=src.get(CV_CAP_PROP_FRAME_HEIGHT);
    auto fw=src.get(CV_CAP_PROP_FRAME_WIDTH);
    auto fps=src.get(CV_CAP_PROP_FPS);
    auto nframes=src.get(CV_CAP_PROP_FRAME_COUNT);

    VideoWriter sink;
    sink.open(argv[3], 0x31637661, fps, Size(fw, fh));
    Mat frame;
    Mat resized_frame;
    Mat inferred_frame;
    int count=1;

    auto inputSize = engine.getInputSize();
    // Create device buffers
    void *data_d, *scores_d, *boxes_d, *classes_d;
    auto num_det = engine.getMaxDetections();
    cudaMalloc(&data_d, 3 * inputSize[0] * inputSize[1] * sizeof(float));
    cudaMalloc(&scores_d, num_det * sizeof(float));
    cudaMalloc(&boxes_d, num_det * 4 * sizeof(float));
    cudaMalloc(&classes_d, num_det * sizeof(float));

    auto scores = new float[num_det];
    auto boxes = new float[num_det * 4];
    auto classes = new float[num_det];

    vector<float> mean {0.485, 0.456, 0.406};
    vector<float> std {0.229, 0.224, 0.225};

    vector<uint8_t> blues {0,63,127,191,255,0}; //colors for bonuding boxes
    vector<uint8_t> greens {0,255,191,127,63,0};
    vector<uint8_t> reds {191,255,0,0,63,127};

    int channels = 3;
    vector<float> img;
    vector<float> data (channels * inputSize[0] * inputSize[1]);

    while (1){
    src >> frame;
    if (frame.empty()){
    cout << "Finished inference!" << endl;
    break;
    }

    cv::resize(frame, resized_frame, Size(inputSize[1], inputSize[0]));
    cv::Mat pixels;
    resized_frame.convertTo(pixels, CV_32FC3, 1.0 / 255, 0);

    img.assign((float*)pixels.datastart, (float*)pixels.dataend);

    for (int c = 0; c < channels; c++) {
    for (int j = 0, hw = inputSize[0] * inputSize[1]; j < hw; j++) {
    data[c * hw + j] = (img[channels * j + 2 - c] - mean[c]) / std[c];
    }
    }

    // Copy image to device
    size_t dataSize = data.size() * sizeof(float);
    cudaMemcpy(data_d, data.data(), dataSize, cudaMemcpyHostToDevice);


    //Do inference
    cout << "Inferring on frame: " << count <<"/" << nframes << endl;
    count++;
    vector<void *> buffers = { data_d, scores_d, boxes_d, classes_d };
    engine.infer(buffers, 1);

    cudaMemcpy(scores, scores_d, sizeof(float) * num_det, cudaMemcpyDeviceToHost);
    cudaMemcpy(boxes, boxes_d, sizeof(float) * num_det * 4, cudaMemcpyDeviceToHost);
    cudaMemcpy(classes, classes_d, sizeof(float) * num_det, cudaMemcpyDeviceToHost);

    // Get back the bounding boxes
    for (int i = 0; i < num_det; i++) {
    // Show results over confidence threshold
    if (scores[i] >= 0.2f) {
    float x1 = boxes[i*4+0];
    float y1 = boxes[i*4+1];
    float x2 = boxes[i*4+2];
    float y2 = boxes[i*4+3];
    int cls=classes[i];
    // Draw bounding box on image
    cv::rectangle(resized_frame, Point(x1, y1), Point(x2, y2), cv::Scalar(blues[cls], greens[cls], reds[cls]));
    }
    }
    //delete[] scores, boxes, classes;
    cv::resize(resized_frame, inferred_frame, Size(fw, fh));
    sink.write(inferred_frame);
    //if (count>100){break;}
    }
    src.release();
    sink.release();
    return 0;
    }