/*****************************************************************************************

        This program reads in a generic directory tree of cropped sample images, and is

        currently only compatible with *nix systems due to file check macros etc.

        It quickly produces a PASCAL VOC (Visual Object Challenge) compatible output,

        and was built to feed OpenCV's sample code like  bagofwords_classification.cpp. 

        2012  Joel Mckay        

        [email protected]

        Disclaimer:

        THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS

        "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 

        LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS

        FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE 

        COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, 

        INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,

        BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 

        LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 

        CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT 

        LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY

        WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY

        OF SUCH DAMAGE. 

 *****************************************************************************************/

#include "global_headers.hpp"

/* How many of the positive set to train on, or comment out to test every file against the set */

//#define POS_TRAIN_SET_RATIO_OF_POSITIVE_SAMPLES 1

//#define NEG_TRAIN_SET_RATIO_OF_NEGAITIVE_SAMPLES 1

/* Notes:  additional negative training samples can be kept in another "negative" folder */

#define NEG_TRAIN_SET_GENERIC_PATH  "negative"

/* + all imported items will be listed as valid for their class */

/* + the file handles are freed with each scan as there may be a great number of classes */

/* + all names must be  [A-Za-Z0-9\.] to avoid VOC import format problems later */

/*****************************************************************************************/

/*

------------------------------------------------------------------------------------------------------------------------------

For classification and detection there are four sets for VOC:

------------------------------------------------------------------------------------------------------------------------------

"train": 

Training data

"val": 

Validation data (suggested). 

The validation data may be used as addi-tional training data (see below).

"trainval": 

The union of train and val.

"test": 

Test data. 

The test set is not provided in the development kit. 

It will be released in good time before the deadline for submission of results.

OpenCV BoG default file suffix used when the new VOC class loads with VocData vocData( vocPath, false ):

        m_train_set = "train";

        m_test_set = "val";

These file suffix used when the new VOC class loads with VocData vocData( vocPath, true ):

        m_train_set = "trainval";

        m_test_set = "test";

---------------------------------------------------------------

There are three ground truth labels:

---------------------------------------------------------------

-1: Negative: 

The image contains no objects of the class of interest. 

A classifier should give a 'negative' output.

1: Positive: 

The image contains at least one object of the class of interest.

A classifier should give a 'positive' output.

0: "Difficult": 

The image contains only objects of the class of interest marked as 'difficult'.

------------------------------------------------------------------------------------------------------------------------------

*/

/********************************************************************************************************/

//This function scans the class dir to push the specific file names one at a time to a vector list

void loadListFromDir( string dir , vector<string>* m_object_filenames)        

{

string filepath;

string filename;

string filebasename;

string fileext; 

int arrTemplatesCount;

DIR *dp;

struct dirent *dirp;

struct stat filestat; 

arrTemplatesCount=0;

dp = opendir( dir.c_str() );        //try to open the directory

if (dp == NULL)

{

    cout << "Error opening " << dir << endl;

    return;

}

while ((dirp = readdir( dp )) && (arrTemplatesCount < 100000)) //scan the dir

{

        filename=dirp->d_name;

        filepath = dir + "/" + filename; 

        // file invalid? we'll skip it... 

        if (stat( filepath.c_str(), &filestat )) continue;

        //is a file? 

        if (S_ISREG( filestat.st_mode ))  //is a real file?

        {

                fileext = filename.substr(filename.find_last_of(".") + 1);

                std::transform(fileext.begin(), fileext.end(),fileext.begin(), ::tolower); //str ::tolower   ::toupper 

                if(( fileext == "png") || ( fileext == "jpg") || ( fileext == "bmp"))        //image?

                {

                        filebasename = filename.substr(0, filename.find_first_of(".") );

                        #if defined(DEBUG_MODE)

                                cout << "Loaded: " << filebasename << "  " << filepath.c_str() << endl;

                        #endif

                        (*m_object_filenames).push_back(filepath);

                }

        }        

}

closedir( dp );        //done loading list

}

/********************************************************************************************************/

//This function scans the class dir to push the specific DIR names one at a time to a vector list

void loadClassListFromDir( string dir , vector<string>* m_object_classes,  vector<string>* m_object_classes_paths)        

{

string filepath;

string filename;

string filebasename;

string fileext; 

int arrTemplatesCount;

DIR *dp;

struct dirent *dirp;

struct stat filestat; 

arrTemplatesCount=0;

dp = opendir( dir.c_str() );        //try to open the directory

if (dp == NULL)

{

    cout << "Error opening " << dir << endl;

    return;

}

while ((dirp = readdir( dp )) && (arrTemplatesCount < 100000)) //scan the dir

{

        filename=dirp->d_name;

        filepath = dir + "/" + filename; 

        // file invalid? we'll skip it... 

        if (stat( filepath.c_str(), &filestat )) continue;

        //is a dir class name? 

        if (S_ISDIR( filestat.st_mode ))  //is a real dir?

        {

                fileext = filename;

                std::transform(fileext.begin(), fileext.end(),fileext.begin(), ::tolower); //str ::tolower   ::toupper 

                if(( fileext != ".") && ( fileext != ".."))        //real dir?

                { 

                        (*m_object_classes).push_back(fileext);

                        (*m_object_classes_paths).push_back(filepath);

                }

        }        

}

closedir( dp );        //done loading list

} 

/********************************************************************************************************/

//This function exports a specific image if it passes QC 

//note: filename is the name without extension

//format string to reduce stack hammering

const string FASTXMLOUTPUTFORMAT="<annotation>\n        <folder>VOC2010</folder>\n        <filename>%s.jpg</filename>\n        <source>\n                <database>The VOC2007 Database</database>\n                <annotation>PASCAL VOC2007</annotation>\n                <image>flickr</image>\n        </source>\n        <size>\n                <width>%i</width>\n                <height>%i</height>\n                <depth>%i</depth>\n        </size>\n        <segmented>1</segmented>\n        <object>\n                <name>%s</name>\n                <pose>Frontal</pose>\n                <truncated>0</truncated>\n                <difficult>0</difficult>\n                <bndbox>\n                        <xmin>%i</xmin>\n                        <ymin>%i</ymin>\n                        <xmax>%i</xmax>\n                        <ymax>%i</ymax>\n                </bndbox>\n        </object>\n</annotation>\n";

int  createImageXMLfile( Mat* img, string* image_path, string* annotation_path, string* filename, string* classname)

{

//create file contents

Size size = (*img).size();

int width=size.width;

int height=size.height;

int colorChans=(*img).channels();

int colordepth=(*img).depth();        //CV_8U ?

int boxXmin=0;

int boxYmin=0;

int boxXmax=size.width;                //assume item is pre-cropped

int boxYmax=size.height;

        //check if BGR or Gray

        if(!((colorChans==3) || (colorChans == 1)))

        {

                cout << "\nError: Bad color channel in: " << (*filename) << endl;

                return -1;

        }

        //check if image too small

        if((width < 40) || (height < 40))

        {

                cout << "\nError: Image too small: " << (*filename) << endl;

                return -1;

        }

        //create new copy of image data

        string jpgFileName= (*image_path) + "/" + (*filename) + ".jpg";

        if(imwrite(jpgFileName, (*img)) == false)

        {

                cout << "\nError: Image write fail: " << (*filename) << endl;

                return -1;  

        }

        //Dump cropped image info to xml file

        FILE * filePointer;

        string xmlFileName= (*annotation_path) + "/" + (*filename) + ".xml"; 

        filePointer = fopen(xmlFileName.c_str(),"w");        //imgfilename.xml

        if( fprintf(filePointer, FASTXMLOUTPUTFORMAT.c_str() ,(*filename).c_str(), width, height, colorChans,(*classname).c_str(),boxXmin,boxYmin,boxXmax,boxYmax) < 0 )

        {

                    cout << "\nError: Could not write xml file: " << xmlFileName << endl;

                    return -1;

        }

        fclose(filePointer);

   cout << "Wrote file: " << xmlFileName << "\r";

   return 0;

}

/********************************************************************************************************/

int main(int argc, char** argv)

{

        if( argc != 3 )

        { 

        cout << "\nUsage" << endl

                << "./build_classification  </path/to/all/cropped/image/classdirs> </path/to/create/new/VOC2010/>" << endl;        

        exit(-1);

        }

        //Mimic minimal VOC2010 dir structure

        string r_vocPath = argv[2];

        if( mkdir(r_vocPath.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH)  != 0 )

        {

                cout << "\nError: Could not create the path: \n" << r_vocPath << endl;        

        }

        string r_annotation_path = r_vocPath + "/Annotations";

        if( mkdir(r_annotation_path.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH)  != 0 )

        {

                cout << "\nError: Could not create the path: \n" << r_annotation_path << endl;        

                exit(-1);

        }

        string r_image_path = r_vocPath + "/JPEGImages";

        if( mkdir(r_image_path.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH)  != 0 )

        {

                cout << "\nCould not create JPEGImages path: \n" << r_image_path << endl;        

                exit(-1);

        }

        string r_imageset_path = r_vocPath + "/ImageSets";

        if( mkdir(r_imageset_path.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH)  != 0 )

        {

                cout << "\nCould not create ImageSets path: \n" << r_imageset_path << endl;        

                exit(-1);

        }

        string r_class_imageset_path = r_vocPath + "/ImageSets/Main";

        if( mkdir(r_class_imageset_path.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH)  != 0 )

        {

                cout << "\nCould not create ImageSets/Main path: \n" << r_class_imageset_path << endl;        

                exit(-1);

        }

        // initialize main classification/detection template paths

        string template_imageset_path = r_vocPath + "/ImageSets/Main/%s.txt";

        string template_class_imageset_path = r_vocPath + "/ImageSets/Main/%s_%s.txt";

        char trainFileName[4096]="\0";

        char valFileName[4096]="\0"; 

        char trainvalFileName[4096]="\0";

        char testFileName[4096]="\0";

        //Scan all of the class dirs

        string classesSrcBaseDir=argv[1];  // classdirs 

        vector<string> my_classes;  

        vector<string> my_classes_dir;   

        vector<vector<string> > my_dataset_files;  

        loadClassListFromDir(classesSrcBaseDir, &my_classes, &my_classes_dir);

        unsigned int fileCounter=0;

        //check all dirs/classnames in path for positive samples 

        for( size_t classIdx = 0; (classIdx < my_classes.size()); ++classIdx )

        {

                /* first check for file list for the current class has been saved to file */

                vector<string> my_files;  

                vector<string> img_data_names;  

                loadListFromDir(my_classes_dir[classIdx], &my_files);

                //parse out each dir as a new class name

                string dirclassname=my_classes[classIdx];//new class to read?

                cout << "\n---------------------------- Create pos set: " << dirclassname << "-----------------------"<<endl;

                //create file index for each scanned dir class

                sprintf(trainFileName, template_class_imageset_path.c_str(), dirclassname.c_str(), "train");

                sprintf(valFileName, template_class_imageset_path.c_str(), dirclassname.c_str(), "val");

                sprintf(trainvalFileName, template_class_imageset_path.c_str(), dirclassname.c_str(), "trainval");

                sprintf(testFileName, template_class_imageset_path.c_str(), dirclassname.c_str(), "test");

                FILE * trainFilePointer;

                FILE * trainvalFilePointer;

                FILE * valFilePointer;

                FILE * testFilePointer;

                trainFilePointer = fopen(trainFileName,"a");

                trainvalFilePointer = fopen(trainvalFileName,"a");

                valFilePointer = fopen(valFileName,"a");

                testFilePointer = fopen(testFileName,"a");

                //check all files in specific class dir

                for(size_t fileIdx = 0; (fileIdx < my_files.size()); ++fileIdx)

                { 

                        Mat img = imread(my_files[fileIdx] .c_str());                //read each file into the buffer

                        //create new file name

                        std::ostringstream namerender;

                        namerender << "0000_" << fileCounter;                        //bogus release prefix to prevent overlap

                        string imgname =  namerender.str();

                        #if defined(DEBUG_MODE)

                        cout << "\nREAD: " << my_files[fileIdx] << endl

                                << "\nr_image_path:" << r_image_path << endl

                                << "\nr_annotation_path:" << r_annotation_path <<endl 

                                << "\nimgname:" << imgname <<endl

                                << "\ndirclassname:" << dirclassname << endl;

                        #endif

                        //if item imported write to the log list format

                        if( createImageXMLfile( &img, &r_image_path, &r_annotation_path, &imgname, &dirclassname) == 0 )

                        {

                                        //Sample valid image name to index files only if not in negative set???

                                        if( imgname.compare( NEG_TRAIN_SET_GENERIC_PATH ) != 0)

                                        {

                                        #if defined(POS_TRAIN_SET_RATIO_OF_POSITIVE_SAMPLES)

                                        if((fileCounter%POS_TRAIN_SET_RATIO_OF_POSITIVE_SAMPLES)==0)

                                        #endif

                                        {

                                                if( fprintf(trainFilePointer, "%s  1\n" ,imgname.c_str() ) < 0 )

                                                {

                                                    cout << "\nError: Could not write file: " << trainFilePointer << endl;

                                                    exit(-1);

                                                }

                                        }

                                        #if defined(POS_TRAIN_SET_RATIO_OF_POSITIVE_SAMPLES)

                                        else

                                        #endif

                                        {

                                                if( fprintf(valFilePointer, "%s  1\n" ,imgname.c_str() ) < 0 )

                                                {

                                                    cout << "\nError: Could not write file: " << trainFilePointer << endl;

                                                    exit(-1);

                                                }

                                        }

                                        if( fprintf(trainvalFilePointer, "%s  1\n" ,imgname.c_str() ) < 0 )

                                        {

                                                    cout << "\nError: Could not write file: " << trainvalFilePointer << endl;

                                                    exit(-1);

                                        }

                                        }

                                        //create data name list

                                        img_data_names.push_back(imgname);

                                        //show new file name for user

                                        cout << imgname   <<".jpg OK!                                                \r";

                                        fileCounter++; 

                        } 

                } 

                //store output marker for negative training data later

                my_dataset_files.push_back(img_data_names);

                fclose(trainFilePointer);

                fclose(trainvalFilePointer);

                fclose(valFilePointer);

                fclose(testFilePointer);

        }

        fileCounter=0;         //reset counter

        //append all classnames negative samples to disjoint sets

        for( size_t classIdx = 0; (classIdx < my_classes.size()); ++classIdx )

        {

                //parse out each dir as a new class name

                string dirclassname=my_classes[classIdx];//new class to read?

                cout << "\n---------------------------- Create neg set: " << dirclassname << " -----------------------"<<endl;

                //Re-open file index for each scanned dir class

                sprintf(trainFileName, template_class_imageset_path.c_str(), dirclassname.c_str(), "train");

                sprintf(trainvalFileName, template_class_imageset_path.c_str(), dirclassname.c_str(), "trainval");

                sprintf(valFileName, template_class_imageset_path.c_str(), dirclassname.c_str(), "val");

                sprintf(testFileName, template_class_imageset_path.c_str(), dirclassname.c_str(), "test");

                FILE * trainFilePointer;

                FILE * trainvalFilePointer;

                FILE * valFilePointer;

                FILE * testFilePointer;

                trainFilePointer = fopen(trainFileName,"a");

                trainvalFilePointer = fopen(trainvalFileName,"a");

                valFilePointer = fopen(valFileName,"a");

                testFilePointer = fopen(testFileName,"a"); 

                //append all file names if in disjoint class

                for( size_t classIter = 0; (classIter < my_classes.size()); ++classIter )

                {

                        if(classIdx != classIter)

                        {

                                for(size_t fileIdx = 0; (fileIdx < my_dataset_files[classIter].size()) ; ++fileIdx)

                                {   

                                        //if item imported write to the log list format

                                        //Sample valid neg index files

                                                #if defined(NEG_TRAIN_SET_RATIO_OF_NEGAITIVE_SAMPLES)

                                                if((fileCounter%NEG_TRAIN_SET_RATIO_OF_NEGAITIVE_SAMPLES)==0)

                                                #endif

                                                {

                                                        if( fprintf(trainFilePointer, "%s -1\n" ,my_dataset_files[classIter][fileIdx].c_str() ) < 0 )

                                                        {

                                                                cout << "\nError: Could not write file: " << trainFilePointer << endl;

                                                                exit(-1);

                                                        }

                                                }

                                                #if defined(NEG_TRAIN_SET_RATIO_OF_NEGAITIVE_SAMPLES)

                                                else

                                                #endif

                                                {

                                                        if( fprintf(valFilePointer, "%s -1\n" ,my_dataset_files[classIter][fileIdx].c_str() ) < 0 )

                                                        {

                                                            cout << "\nError: Could not write file: " << trainFilePointer << endl;

                                                            exit(-1);

                                                        }

                                                 }

                                                if( fprintf(trainvalFilePointer, "%s -1\n" ,my_dataset_files[classIter][fileIdx].c_str() ) < 0 )

                                                {

                                                        cout << "\nError: Could not write file: " << trainvalFilePointer << endl;

                                                        exit(-1);

                                                }

                                                fileCounter++; 

                                } 

                                cout << ".";         //progress bar nibbles

                        }

                } 

                fclose(trainFilePointer);

                fclose(trainvalFilePointer);

                fclose(valFilePointer);

                fclose(testFilePointer);

        }

        cout << "\n Done!" << endl;

        return 0;

}