DNN based super-resolution phase 2

modules/README.md

@@ -22,7 +22,9 @@ $ cmake -D OPENCV_EXTRA_MODULES_PATH=<opencv_contrib>/modules -D BUILD_opencv_<r
 
 - **datasets**: Datasets Reader -- Code for reading existing computer vision databases and samples of using the readers to train, test and run using that dataset's data.
 
-- **dnn_objdetect**: Object Detection using CNNs -- Implements compact CNN Model for object detection. Trained using Caffe but uses opencv_dnn modeule.
+- **dnn_objdetect**: Object Detection using CNNs -- Implements compact CNN Model for object detection. Trained using Caffe but uses opencv_dnn module.
+
+- **dnn_superres**: Superresolution using CNNs -- Contains four trained convolutional neural networks to upscale images.
 
 - **dnns_easily_fooled**: Subvert DNNs -- This code can use the activations in a network to fool the networks into recognizing something else.
 

modules/dnn_superres/CMakeLists.txt

@@ -0,0 +1,5 @@
+set(the_description "Super Resolution using CNNs")
+
+ocv_define_module(dnn_superres opencv_core opencv_imgproc opencv_dnn
+    OPTIONAL opencv_highgui opencv_imgcodecs  # samples
+)

modules/dnn_superres/README.md

@@ -0,0 +1,61 @@
+# Super Resolution using Convolutional Neural Networks
+
+This module contains several learning-based algorithms for upscaling an image.
+
+## Dependencies
+-
+-
+
+## Building this module
+
+Run the following command to build this module:
+
+```make
+cmake -DOPENCV_EXTRA_MODULES_PATH=<opencv_contrib>/modules -Dopencv_dnn_superres=ON <opencv_source_dir>
+```
+
+## Models
+
+There are four models which are trained. (Not yet implemented!!)
+
+#### EDSR
+
+- Size of the model:
+- This model was trained for <> iterations with a batch size of <>
+- Link to model:
+- Advantage:
+- Disadvantage:
+- Speed:
+
+#### ESPCN
+
+Trained models can be downloaded from [here](https://github.com/fannymonori/TF-ESPCN/tree/master/export).
+
+- Size of the model: ~100kb
+- This model was trained for 100 iterations with a batch size of 32
+- Link to implementation code: https://github.com/fannymonori/TF-ESPCN
+- x2, x3, x4 trained models available
+- Advantage: It is tiny, and fast, and still perform well.
+- Disadvantage: Perform worse visually than newer, more robust models.
+- Speed:
+
+#### FSRCNN
+
+- Size of the model:
+- This model was trained for <> iterations with a batch size of <>
+- Link to model:
+- Advantage:
+- Disadvantage:
+- Speed:
+
+#### LapSRN
+
+Trained models can be downloaded from [here](https://github.com/fannymonori/TF-LapSRN/tree/master/export).
+
+- Size of the model: between 1-5Mb
+- This model was trained for ~50 iterations with a batch size of 32
+- Link to implementation code: https://github.com/fannymonori/TF-LAPSRN
+- x2, x4, x8 trained models available
+- Advantage: The model can do multi-scale super-resolution with one forward pass. It can now support 2x, 4x, 8x, and [2x, 4x] and [2x, 4x, 8x] super-resolution.
+- Disadvantage: It is a slower model.
+- Speed

modules/dnn_superres/include/opencv2/dnn_superres.hpp

@@ -0,0 +1,158 @@
+// This file is part of OpenCV project.
+// It is subject to the license terms in the LICENSE file found in the top-level directory
+// of this distribution and at http://opencv.org/license.html.
+
+#ifndef _OPENCV_DNN_SUPERRES_DNNSUPERRESIMPL_HPP_
+#define _OPENCV_DNN_SUPERRES_DNNSUPERRESIMPL_HPP_
+
+#include <opencv2/highgui.hpp>
+#include <opencv2/imgproc.hpp>
+#include "opencv2/dnn.hpp"
+
+/** @defgroup dnn_superres DNN used for super resolution
+
+This module contains functionality for upscaling an image via convolutional neural networks.
+The following four models are implemented:
+
+- EDSR
+- ESPCN
+- FSRCNN
+- LapSRN
+
+There is also functionality for simply upscaling by bilinear or bicubic interpolation.
+
+*/
+
+namespace cv
+{
+namespace dnn
+{
+namespace dnn_superres
+{
+    //! @addtogroup dnn_superres
+    //! @{
+
+    /** @brief A class to upscale images via convolutional neural networks.
+    The following four models are implemented:
+
+    - edsr
+    - espcn
+    - fsrcnn
+    - lapsrn
+     */
+    class CV_EXPORTS DnnSuperResImpl
+    {
+        private:
+
+            /** @brief Net which holds the desired neural network
+             */
+            Net net;
+
+            std::string alg; //algorithm
+
+            int sc; //scale factor
+
+            /// @private
+            static int layer_loaded;
+
+            void registerLayers();
+
+            void preprocess(const Mat inpImg, Mat &outpImg);
+
+            void reconstruct_YCrCb(const Mat inpImg, const Mat origImg, Mat &outpImg, int scale);
+
+            void reconstruct_YCrCb(const Mat inpImg, const Mat origImg, Mat &outpImg);
+
+            void preprocess_YCrCb(const Mat inpImg, Mat &outpImg);
+
+        public:
+
+            /** @brief Empty constructor
+             */
+            DnnSuperResImpl();
+
+            /** @brief Constructor which immediately sets the desired model
+            @param algo String containing one of the desired models:
+                - __edsr__
+                - __espcn__
+                - __fsrcnn__
+                - __lapsrn__
+            @param scale Integer specifying the upscale factor
+             */
+            DnnSuperResImpl(std::string algo, int scale);
+
+            /** @brief Read the model from the given path
+            @param path Path to the model file.
+            */
+            void readModel(std::string path);
+
+            /** @brief Read the model from the given path
+            @param weights Path to the model weights file.
+            @param definition Path to the model definition file.
+            */
+            void readModel(std::string weights, std::string definition);
+
+            /** @brief Set desired model
+            @param algo String containing one of the desired models:
+                - __edsr__
+                - __espcn__
+                - __fsrcnn__
+                - __lapsrn__
+            @param scale Integer specifying the upscale factor
+             */
+            void setModel(std::string algo, int scale);
+
+            /** @brief Upsample via neural network
+            @param img Image to upscale
+            @param img_new Destination upscaled image
+             */
+            void upsample(Mat img, Mat &img_new);
+
+            /** @brief Upsample via neural network of multiple outputs
+            @param img Image to upscale
+            @param imgs_new Destination upscaled images
+            @param scale_factors Scaling factors of the output nodes
+            @param node_names Names of the output nodes in the neural network
+            */
+            void upsample_multioutput(Mat img, std::vector<Mat> &imgs_new, std::vector<int> scale_factors, std::vector<String> node_names);
+
+            /** @brief Returns the scale factor of the model:
+            @return Current scale factor.
+            */
+            int getScale();
+
+            /** @brief Returns the scale factor of the model:
+            @return Current algorithm.
+            */
+            std::string getAlgorithm();
+
+            private:
+            /** @brief Class for importing DepthToSpace layer from the ESPCN model
+            */
+            class DepthToSpace CV_FINAL : public cv::dnn::Layer
+            {
+                public:
+
+                /// @private
+                DepthToSpace(const cv::dnn::LayerParams &params);
+
+                /// @private
+                static cv::Ptr<cv::dnn::Layer> create(cv::dnn::LayerParams& params);
+
+                /// @private
+                virtual bool getMemoryShapes(const std::vector<std::vector<int> > &inputs,
+                                             const int,
+                                             std::vector<std::vector<int> > &outputs,
+                                             std::vector<std::vector<int> > &) const CV_OVERRIDE;
+
+                /// @private
+                virtual void forward(cv::InputArrayOfArrays inputs_arr,
+                                     cv::OutputArrayOfArrays outputs_arr,
+                                     cv::OutputArrayOfArrays) CV_OVERRIDE;
+            };
+    };
+    //! @}
+}
+}
+}
+#endif

modules/dnn_superres/samples/dnn_superres.cpp

@@ -0,0 +1,76 @@
+#include <iostream>
+
+#include <opencv2/dnn_superres.hpp>
+
+using namespace std;
+using namespace cv;
+using namespace dnn;
+using namespace dnn_superres;
+
+int main(int argc, char *argv[])
+{
+    // Check for valid command line arguments, print usage
+    // if insufficient arguments were given.
+    if (argc < 4) {
+        cout << "usage:   Arg 1: image     | Path to image" << endl;
+        cout << "\t Arg 2: algorithm | bilinear, bicubic, edsr, espcn, fsrcnn or lapsrn" << endl;
+        cout << "\t Arg 3: scale     | 2, 3 or 4 \n";
+        cout << "\t Arg 4: path to model file \n";
+        return -1;
+    }
+
+    string img_path = string(argv[1]);
+    string algorithm = string(argv[2]);
+    int scale = atoi(argv[3]);
+    string path = "";
+
+    if( argc > 4)
+        path = string(argv[4]);
+
+    // Load the image
+    Mat img = cv::imread(img_path);
+    Mat original_img(img);
+    if (img.empty())
+    {
+        std::cerr << "Couldn't load image: " << img << "\n";
+        return -2;
+    }
+
+    //Make dnn super resolution instance
+    DnnSuperResImpl sr;
+
+    Mat img_new;
+
+    if(algorithm == "bilinear"){
+        resize(img, img_new, Size(), scale, scale, 2);
+    }
+    else if(algorithm == "bicubic")
+    {
+        resize(img, img_new, Size(), scale, scale, 3);
+    }
+    else if(algorithm == "espcn" || algorithm == "lapsrn")
+    {
+        sr.readModel(path);
+        sr.setModel(algorithm, scale);
+        sr.upsample(img, img_new);
+    }
+    else{ //one of the neural networks
+        sr.setModel(algorithm, scale);
+        sr.upsample(img, img_new);
+    }
+
+    if (img_new.empty())
+    {
+        std::cerr << "Upsampling failed. \n";
+        return -3;
+    }
+    cout << "Upsampling succeeded. \n";
+
+    // Display image
+    cv::namedWindow("Initial Image", WINDOW_AUTOSIZE);
+    cv::imshow("Initial Image", img_new);
+    //cv::imwrite("./saved.jpg", img_new);
+    cv::waitKey(0);
+
+    return 0;
+}

modules/dnn_superres/samples/dnn_superres_multioutput.cpp

@@ -0,0 +1,68 @@
+#include <iostream>
+#include <sstream>
+#include <opencv2/dnn_superres.hpp>
+
+using namespace std;
+using namespace cv;
+using namespace dnn;
+using namespace dnn_superres;
+
+int main(int argc, char *argv[])
+{
+    // Check for valid command line arguments, print usage
+    // if insufficient arguments were given.
+    if (argc < 4) {
+        cout << "usage:   Arg 1: image     | Path to image" << endl;
+        cout << "\t Arg 2: scales in a format of 2,4,8\n";
+        cout << "\t Arg 3: output node names in a format of nchw_output_0,nchw_output_1\n";
+        cout << "\t Arg 4: path to model file \n";
+        return -1;
+    }
+
+    string img_path = string(argv[1]);
+    string scales_str = string(argv[2]);
+    string output_names_str = string(argv[3]);
+    std::string path = string(argv[4]);
+
+    std::stringstream ss(scales_str);
+    std::vector<int> scales;
+    std::string token;
+    char delim = ',';
+    while (std::getline(ss, token, delim)) {
+        scales.push_back(atoi(token.c_str()));
+    }
+
+    ss = std::stringstream(output_names_str);
+    std::vector<String> node_names;
+    while (std::getline(ss, token, delim)) {
+        node_names.push_back(token);
+    }
+
+    // Load the image
+    Mat img = cv::imread(img_path);
+    Mat original_img(img);
+    if (img.empty())
+    {
+        std::cerr << "Couldn't load image: " << img << "\n";
+        return -2;
+    }
+
+    //Make dnn super resolution instance
+    DnnSuperResImpl sr;
+    int scale = *max_element(scales.begin(), scales.end());
+    std::vector<Mat> outputs;
+    sr.readModel(path);
+    sr.setModel("lapsrn", scale);
+
+    sr.upsample_multioutput(img, outputs, scales, node_names);
+
+    for(unsigned int i = 0; i < outputs.size(); i++)
+    {
+        cv::namedWindow("Upsampled image", WINDOW_AUTOSIZE);
+        cv::imshow("Upsampled image", outputs[i]);
+        //cv::imwrite("./saved.jpg", img_new);
+        cv::waitKey(0);
+    }
+
+    return 0;
+}