main.cpp 3.89 KB
Newer Older
1

2
3
#include "image_dataset_metadata.h"

4
5
#include <iostream>
#include <fstream>
6
#include <string>
7
8

#include <dlib/cmd_line_parser.h>
9
#include <dlib/dir_nav.h>
10
11


12
13
14
15

using namespace std;
using namespace dlib;

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
std::string strip_path (
    const std::string& str,
    const std::string& prefix
)
{
    unsigned long i;
    for (i = 0; i < str.size() && i < prefix.size(); ++i)
    {
        if (str[i] != prefix[i]) 
            return str;
    }

    if (i < str.size() && (str[i] == '/' || str[i] == '\\'))
        ++i;

    return str.substr(i);
}

// ----------------------------------------------------------------------------------------

void make_empty_file (
    const std::string& filename
)
{
    ofstream fout(filename.c_str());
    if (!fout)
        throw dlib::error("ERROR: Unable to open " + filename + " for writing.");
}

// ----------------------------------------------------------------------------------------


50
51
52
53
54
55
56
57
58
59
int main(int argc, char** argv)
{
    try
    {
        typedef dlib::cmd_line_parser<char>::check_1a_c parser_type;

        parser_type parser;

        parser.add_option("h","Displays this information.");
        parser.add_option("c","Create an XML file named <arg> listing a set of images.",1);
60
        parser.add_option("r","Search directories recursively for images.");
61
62
63

        parser.parse(argc, argv);

64
        const char* singles[] = {"h","c","r"};
65
        parser.check_one_time_options(singles);
66
        parser.check_sub_option("c", "r");
67
68
69

        if (parser.option("h"))
        {
Davis King's avatar
Davis King committed
70
            cout << "Usage: imglab [options] <image files/directories or XML file list>\n";
71
            parser.print_options(cout);
Davis King's avatar
Davis King committed
72
            cout << endl << endl;
73
74
75
76
77
            return EXIT_SUCCESS;
        }

        if (parser.option("c"))
        {
Davis King's avatar
Davis King committed
78
            using namespace dlib::image_dataset_metadata;
79

80
81
82
83
84
85
            const std::string filename = parser.option("c").argument();
            // make sure the file exists so we can use the get_parent_directory() command to
            // figure out it's parent directory.
            make_empty_file(filename);
            const std::string parent_dir = get_parent_directory(file(filename)).full_name();

86
87
88
89
            unsigned long depth = 0;
            if (parser.option("r"))
                depth = 30;

Davis King's avatar
Davis King committed
90
91
92
            dataset meta;
            meta.name = "imglab dataset";
            meta.comment = "Created by imglab tool.";
93
94
            for (unsigned long i = 0; i < parser.number_of_arguments(); ++i)
            {
95
96
97
                try
                {
                    const string temp = strip_path(file(parser[i]).full_name(), parent_dir);
Davis King's avatar
Davis King committed
98
                    meta.images.push_back(image(temp));
99
100
101
102
                }
                catch (dlib::file::file_not_found&)
                {
                    // then parser[i] should be a directory
103

104
105
106
107
                    std::vector<file> files = get_files_in_directory_tree(parser[i], 
                                                                          match_endings(".png .PNG .jpeg .JPEG .jpg .JPG .bmp .BMP .dng .DNG"),
                                                                          depth);
                    sort(files.begin(), files.end());
108

109
110
                    for (unsigned long j = 0; j < files.size(); ++j)
                    {
Davis King's avatar
Davis King committed
111
                        meta.images.push_back(image(strip_path(files[j].full_name(), parent_dir)));
112
                    }
113
                }
114
115
            }

Davis King's avatar
Davis King committed
116
            save_image_dataset_metadata(meta, filename);
117
118
119
120

            return EXIT_SUCCESS;
        }

121
122
        if (parser.number_of_arguments() == 1)
        {
Davis King's avatar
Davis King committed
123
124
125
            dlib::image_dataset_metadata::dataset meta;
            load_image_dataset_metadata(meta, parser[0]);
            save_image_dataset_metadata(meta, "out.xml");
126
        }
127
128
129
130
131
132
133
134
    }
    catch (exception& e)
    {
        cout << e.what() << endl;
        return EXIT_FAILURE;
    }
}