main.cpp 6.77 KB
Newer Older
1

2
#include "image_dataset_metadata.h"
3
#include "metadata_editor.h"
4
#include "convert_pascal_voc.h"
5

6
7
#include <iostream>
#include <fstream>
8
#include <string>
9
#include <set>
10

11
#include <dlib/dir_nav.h>
12
13


Davis King's avatar
Davis King committed
14
15
16
const char* VERSION = "0.1";


17
18
19
20

using namespace std;
using namespace dlib;

21
22
// ----------------------------------------------------------------------------------------

Davis King's avatar
Davis King committed
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
void create_new_dataset (
    const parser_type& parser
)
{
    using namespace dlib::image_dataset_metadata;

    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();

    unsigned long depth = 0;
    if (parser.option("r"))
        depth = 30;

    dataset meta;
    meta.name = "imglab dataset";
    meta.comment = "Created by imglab tool.";
    for (unsigned long i = 0; i < parser.number_of_arguments(); ++i)
    {
        try
        {
            const string temp = strip_path(file(parser[i]).full_name(), parent_dir);
            meta.images.push_back(image(temp));
        }
        catch (dlib::file::file_not_found&)
        {
            // then parser[i] should be a directory

            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());

            for (unsigned long j = 0; j < files.size(); ++j)
            {
                meta.images.push_back(image(strip_path(files[j].full_name(), parent_dir)));
            }
        }
    }

    save_image_dataset_metadata(meta, filename);
}

// ----------------------------------------------------------------------------------------
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
void print_all_labels (
    const dlib::image_dataset_metadata::dataset& data
)
{
    std::set<std::string> labels;
    for (unsigned long i = 0; i < data.images.size(); ++i)
    {
        for (unsigned long j = 0; j < data.images[i].boxes.size(); ++j)
        {
            labels.insert(data.images[i].boxes[j].label);
        }
    }

    for (std::set<std::string>::iterator i = labels.begin(); i != labels.end(); ++i)
    {
        if (i->size() != 0)
        {
            cout << *i << endl;
        }
    }
}

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

Davis King's avatar
Davis King committed
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
void rename_labels (
    dlib::image_dataset_metadata::dataset& data,
    const std::string& from,
    const std::string& to
)
{
    for (unsigned long i = 0; i < data.images.size(); ++i)
    {
        for (unsigned long j = 0; j < data.images[i].boxes.size(); ++j)
        {
            if (data.images[i].boxes[j].label == from)
                data.images[i].boxes[j].label = to;
        }
    }

}

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

113
114
115
116
117
118
119
120
121
int main(int argc, char** argv)
{
    try
    {

        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);
122
        parser.add_option("r","Search directories recursively for images.");
123
        parser.add_option("l","List all the labels in the given XML file.");
Davis King's avatar
Davis King committed
124
        parser.add_option("rename", "Rename all labels of <arg1> to <arg2>.",2);
Davis King's avatar
Davis King committed
125
        parser.add_option("v","Display version.");
126
127
        parser.add_option("convert","Convert foreign image Annotations from <arg> format to the imglab format. "
                          "Supported formats: pascal-voc",1);
128
129
130

        parser.parse(argc, argv);

131
        const char* singles[] = {"h","c","r","l","convert"};
132
        parser.check_one_time_options(singles);
133
134
        const char* c_sub_ops[] = {"r", "convert"};
        parser.check_sub_options("c", c_sub_ops);
Davis King's avatar
Davis King committed
135
136
        parser.check_incompatible_options("c", "l");
        parser.check_incompatible_options("c", "rename");
Davis King's avatar
Davis King committed
137
        parser.check_incompatible_options("l", "rename");
138
139
140
141
        parser.check_incompatible_options("convert", "l");
        parser.check_incompatible_options("convert", "rename");
        const char* convert_args[] = {"pascal-voc"};
        parser.check_option_arg_range("convert", convert_args);
142
143
144

        if (parser.option("h"))
        {
145
            cout << "Usage: imglab [options] <image files/directories or XML file>\n";
146
            parser.print_options(cout);
Davis King's avatar
Davis King committed
147
            cout << endl << endl;
148
149
150
            return EXIT_SUCCESS;
        }

Davis King's avatar
Davis King committed
151
152
153
154
155
156
157
158
159
        if (parser.option("v"))
        {
            cout << "imglab v" << VERSION 
                 << "\nCompiled: " << __TIME__ << " " << __DATE__ 
                 << "\nWritten by Davis King\n";
            cout << "Check for updates at http://dlib.net\n\n";
            return EXIT_SUCCESS;
        }

160
161
        if (parser.option("c"))
        {
162
163
164
165
166
167
168
169
170
171
172
            if (parser.option("convert"))
            {
                if (parser.option("convert").argument() == "pascal-voc")
                {
                    convert_pascal_voc(parser);
                }
            }
            else
            {
                create_new_dataset(parser);
            }
173
174
175
            return EXIT_SUCCESS;
        }

176
177
178
179
180
181
182
183
184
185
186
187
188
189
        if (parser.option("l"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The -l option requires you to give one XML file on the command line." << endl;
                return EXIT_FAILURE;
            }

            dlib::image_dataset_metadata::dataset data;
            load_image_dataset_metadata(data, parser[0]);
            print_all_labels(data);
            return EXIT_SUCCESS;
        }

Davis King's avatar
Davis King committed
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
        if (parser.option("rename"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --rename option requires you to give one XML file on the command line." << endl;
                return EXIT_FAILURE;
            }

            dlib::image_dataset_metadata::dataset data;
            load_image_dataset_metadata(data, parser[0]);
            for (unsigned long i = 0; i < parser.option("rename").count(); ++i)
            {
                rename_labels(data, parser.option("rename").argument(0,i), parser.option("rename").argument(1,i));
            }
            save_image_dataset_metadata(data, parser[0]);
            return EXIT_SUCCESS;
        }

208
209
        if (parser.number_of_arguments() == 1)
        {
210
211
            metadata_editor editor(parser[0]);
            editor.wait_until_closed();
212
        }
213
214
215
    }
    catch (exception& e)
    {
216
        cerr << e.what() << endl;
217
218
219
220
        return EXIT_FAILURE;
    }
}

Davis King's avatar
Davis King committed
221
222
// ----------------------------------------------------------------------------------------