Hugging Face's logo

Riser
/
YOLOP

Object Detection
Model card Files Community
YOLOP
File size: 5,315 Bytes 
67bb36a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
 
#ifndef TRTX_YOLOV5_UTILS_H_
#define TRTX_YOLOV5_UTILS_H_

#include <dirent.h>
#include <opencv2/opencv.hpp>
#include <opencv2/cudawarping.hpp>
#include <opencv2/cudaimgproc.hpp>
#include <opencv2/cudaarithm.hpp>
#include <opencv2/highgui.hpp>
#include <iostream>
#include "common.hpp"

#define SHOW_IMG

static inline cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) {
    int w, h, x, y;
    float r_w = input_w / (img.cols*1.0);
    float r_h = input_h / (img.rows*1.0);
    if (r_h > r_w) {
        w = input_w;
        h = r_w * img.rows;
        x = 0;
        y = (input_h - h) / 2;
    } else {
        w = r_h * img.cols;
        h = input_h;
        x = (input_w - w) / 2;
        y = 0;
    }
    cv::Mat re(h, w, CV_8UC3);
    cv::resize(img, re, re.size(), 0, 0, cv::INTER_LINEAR);
    cv::Mat out(input_h, input_w, CV_8UC3, cv::Scalar(114, 114, 114));
    re.copyTo(out(cv::Rect(x, y, re.cols, re.rows)));
    cv::Mat tensor;
    out.convertTo(tensor, CV_32FC3, 1.f / 255.f);
    
    cv::subtract(tensor, cv::Scalar(0.485, 0.456, 0.406), tensor, cv::noArray(), -1);
    cv::divide(tensor, cv::Scalar(0.229, 0.224, 0.225), tensor, 1, -1);
    // std::cout << cv::format(out, cv::Formatter::FMT_NUMPY)<< std::endl;
    // assert(false);
    // cv::Mat out(input_h, input_w, CV_8UC3);
    // cv::copyMakeBorder(re, out, y, y, x, x, cv::BORDER_CONSTANT, cv::Scalar(128, 128, 128));
    return tensor;
}

void preprocess_img_gpu(cv::cuda::GpuMat& img, float* gpu_input, int input_w, int input_h) {
    int w, h, x, y;
    float r_w = input_w / (img.cols*1.0);
    float r_h = input_h / (img.rows*1.0);
    if (r_h > r_w) {
        w = input_w;
        h = r_w * img.rows;
        x = 0;
        y = (input_h - h) / 2;
    } else {
        w = r_h * img.cols;
        h = input_h;
        x = (input_w - w) / 2;
        y = 0;
    }
    cv::cuda::GpuMat re(h, w, CV_8UC3);
    cv::cuda::resize(img, re, re.size(), 0, 0, cv::INTER_LINEAR);
    cv::cuda::GpuMat out(input_h, input_w, CV_8UC3, cv::Scalar(114, 114, 114));
    cv::cuda::GpuMat tensor;
    re.copyTo(out(cv::Rect(x, y, re.cols, re.rows)));
    out.convertTo(tensor, CV_32FC3, 1.f / 255.f);
    cv::cuda::subtract(tensor, cv::Scalar(0.485, 0.456, 0.406), tensor, cv::noArray(), -1);
    cv::cuda::divide(tensor, cv::Scalar(0.229, 0.224, 0.225), tensor, 1, -1);
    // cv::Mat out(input_h, input_w, CV_8UC3);
    // cv::copyMakeBorder(re, out, y, y, x, x, cv::BORDER_CONSTANT, cv::Scalar(128, 128, 128));

    // to tensor
    std::vector<cv::cuda::GpuMat> chw;
    for (size_t i = 0; i < 3; ++i)
    {
        chw.emplace_back(cv::cuda::GpuMat(tensor.size(), CV_32FC1, gpu_input + i * input_w * input_h));
    }
    cv::cuda::split(tensor, chw);
}

static inline int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_names) {
    DIR *p_dir = opendir(p_dir_name);
    if (p_dir == nullptr) {
        return -1;
    }

    struct dirent* p_file = nullptr;
    while ((p_file = readdir(p_dir)) != nullptr) {
        if (strcmp(p_file->d_name, ".") != 0 &&
            strcmp(p_file->d_name, "..") != 0) {
            //std::string cur_file_name(p_dir_name);
            //cur_file_name += "/";
            //cur_file_name += p_file->d_name;
            std::string cur_file_name(p_file->d_name);
            file_names.push_back(cur_file_name);
        }
    }

    closedir(p_dir);
    return 0;
}

void PrintMat(cv::Mat &A)
{
  for(int i=0; i<A.rows; i++)
  {
    for(int j=0; j<A.cols; j++)
        std::cout << A.at<int>(i,j) << ' ';
    std::cout << std::endl;
  }
  std::cout << std::endl;
}

void visualization(cv::cuda::GpuMat& cvt_img, cv::Mat& seg_res, cv::Mat& lane_res, std::vector<Yolo::Detection>& res, char& key)
{
    static const std::vector<cv::Vec3b> segColor{cv::Vec3b(0, 0, 0), cv::Vec3b(0, 255, 0), cv::Vec3b(255, 0, 0)};
    static const std::vector<cv::Vec3b> laneColor{cv::Vec3b(0, 0, 0), cv::Vec3b(0, 0, 255), cv::Vec3b(0, 0, 0)};
    cv::Mat cvt_img_cpu;
    cvt_img.download(cvt_img_cpu);

    // handling seg and lane results
    for (int row = 0; row < cvt_img_cpu.rows; ++row) {
        uchar* pdata = cvt_img_cpu.data + row * cvt_img_cpu.step;
        for (int col = 0; col < cvt_img_cpu.cols; ++col) {
            int seg_idx = seg_res.at<int>(row, col);
            int lane_idx = lane_res.at<int>(row, col);
            //std::cout << "enter" << ix << std::endl;
            for (int i = 0; i < 3; ++i) {
                if (lane_idx) {
                    if (i != 2)
                        pdata[i] = pdata[i] / 2 + laneColor[lane_idx][i] / 2;
                }
                else if (seg_idx)
                    pdata[i] = pdata[i] / 2 + segColor[seg_idx][i] / 2;
            }
            pdata += 3;
        }
    }

    // handling det results
    for (size_t j = 0; j < res.size(); ++j) {
        cv::Rect r = get_rect(cvt_img_cpu, res[j].bbox);
        cv::rectangle(cvt_img_cpu, r, cv::Scalar(0x27, 0xC1, 0x36), 2);
        cv::putText(cvt_img_cpu, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
    }
    
#ifdef SHOW_IMG
    cv::imshow("img", cvt_img_cpu);
    key = cv::waitKey(1);
#else
    cv::imwrite("../zed_result.jpg", cvt_img_cpu);
#endif
}

#endif  // TRTX_YOLOV5_UTILS_H_