File size: 4,859 Bytes
1387a61 | 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 | #include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/imgcodecs.hpp>
#include <iostream>
#include <string>
#include <vector>
// The exported graph is frozen at 384x384 and always takes two images: the
// picture and its horizontal mirror.
static const int INPUT_SIZE = 384;
static const cv::Scalar MEAN(104.00698793, 116.66876762, 122.67891434);
// The 20 LIP classes, in channel order.
static const char* CLASSES[] = {
"Background", "Hat", "Hair", "Glove", "Sunglasses", "UpperClothes",
"Dress", "Coat", "Socks", "Pants", "Jumpsuits", "Scarf", "Skirt",
"Face", "LeftArm", "RightArm", "LeftLeg", "RightLeg", "LeftShoe",
"RightShoe"
};
static const int NUM_CLASSES = 20;
// LIP palette, RGB.
static const unsigned char COLORS[NUM_CLASSES][3] = {
{0, 0, 0}, {128, 0, 0}, {255, 0, 0}, {0, 85, 0}, {170, 0, 51},
{255, 85, 0}, {0, 0, 85}, {0, 119, 221}, {85, 85, 0}, {0, 85, 85},
{85, 51, 0}, {52, 86, 128}, {0, 128, 0}, {0, 0, 255}, {51, 170, 221},
{0, 255, 255}, {85, 255, 170}, {170, 255, 85}, {255, 255, 0}, {255, 170, 0}
};
static std::string argVal(int argc, char** argv, const std::string& key, const std::string& def)
{
for (int i = 1; i + 1 < argc; ++i)
if (key == argv[i]) return argv[i + 1];
return def;
}
// Mirroring an image swaps the left and right body parts, so the channels of
// the mirrored prediction have to be swapped back before the two are averaged.
static int mirrorChannel(int c)
{
switch (c)
{
case 14: return 15; // LeftArm <-> RightArm
case 15: return 14;
case 16: return 17; // LeftLeg <-> RightLeg
case 17: return 16;
case 18: return 19; // LeftShoe <-> RightShoe
case 19: return 18;
default: return c;
}
}
// Build the [2, 3, 384, 384] blob: the image and its mirror, mean-subtracted.
static cv::Mat preprocess(const cv::Mat& image)
{
cv::Mat resized, mirror;
cv::resize(image, resized, cv::Size(INPUT_SIZE, INPUT_SIZE));
cv::flip(resized, mirror, 1);
std::vector<cv::Mat> images;
images.push_back(resized);
images.push_back(mirror);
return cv::dnn::blobFromImages(images, 1.0, cv::Size(), MEAN);
}
// Average the two predictions and turn them into a label map of `size`.
// `out` is [2, 20, 384, 384].
static cv::Mat postprocess(const cv::Mat& out, const cv::Size& size)
{
cv::Mat labels(size, CV_8U, cv::Scalar(0));
cv::Mat best(size, CV_32F, cv::Scalar(-FLT_MAX));
for (int c = 0; c < NUM_CLASSES; ++c)
{
cv::Mat direct(INPUT_SIZE, INPUT_SIZE, CV_32F,
const_cast<float*>(out.ptr<float>(0, c)));
cv::Mat mirrored(INPUT_SIZE, INPUT_SIZE, CV_32F,
const_cast<float*>(out.ptr<float>(1, mirrorChannel(c))));
cv::Mat unmirrored;
cv::flip(mirrored, unmirrored, 1);
cv::Mat score = 0.5f * (direct + unmirrored);
cv::resize(score, score, size, 0, 0, cv::INTER_LINEAR);
cv::Mat better = score > best;
score.copyTo(best, better);
labels.setTo(c, better);
}
return labels;
}
// Paint a label map with the LIP palette, as a BGR image.
static cv::Mat colorize(const cv::Mat& labels)
{
cv::Mat lut(1, NUM_CLASSES, CV_8UC3);
for (int i = 0; i < NUM_CLASSES; ++i)
lut.at<cv::Vec3b>(0, i) = cv::Vec3b(COLORS[i][2], COLORS[i][1], COLORS[i][0]);
cv::Mat segmentation(labels.size(), CV_8UC3);
for (int y = 0; y < labels.rows; ++y)
{
const unsigned char* src = labels.ptr<unsigned char>(y);
cv::Vec3b* dst = segmentation.ptr<cv::Vec3b>(y);
for (int x = 0; x < labels.cols; ++x)
dst[x] = lut.at<cv::Vec3b>(0, src[x]);
}
return segmentation;
}
int main(int argc, char** argv)
{
std::string model = argVal(argc, argv, "--model", "human_parsing_jppnet_2026sep.onnx");
std::string image = argVal(argc, argv, "--image", "example_outputs/input_image.png");
std::string output = argVal(argc, argv, "--output", "example_outputs/output_image.png");
cv::Mat img = cv::imread(image);
if (img.empty())
{
std::cerr << "could not read image: " << image << std::endl;
return 1;
}
cv::dnn::Net net = cv::dnn::readNetFromONNX(model);
net.setInput(preprocess(img));
cv::Mat out = net.forward();
cv::Mat labels = postprocess(out, img.size());
cv::Mat segmentation = colorize(labels);
std::cout << "human_parsing_jppnet input " << img.cols << "x" << img.rows << std::endl;
for (int c = 0; c < NUM_CLASSES; ++c)
{
int pixels = cv::countNonZero(labels == c);
if (pixels > 0)
std::cout << " " << CLASSES[c] << " " << pixels << " px" << std::endl;
}
cv::imwrite(output, segmentation);
std::cout << "wrote " << output << std::endl;
return 0;
}
|