Skip to content

Instantly share code, notes, and snippets.

@dotchang
Created November 28, 2019 04:09
Show Gist options
  • Save dotchang/21123d1293e9198022af8b9d6554e120 to your computer and use it in GitHub Desktop.
Save dotchang/21123d1293e9198022af8b9d6554e120 to your computer and use it in GitHub Desktop.
// Draw the predicted bounding box, colorize and show the mask on the image
void ThreadedMaskRCNN::drawBox(Mat& target, int classId, float conf, Rect box, Mat& objectMask)
{
//Draw a rectangle displaying the bounding box
rectangle(target, Point(box.x, box.y), Point(box.x + box.width, box.y + box.height), Scalar(255, 178, 50), 3);
//Get the label for the class name and its confidence
string label = format("%.2f", conf);
if (!classes.empty())
{
CV_Assert(classId < (int)classes.size());
label = classes[classId] + ":" + label;
}
std::cout << label << std::endl;
//Display the label at the top of the bounding box
int baseLine;
Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
box.y = max(box.y, labelSize.height);
rectangle(target, Point(box.x, box.y - round(1.5 * labelSize.height)), Point(box.x + round(1.5 * labelSize.width), box.y + baseLine), Scalar(255, 255, 255), FILLED);
putText(target, label, Point(box.x, box.y), FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 0, 0), 1);
Scalar color = colors[classId % colors.size()];
// Resize the mask, threshold, color and apply it on the image
resize(objectMask, objectMask, Size(box.width, box.height));
Mat mask = (objectMask > maskThreshold);
// bug fix
cv::Rect mask_roi;
mask_roi = cv::Rect(0, 0, box.width, box.height);
if (!(0 <= box.x && 0 <= box.width && box.x + box.width <= target.cols &&
0 <= box.y && 0 <= box.height && box.y + box.height <= target.rows))
{
if (box.x + box.width > target.cols) {
int last_width = box.width;
box.width = target.cols - box.x;
mask_roi.width -= last_width - box.width;
}
if (box.y + box.height > target.rows) {
int last_height = box.height;
box.height = target.rows - box.y;
mask_roi.height -= last_height - box.height;
}
std::cout << "roi error:" << box.x << ", " << box.y << ", " << box.width << ", " << box.height << ", "
<< target.cols << ", " << target.rows << std::endl;
}
Mat coloredRoi = (0.3 * color + 0.7 * target(box));
coloredRoi.convertTo(coloredRoi, CV_8UC3);
// Draw the contours on the image
vector<Mat> contours;
Mat hierarchy;
mask.convertTo(mask, CV_8U);
findContours(mask, contours, hierarchy, RETR_CCOMP, CHAIN_APPROX_SIMPLE);
drawContours(coloredRoi, contours, -1, color, 5, LINE_8, hierarchy, 100);
coloredRoi.copyTo(target(box), mask(mask_roi));
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment