void Processor::nms()

in models/01_YoloV5/01_Pytorch/processing_cpp/src/processor.cc [52:91]


void Processor::nms(std::tuple<vecf,vecf,vecf>& detections, float iou_threshold) {
    vecf& bboxes = xt::get<0>(detections);
    vecf& scores = xt::get<1>(detections);
    vecf& cids = xt::get<2>(detections);
    // extract each axis of the bboxes
    vecf x1 = xt::view(bboxes, xt::all(), 0);
    vecf y1 = xt::view(bboxes, xt::all(), 1);
    vecf x2 = xt::view(bboxes, xt::all(), 2);
    vecf y2 = xt::view(bboxes, xt::all(), 3);
    // compute the areas of each bbox
    vecf areas = (x2 - x1 + 1) * (y2 - y1 + 1);
    // sort the scores - reverse order
    vecf ord = xt::argsort( xt::flatten(scores) );
    vecf order(ord.shape());
    std::copy(ord.crbegin(), ord.crend(), order.begin());

    // compute intersections
    std::vector<size_t> keep;
    while (order.size() > 0 ) {
        int i = order[0];
        keep.push_back(i);
        vecs rest = xt::view(order, xt::range(1, xt::placeholders::_));
        vecf xx1 = xt::maximum(x1[i], xt::view(x1, xt::keep(rest)));
        vecf yy1 = xt::maximum(y1[i], xt::view(y1, xt::keep(rest)));
        vecf xx2 = xt::minimum(x2[i], xt::view(x2, xt::keep(rest)));
        vecf yy2 = xt::minimum(y2[i], xt::view(y2, xt::keep(rest)));

        vecf w = xt::maximum(0.0, xx2 - xx1 + 1);
        vecf h = xt::maximum(0.0, yy2 - yy1 + 1);
        vecf inter = w * h;
        vecf iou = inter / (areas[i] + xt::view(areas, xt::keep(rest)) - inter);
        vecs inds = xt::ravel_indices(xt::argwhere(iou <= iou_threshold), iou.shape());
        inds += 1;
        order = xt::view(order, xt::keep(inds));
    } // while
    // keep only the bboxes selected by nms
    bboxes = xt::view(bboxes, xt::keep(keep));
    scores = xt::view(scores, xt::keep(keep));
    cids = xt::view(cids, xt::keep(keep));
}