Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
main.cpp 16.7 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
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
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
94
95
96
97
98
99 
// ----------------------------------------------------------------------------------------

namespace dlib
{
    namespace imglab
    {
        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;
        };

        struct image_dataset_metadata
        {
            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 (
            const image_dataset_metadata& metadata,
            const std::string& filename
        )
        {
            create_image_metadata_stylesheet_file();

            const std::vector<image>& images = metadata.images;

            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";
            fout << "<name>" << metadata.name << "</name>\n";
            fout << "<comment>" << metadata.comment << "</comment>\n";
            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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
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
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
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 
    // ------------------------------------------------------------------------------------
    // ------------------------------------------------------------------------------------
    // ------------------------------------------------------------------------------------

        class doc_handler : public document_handler
        {
            std::vector<std::string> ts;
            image temp_image;
            box temp_box;

            image_dataset_metadata& metadata;

        public:

            doc_handler(
                image_dataset_metadata& metadata_
            ):
                metadata(metadata_) 
            {}


            virtual void start_document (
            )
            {
                metadata = image_dataset_metadata();
                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")
                {
                    metadata.images.push_back(temp_image);
                    temp_image = image();
                }
            }

            virtual void characters ( 
                const std::string& data
            )
            {
                if (ts.size() == 2 && ts[1] == "name")
                {
                    metadata.name = trim(data);
                }
                else if (ts.size() == 2 && ts[1] == "comment")
                {
                    metadata.comment = trim(data);
                }
                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
Added code to read the image metadata XML files  
Davis King committed
297 
            image_dataset_metadata& metadata,
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
302
303
304
305
306
307
308
309
310
311 
            xml_error_handler eh;
            doc_handler dh(metadata);

            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
413 
        parser.add_option("d","Include all images files in directory <arg> in the new image file list.",1);
        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
414
415
416 

        parser.parse(argc, argv);
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
417 
        const char* singles[] = {"h","c","r"};
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
418 
        parser.check_one_time_options(singles);
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
419
420 
        parser.check_sub_option("c", "d");
        parser.check_sub_option("d", "r");
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
421
422
423
424
425
426
427
428
429
430
431 

        if (parser.option("h"))
        {
            cout << "Options:\n";
            parser.print_options(cout);
            cout << endl;
            return EXIT_SUCCESS;
        }

        if (parser.option("c"))
        {
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
432
433 
            using namespace dlib::imglab;
Davis King's avatar
Made the initial file selection more flexible.  
Davis King committed
434
435
436
437
438
439 
            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
Added some code so create initial image lists  
Davis King committed
440
441
442
443
444 
            image_dataset_metadata metadata;
            metadata.name = "imglab dataset";
            metadata.comment = "Created by imglab tool.";
            for (unsigned long i = 0; i < parser.number_of_arguments(); ++i)
            {
Davis King's avatar
Made the initial file selection more flexible.  
Davis King committed
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464 
                const string temp = strip_path(file(parser[i]).full_name(), parent_dir);
                metadata.images.push_back(image(temp));
            }


            for (unsigned long i = 0; i < parser.option("d").count(); ++i)
            {
                unsigned long depth = 0;
                if (parser.option("r"))
                    depth = 30;

                std::vector<file> files = get_files_in_directory_tree(parser.option("d").argument(0,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)
                {
                    metadata.images.push_back(image(strip_path(files[j].full_name(), parent_dir)));
                }
Davis King's avatar
Added some code so create initial image lists  
Davis King committed
465
466 
            }
Davis King's avatar
Made the initial file selection more flexible.  
Davis King committed
467 
            save_image_dataset_metadata(metadata, filename);
Davis King's avatar
Adding beginnings of a tool for dealing with labeled image datasets.  
Davis King committed
468
469
470
471 

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