Merge pull request #1127 from arrybn:accuracy_scripts

modules/dnn/samples/torch_enet.cpp

@@ -26,9 +26,9 @@ const String keys =
         "{o_blob    || output blob's name. If empty, last blob's name in net is used}"
         ;
 
-std::vector<String> readClassNames(const char *filename);
 static void colorizeSegmentation(const Mat &score, Mat &segm,
-                                 Mat &legend, vector<String> &classNames);
+                                 Mat &legend, vector<String> &classNames, vector<Vec3b> &colors);
+static vector<Vec3b> readColors(const String &filename, vector<String>& classNames);
 
 int main(int argc, char **argv)
 {
@@ -52,43 +52,21 @@ int main(int argc, char **argv)
     String classNamesFile = parser.get<String>("c_names");
     String resultFile = parser.get<String>("result");
 
-    //! [Create the importer of TensorFlow model]
-    Ptr<dnn::Importer> importer;
-    try                                     //Try to import TensorFlow AlexNet model
-    {
-        importer = dnn::createTorchImporter(modelFile);
-    }
-    catch (const cv::Exception &err)        //Importer can throw errors, we will catch them
-    {
-        std::cerr << err.msg << std::endl;
-    }
-    //! [Create the importer of Caffe model]
-
-    if (!importer)
-    {
-        std::cerr << "Can't load network by using the mode file: " << std::endl;
-        std::cerr << modelFile << std::endl;
-        exit(-1);
-    }
-
-    //! [Initialize network]
-    dnn::Net net;
-    importer->populateNet(net);
-    importer.release();                     //We don't need importer anymore
-    //! [Initialize network]
+    //! [Read model and initialize network]
+    dnn::Net net = dnn::readNetFromTorch(modelFile);
 
     //! [Prepare blob]
-    Mat img = imread(imageFile, 1);
-
+    Mat img = imread(imageFile), input;
     if (img.empty())
     {
         std::cerr << "Can't read image from the file: " << imageFile << std::endl;
         exit(-1);
     }
 
-    Size inputImgSize(512, 512);
+    Size origSize = img.size();
+    Size inputImgSize = cv::Size(1024, 512);
 
-    if (inputImgSize != img.size())
+    if (inputImgSize != origSize)
         resize(img, img, inputImgSize);       //Resize image to input size
 
     Mat inputBlob = blobFromImage(img, 1./255, true);   //Convert Mat to image batch
@@ -135,20 +113,18 @@ int main(int argc, char **argv)
 
     if (parser.has("show"))
     {
-        size_t nclasses = result.size[1];
         std::vector<String> classNames;
+        vector<cv::Vec3b> colors;
         if(!classNamesFile.empty()) {
-            classNames = readClassNames(classNamesFile.c_str());
-            if (classNames.size() > nclasses)
-                classNames = std::vector<String>(classNames.begin() + classNames.size() - nclasses,
-                                                 classNames.end());
+            colors = readColors(classNamesFile, classNames);
         }
         Mat segm, legend;
-        colorizeSegmentation(result, segm, legend, classNames);
+        colorizeSegmentation(result, segm, legend, classNames, colors);
 
         Mat show;
-        addWeighted(img, 0.2, segm, 0.8, 0.0, show);
+        addWeighted(img, 0.1, segm, 0.9, 0.0, show);
 
+        cv::resize(show, show, origSize, 0, 0, cv::INTER_NEAREST);
         imshow("Result", show);
         if(classNames.size())
             imshow("Legend", legend);
@@ -158,44 +134,16 @@ int main(int argc, char **argv)
     return 0;
 } //main
 
-
-std::vector<String> readClassNames(const char *filename)
-{
-    std::vector<String> classNames;
-
-    std::ifstream fp(filename);
-    if (!fp.is_open())
-    {
-        std::cerr << "File with classes labels not found: " << filename << std::endl;
-        exit(-1);
-    }
-
-    std::string name;
-    while (!fp.eof())
-    {
-        std::getline(fp, name);
-        if (name.length())
-            classNames.push_back(name);
-    }
-
-    fp.close();
-    return classNames;
-}
-
-static void colorizeSegmentation(const Mat &score, Mat &segm, Mat &legend, vector<String> &classNames)
+static void colorizeSegmentation(const Mat &score, Mat &segm, Mat &legend, vector<String> &classNames, vector<Vec3b> &colors)
 {
     const int rows = score.size[2];
     const int cols = score.size[3];
     const int chns = score.size[1];
 
-    vector<Vec3i> colors;
-    RNG rng(12345678);
-
     cv::Mat maxCl(rows, cols, CV_8UC1);
     cv::Mat maxVal(rows, cols, CV_32FC1);
     for (int ch = 0; ch < chns; ch++)
     {
-        colors.push_back(Vec3i(rng.uniform(0, 256), rng.uniform(0, 256), rng.uniform(0, 256)));
         for (int row = 0; row < rows; row++)
         {
             const float *ptrScore = score.ptr<float>(0, ch, row);
@@ -235,3 +183,38 @@ static void colorizeSegmentation(const Mat &score, Mat &segm, Mat &legend, vecto
         }
     }
 }
+
+static vector<Vec3b> readColors(const String &filename, vector<String>& classNames)
+{
+    vector<cv::Vec3b> colors;
+    classNames.clear();
+
+    ifstream fp(filename.c_str());
+    if (!fp.is_open())
+    {
+        cerr << "File with colors not found: " << filename << endl;
+        exit(-1);
+    }
+
+    string line;
+    while (!fp.eof())
+    {
+        getline(fp, line);
+        if (line.length())
+        {
+            stringstream ss(line);
+
+            string name; ss >> name;
+            int temp;
+            cv::Vec3b color;
+            ss >> temp; color[0] = temp;
+            ss >> temp; color[1] = temp;
+            ss >> temp; color[2] = temp;
+            classNames.push_back(name);
+            colors.push_back(color);
+        }
+    }
+
+    fp.close();
+    return colors;
+}