Untitled

 avatar
unknown
plain_text
a year ago
2.3 kB
4
Indexable
int main() {
    Yolov8SegOnnx task_segment_onnx;

    vector<Scalar> color;
    srand(time(0));
    for (int i = 0; i < 80; i++) {
        int b = rand() % 256;
        int g = rand() % 256;
        int r = rand() % 256;
        color.push_back(Scalar(b, g, r));
    }

    for (const auto& entry : std::filesystem::directory_iterator(img_folder)) {
        string img_path = entry.path().string();
        Mat src = imread(img_path);
        std::string file_name = entry.path().filename().string();

        if (!src.empty()) {
            cout << "Image is read!!" << endl;
            cout << file_name << endl;

            if (task_segment_onnx.ReadModel(seg_model_path, false)) {
                cout << "Model loaded successfully!" << endl;
                vector<OutputSeg> result;
                if (task_segment_onnx.OnnxDetect(src, result)) {
                    cout << "Detection successful!" << endl;
                    DrawPred(src, result, task_segment_onnx._className, color);
                    cout << "Drawing successful!" << endl;

                    cv::Mat canvas(src.rows, src.cols, CV_8UC1, cv::Scalar(0, 0, 0));
                    if (!result.empty()) {
                        std::cout << result.size() << std::endl;

                        for (const auto& OutputSeg : result) {
                            cv::Mat mask = OutputSeg.boxMask;
                            cv::Rect roi = OutputSeg.box;
                            cv::Mat canvasROI = canvas(roi);
                            cv::bitwise_or(canvasROI, mask, canvasROI);
                        }

                        cv::imwrite(output_path_mask + file_name, canvas);
                        cout << "Mask saved successfully!" << endl;
                    }

                    imwrite(output_path + file_name, src);
                    cout << "Output image saved successfully!" << endl;
                } else {
                    cout << "Detection failed!" << endl;
                }
            } else {
                cout << "Model loading failed!" << endl;
            }
        } else {
            cout << file_name << endl;
            cout << "Image could not be read!" << endl;
        }
    }

    return 0;
}
Editor is loading...
Leave a Comment