Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
main.cpp 16.6 KB
Newer Older
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
1
2
3 

#include <iostream>
#include <fstream>
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
4 
#include <string>
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
5
6 

#include <dlib/cmd_line_parser.h>
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
7 
#include <dlib/geometry.h>
Davis King's avatar
Made the initial file selection more flexible.  
Davis King committed
8 
#include <dlib/dir_nav.h>
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
9
10
11
12 

#include <sstream>
#include <dlib/compress_stream.h>
#include <dlib/base64.h>
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
13 
#include <dlib/xml_parser.h>
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
14
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
15
16
17
18 

using namespace std;
using namespace dlib;
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
19
20
21
22 
// ----------------------------------------------------------------------------------------

namespace dlib
{
Davis King's avatar
renamed some things  
Davis King committed
23 
    namespace image_dataset_metadata
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47 
    {
        struct box
        {
            box() : head(-0xFFFF,-0xFFFF) {}

            rectangle rect;

            // optional fields
            std::string label;
            point head; // a value of (-0xFFFF,-0xFFFF) indicates the field not supplied

            bool has_head() const { return head != point(-0xFFFF,-0xFFFF); }
            bool has_label() const { return label.size() != 0; }
        };

        struct image
        {
            image() {}
            image(const std::string& f) : filename(f) {}

            std::string filename;
            std::vector<box> boxes;
        };
Davis King's avatar
renamed some things  
Davis King committed
48 
        struct dataset
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70 
        {
            std::vector<image> images;
            std::string comment;
            std::string name;
        };

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

        const std::string get_decoded_string();
        void create_image_metadata_stylesheet_file()
        {
            ofstream fout("image_metadata_stylesheet.xsl");
            if (!fout)
                throw dlib::error("ERROR: Unable to open image_metadata_stylesheet.xsl for writing.");

            fout << get_decoded_string();

            if (!fout)
                throw dlib::error("ERROR: Unable to write to image_metadata_stylesheet.xsl.");
        }

        void save_image_dataset_metadata (
Davis King's avatar
renamed some things  
Davis King committed
71 
            const dataset& meta,
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
72
73
74
75
76 
            const std::string& filename
        )
        {
            create_image_metadata_stylesheet_file();
Davis King's avatar
renamed some things  
Davis King committed
77 
            const std::vector<image>& images = meta.images;
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
78
79
80
81
82
83
84
85 

            ofstream fout(filename.c_str());
            if (!fout)
                throw dlib::error("ERROR: Unable to open " + filename + " for writing.");

            fout << "<?xml version='1.0' encoding='ISO-8859-1'?>\n";
            fout << "<?xml-stylesheet type='text/xsl' href='image_metadata_stylesheet.xsl'?>\n";
            fout << "<dataset>\n";
Davis King's avatar
renamed some things  
Davis King committed
86
87 
            fout << "<name>" << meta.name << "</name>\n";
            fout << "<comment>" << meta.comment << "</comment>\n";
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
88
89
90
91
92
93
94
95
96
97
98
99 
            fout << "<images>\n";
            for (unsigned long i = 0; i < images.size(); ++i)
            {
                fout << "  <image file='" << images[i].filename << "'>\n";

                // save all the boxes
                for (unsigned long j = 0; j < images[i].boxes.size(); ++j)
                {
                    const box& b = images[i].boxes[j];
                    fout << "    <box top='" << b.rect.top() << "' "
                                 << "left='" << b.rect.left() << "' "
                                << "width='" << b.rect.width() << "' "
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116 
                               << "height='" << b.rect.height() << "'";

                    if (b.has_label() || b.has_head())
                    {
                        fout << ">\n";

                        if (b.has_label())
                            fout << "      <label>" << b.label << "</label>\n";
                        if (b.has_head())
                            fout << "      <head x='"<< b.head.x() <<"' y='"<< b.head.y() <<"'/>\n";

                        fout << "    </box>\n";
                    }
                    else
                    {
                        fout << "/>\n";
                    }
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
117
118
119
120
121
122
123
124
125 
                }



                fout << "  </image>\n";

                if (!fout)
                    throw dlib::error("ERROR: Unable to write to " + filename + ".");
            }
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
126 
            fout << "</images>\n";
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
127
128
129 
            fout << "</dataset>";
        }
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
130
131
132
133
134
135
136
137
138
139 
    // ------------------------------------------------------------------------------------
    // ------------------------------------------------------------------------------------
    // ------------------------------------------------------------------------------------

        class doc_handler : public document_handler
        {
            std::vector<std::string> ts;
            image temp_image;
            box temp_box;
Davis King's avatar
renamed some things  
Davis King committed
140 
            dataset& meta;
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
141
142
143
144 

        public:

            doc_handler(
Davis King's avatar
renamed some things  
Davis King committed
145 
                dataset& metadata_
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
146 
            ):
Davis King's avatar
renamed some things  
Davis King committed
147 
                meta(metadata_)
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
148
149
150
151
152
153 
            {}


            virtual void start_document (
            )
            {
Davis King's avatar
renamed some things  
Davis King committed
154 
                meta = dataset();
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238 
                ts.clear();
                temp_image = image();
                temp_box = box();
            }

            virtual void end_document (
            )
            {
            }

            virtual void start_element ( 
                const unsigned long line_number,
                const std::string& name,
                const dlib::attribute_list& atts
            )
            {
                if (ts.size() == 0) 
                {
                    if (name != "dataset")
                    {
                        std::ostringstream sout;
                        sout << "Invalid XML document.  Root tag must be <dataset>.  Found <" << name << "> instead.";
                        throw dlib::error(sout.str());
                    }
                    else
                    {
                        ts.push_back(name);
                        return;
                    }
                }


                if (name == "box")
                {
                    if (atts.is_in_list("top")) temp_box.rect.top() = sa = atts["top"];
                    else throw dlib::error("<box> missing required attribute 'top'");

                    if (atts.is_in_list("left")) temp_box.rect.left() = sa = atts["left"];
                    else throw dlib::error("<box> missing required attribute 'left'");

                    if (atts.is_in_list("width")) temp_box.rect.right() = sa = atts["width"];
                    else throw dlib::error("<box> missing required attribute 'width'");

                    if (atts.is_in_list("height")) temp_box.rect.bottom() = sa = atts["height"];
                    else throw dlib::error("<box> missing required attribute 'height'");

                    temp_box.rect.bottom() += temp_box.rect.top()-1;
                    temp_box.rect.right() += temp_box.rect.left()-1;
                }
                else if (name == "head" && ts.back() == "box")
                {
                    if (atts.is_in_list("x")) temp_box.head.x() = sa = atts["x"];
                    else throw dlib::error("<head> missing required attribute 'x'");

                    if (atts.is_in_list("y")) temp_box.head.y() = sa = atts["y"];
                    else throw dlib::error("<head> missing required attribute 'y'");
                }
                else if (name == "image")
                {
                    temp_image.boxes.clear();

                    if (atts.is_in_list("file")) temp_image.filename = atts["file"];
                    else throw dlib::error("<image> missing required attribute 'file'");
                }

                ts.push_back(name);
            }

            virtual void end_element ( 
                const unsigned long line_number,
                const std::string& name
            )
            {
                ts.pop_back();
                if (ts.size() == 0)
                    return;

                if (name == "box" && ts.back() == "image")
                {
                    temp_image.boxes.push_back(temp_box);
                    temp_box = box();
                }
                else if (name == "image" && ts.back() == "images")
                {
Davis King's avatar
renamed some things  
Davis King committed
239 
                    meta.images.push_back(temp_image);
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
240
241
242
243
244
245
246
247
248
249 
                    temp_image = image();
                }
            }

            virtual void characters ( 
                const std::string& data
            )
            {
                if (ts.size() == 2 && ts[1] == "name")
                {
Davis King's avatar
renamed some things  
Davis King committed
250 
                    meta.name = trim(data);
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
251
252
253 
                }
                else if (ts.size() == 2 && ts[1] == "comment")
                {
Davis King's avatar
renamed some things  
Davis King committed
254 
                    meta.comment = trim(data);
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293 
                }
                else if (ts.size() >= 2 && ts[ts.size()-1] == "label" && 
                                           ts[ts.size()-2] == "box")
                {
                    temp_box.label = trim(data);
                }
            }

            virtual void processing_instruction (
                const unsigned long ,
                const std::string& ,
                const std::string& 
            )
            {
            }
        };

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

        class xml_error_handler : public error_handler
        {
        public:
            virtual void error (
                const unsigned long line_number
            )
            {
                cout << "There is a non-fatal error on line " << line_number << " in the file we are parsing." << endl;
            }

            virtual void fatal_error (
                const unsigned long line_number
            )
            {
                std::ostringstream sout;
                sout << "There is a fatal error on line " << line_number << " so parsing will now halt.";
                throw dlib::error(sout.str());
            }
        };
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
294
295
296 
    // ------------------------------------------------------------------------------------

        void load_image_dataset_metadata (
Davis King's avatar
renamed some things  
Davis King committed
297 
            dataset& meta,
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
298
299
300 
            const std::string& filename
        )
        {
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
301 
            xml_error_handler eh;
Davis King's avatar
renamed some things  
Davis King committed
302 
            doc_handler dh(meta);
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
303
304
305
306
307
308
309
310
311 

            std::ifstream fin(filename.c_str());
            if (!fin)
                throw dlib::error("ERROR: unable to open " + filename + " for reading.");

            xml_parser::kernel_1a parser;
            parser.add_document_handler(dh);
            parser.add_error_handler(eh);
            parser.parse(fin);
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
312
313 
        }
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
314
315 
    // ------------------------------------------------------------------------------------
    // ------------------------------------------------------------------------------------
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369 
    // ------------------------------------------------------------------------------------

        // This function returns the contents of the file 'images.xsl'
        const std::string get_decoded_string()
        {
            dlib::base64 base64_coder;
            dlib::compress_stream::kernel_1ea compressor;
            std::ostringstream sout;
            std::istringstream sin;

            // The base64 encoded data from the file 'image_metadata_stylesheet.xsl' we want to decode and return.
            sout << "PFWfgmWfCHr1DkV63lbjjeY2dCc2FbHDOVh0Kd7dkvaOfRYrOG24f0x77/5iMVq8FtE3UBxtGwSd";
            sout << "1ZHOHRSHgieNoeBv8ssJQ75RRxYtFKRY3OTPX5eKQoCN9jUaUnHnR4QZtEHgmKqXSs50Yrdd+2Ah";
            sout << "gNyarPZCiR6nvqNvCjtP2MP5FxleqNf8Fylatm2KdsXmrv5K87LYVN7i7JMkmZ++cTXYSOxDmxZi";
            sout << "OiCH8funXUdF9apDW547gCjz9HOQUI6dkz5dYUeFjfp6dFugpnaJyyprFLKq048Qk7+QiL4CNF/G";
            sout << "7e0VpBw8dMpiyRNi2fSQGSZGfIAUQKKT6+rPwQoRH2spdjsdXVWj4XQAqBX87nmqMnqjMhn/Vd1s";
            sout << "W5aoC0drwRGu3Xe3gn9vBL8hBkRXcKndVwPFMPUeADJ2Z2tc8cb9MqwkOCPBKdfQ+w24u/kvjdsQ";
            sout << "Kq20+MI2uPOSGbCebui/yMveNYUeot5b4s+ZyRIdcbURUHlQQ68BZQVj0WTw+8ohK7VUpjidF/T+";
            sout << "gawZzSnr/HrocEedQL7DyjyVuC03qoFGlwe6Q/4ynndgscVhXyVleneSpOlx/CGOlBKXzwyEpEi8";
            sout << "JV4kPhoJ21N5va7iyDgxuikbDT94tFYje9GtXQOPGPERKzxnp6dCor3a5dPLOfmVsvAdEDtr7B3e";
            sout << "WJpFbfD4XtFfllWmwxggCqyfvz/q7kMQzhsH623CLRDre9AsIRg830e6T992uKgtXx8QM4GPhsjc";
            sout << "leFCNzwDqpeORW7oJ72COj3pgGIn2/BGDWrr1oqu0ACKy7vY32VcgvyFcWbqqdEX71SQ9LcxCbet";
            sout << "iwrBJY8n//Whq55kD0hpSffmW+uup6sV8sbAI7rwnCYpfUNQHXgN+5sLot67jhO03FlafaEmrgfT";
            sout << "XKjBU+z64tCCy1uwJj8gx0CBVFhr8MPzGkKp3rAaeKzVYZJbC2TkZPL+PMjohL0SDhaXxiwX+pKo";
            sout << "VHNRqCmlnrxmB6kikuS5zN5Z6feHW6KNhhcEKuXvDawYNlppJvMt1lE0Q9oKx+JL3Atlm+V8/Q2R";
            sout << "fI5DQAaCaxMXdzJcDfNPMHOIcYlkQPIk9cdLqScrbFu9VjFpqGyGoTGOUiiJP4d3peJvEbRgZZ8d";
            sout << "SUaF4cycZh4yRSI06rrTGi3wyS5HjRFgS+giS2p0ZUi+7YAt1opbkDhcovTxZGUkuavZBCsjZw3C";
            sout << "1CNGurdvhw9mOOL9JY3Vn/wfb9t4ScwaGGrUoeyZHXMARoHYVF1ST7gGOHL2qMRMdXr+k0P3OLxO";
            sout << "4DVe0r3v93vpyqd5KgA=";


            // Put the data into the istream sin
            sin.str(sout.str());
            sout.str("");

            // Decode the base64 text into its compressed binary form
            base64_coder.decode(sin,sout);
            sin.clear();
            sin.str(sout.str());
            sout.str("");

            // Decompress the data into its original form
            compressor.decompress(sin,sout);

            // Return the decoded and decompressed data
            return sout.str();
        }


    }
}

// ----------------------------------------------------------------------------------------
Davis King's avatar
Made the initial file selection more flexible.  
Davis King committed
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401 
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.");
}

// ----------------------------------------------------------------------------------------
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
402
403
404
405
406
407
408
409
410
411 
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);
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
412 
        parser.add_option("r","Search directories recursively for images.");
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
413
414
415 

        parser.parse(argc, argv);
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
416 
        const char* singles[] = {"h","c","r"};
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
417 
        parser.check_one_time_options(singles);
Davis King's avatar
Further simplified initial image selection.  
Davis King committed
418 
        parser.check_sub_option("c", "r");
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
419
420
421 

        if (parser.option("h"))
        {
Davis King's avatar
Improved -h output  
Davis King committed
422 
            cout << "Usage: imglab [options] <image files/directories or XML file list>\n";
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
423 
            parser.print_options(cout);
Davis King's avatar
Improved -h output  
Davis King committed
424 
            cout << endl << endl;
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
425
426
427
428
429 
            return EXIT_SUCCESS;
        }

        if (parser.option("c"))
        {
Davis King's avatar
renamed some things  
Davis King committed
430 
            using namespace dlib::image_dataset_metadata;
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
431
Davis King's avatar
Made the initial file selection more flexible.  
Davis King committed
432
433
434
435
436
437 
            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();
Davis King's avatar
Further simplified initial image selection.  
Davis King committed
438
439
440
441 
            unsigned long depth = 0;
            if (parser.option("r"))
                depth = 30;
Davis King's avatar
renamed some things  
Davis King committed
442
443
444 
            dataset meta;
            meta.name = "imglab dataset";
            meta.comment = "Created by imglab tool.";
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
445
446 
            for (unsigned long i = 0; i < parser.number_of_arguments(); ++i)
            {
Davis King's avatar
Further simplified initial image selection.  
Davis King committed
447
448
449 
                try
                {
                    const string temp = strip_path(file(parser[i]).full_name(), parent_dir);
Davis King's avatar
renamed some things  
Davis King committed
450 
                    meta.images.push_back(image(temp));
Davis King's avatar
Further simplified initial image selection.  
Davis King committed
451
452
453
454 
                }
                catch (dlib::file::file_not_found&)
                {
                    // then parser[i] should be a directory
Davis King's avatar
Made the initial file selection more flexible.  
Davis King committed
455
Davis King's avatar
Further simplified initial image selection.  
Davis King committed
456
457
458
459 
                    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());
Davis King's avatar
Made the initial file selection more flexible.  
Davis King committed
460
Davis King's avatar
Further simplified initial image selection.  
Davis King committed
461
462 
                    for (unsigned long j = 0; j < files.size(); ++j)
                    {
Davis King's avatar
renamed some things  
Davis King committed
463 
                        meta.images.push_back(image(strip_path(files[j].full_name(), parent_dir)));
Davis King's avatar
Further simplified initial image selection.  
Davis King committed
464 
                    }
Davis King's avatar
Made the initial file selection more flexible.  
Davis King committed
465 
                }
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
466
467 
            }
Davis King's avatar
renamed some things  
Davis King committed
468 
            save_image_dataset_metadata(meta, filename);
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
469
470
471
472 

            return EXIT_SUCCESS;
        }
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
473
474 
        if (parser.number_of_arguments() == 1)
        {
Davis King's avatar
renamed some things  
Davis King committed
475
476
477 
            dlib::image_dataset_metadata::dataset meta;
            load_image_dataset_metadata(meta, parser[0]);
            save_image_dataset_metadata(meta, "out.xml");
Davis King's avatar
Added code to read the image metadata XML files  
Davis King committed
478 
        }
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
479
480
481
482
483
484
485
486 
    }
    catch (exception& e)
    {
        cout << e.what() << endl;
        return EXIT_FAILURE;
    }
}