main.cpp 22.4 KB
Newer Older
1

2
#include "dlib/data_io.h"
3
#include "dlib/string.h"
4
#include "metadata_editor.h"
Davis King's avatar
Davis King committed
5
#include "convert_pascal_xml.h"
6
#include "convert_pascal_v1.h"
7
#include "convert_idl.h"
Davis King's avatar
Davis King committed
8
#include <dlib/cmd_line_parser.h>
Davis King's avatar
Davis King committed
9
#include <dlib/image_transforms.h>
10
#include <dlib/svm.h>
11
#include <dlib/console_progress_indicator.h>
12

13
14
#include <iostream>
#include <fstream>
15
#include <string>
16
#include <set>
17

18
#include <dlib/dir_nav.h>
19
20


21
const char* VERSION = "1.1";
Davis King's avatar
Davis King committed
22
23


24
25
26
27

using namespace std;
using namespace dlib;

28
29
// ----------------------------------------------------------------------------------------

Davis King's avatar
Davis King committed
30
void create_new_dataset (
Davis King's avatar
Davis King committed
31
    const command_line_parser& parser
Davis King's avatar
Davis King committed
32
33
34
35
36
37
38
39
)
{
    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);
Davis King's avatar
Davis King committed
40
    const std::string parent_dir = get_parent_directory(file(filename));
Davis King's avatar
Davis King committed
41
42
43
44
45
46
47
48
49
50
51
52

    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
        {
Davis King's avatar
Davis King committed
53
            const string temp = strip_path(file(parser[i]), parent_dir);
Davis King's avatar
Davis King committed
54
55
56
57
58
59
60
61
62
63
64
65
66
            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)
            {
Davis King's avatar
Davis King committed
67
                meta.images.push_back(image(strip_path(files[j], parent_dir)));
Davis King's avatar
Davis King committed
68
69
70
71
72
73
74
75
            }
        }
    }

    save_image_dataset_metadata(meta, filename);
}

// ----------------------------------------------------------------------------------------
76

77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
int split_dataset (
    const command_line_parser& parser
)
{
    if (parser.number_of_arguments() != 1)
    {
        cerr << "The --split option requires you to give one XML file on the command line." << endl;
        return EXIT_FAILURE;
    }

    const std::string label = parser.option("split").argument();

    dlib::image_dataset_metadata::dataset data, data_with, data_without;
    load_image_dataset_metadata(data, parser[0]);

    data_with.name = data.name;
    data_with.comment = data.comment;
    data_without.name = data.name;
    data_without.comment = data.comment;

    for (unsigned long i = 0; i < data.images.size(); ++i)
    {
        dlib::image_dataset_metadata::image temp = data.images[i];

        bool has_the_label = false;
        // check for the label we are looking for
        for (unsigned long j = 0; j < temp.boxes.size(); ++j)
        {
            if (temp.boxes[j].label == label)
            {
                has_the_label = true;
                break;
            }
        }

        if (has_the_label)
        {
            std::vector<dlib::image_dataset_metadata::box> boxes;
            // remove other labels
            for (unsigned long j = 0; j < temp.boxes.size(); ++j)
            {
                if (temp.boxes[j].label == label)
                {
                    // put only the boxes with the label we want into boxes
                    boxes.push_back(temp.boxes[j]);
                }
            }
            temp.boxes = boxes;
            data_with.images.push_back(temp);
        }
        else
        {
            data_without.images.push_back(temp);
        }
    }


    save_image_dataset_metadata(data_with, left_substr(parser[0],".") + "_with_"+label + ".xml");
    save_image_dataset_metadata(data_without, left_substr(parser[0],".") + "_without_"+label + ".xml");

    return EXIT_SUCCESS;
}

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

142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
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;
        }
    }
}

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

166
167
168
169
170
void print_all_label_stats (
    const dlib::image_dataset_metadata::dataset& data
)
{
    std::map<std::string, running_stats<double> > area_stats, aspect_ratio;
171
    std::map<std::string, int> image_hits;
172
173
174
    std::set<std::string> labels;
    for (unsigned long i = 0; i < data.images.size(); ++i)
    {
175
        std::set<std::string> temp;
176
177
178
        for (unsigned long j = 0; j < data.images[i].boxes.size(); ++j)
        {
            labels.insert(data.images[i].boxes[j].label);
179
            temp.insert(data.images[i].boxes[j].label);
180
181
182
183
184

            area_stats[data.images[i].boxes[j].label].add(data.images[i].boxes[j].rect.area());
            aspect_ratio[data.images[i].boxes[j].label].add(data.images[i].boxes[j].rect.width()/
                                                    (double)data.images[i].boxes[j].rect.height());
        }
185
186
187
188

        // count the number of images for each label
        for (std::set<std::string>::iterator i = temp.begin(); i != temp.end(); ++i)
            image_hits[*i] += 1;
189
190
    }

191
    cout << "Number of images: "<< data.images.size() << endl;
192
193
194
195
    cout << "Number of different labels: "<< labels.size() << endl << endl;

    for (std::set<std::string>::iterator i = labels.begin(); i != labels.end(); ++i)
    {
196
197
198
        if (i->size() == 0)
            cout << "Unlabeled Boxes:" << endl;
        else
199
            cout << "Label: "<< *i << endl;
200
201
202
203
204
205
206
207
208
        cout << "   number of images:      " << image_hits[*i] << endl;
        cout << "   number of occurrences: " << area_stats[*i].current_n() << endl;
        cout << "   min box area:    " << area_stats[*i].min() << endl;
        cout << "   max box area:    " << area_stats[*i].max() << endl;
        cout << "   mean box area:   " << area_stats[*i].mean() << endl;
        cout << "   stddev box area: " << area_stats[*i].stddev() << endl;
        cout << "   mean width/height ratio:   " << aspect_ratio[*i].mean() << endl;
        cout << "   stddev width/height ratio: " << aspect_ratio[*i].stddev() << endl;
        cout << endl;
209
210
211
212
213
    }
}

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

Davis King's avatar
Davis King committed
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
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;
        }
    }

}

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

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
void merge_metadata_files (
    const command_line_parser& parser
)
{
    image_dataset_metadata::dataset src, dest;
    load_image_dataset_metadata(src, parser.option("add").argument(0));
    load_image_dataset_metadata(dest, parser.option("add").argument(1));

    std::map<string,image_dataset_metadata::image> merged_data;
    for (unsigned long i = 0; i < dest.images.size(); ++i)
        merged_data[dest.images[i].filename] = dest.images[i];
    // now add in the src data and overwrite anything if there are duplicate entries.
    for (unsigned long i = 0; i < src.images.size(); ++i)
        merged_data[src.images[i].filename] = src.images[i];

    // copy merged data into dest
    dest.images.clear();
    for (std::map<string,image_dataset_metadata::image>::const_iterator i = merged_data.begin(); 
        i != merged_data.end(); ++i)
    {
        dest.images.push_back(i->second);
    }

    save_image_dataset_metadata(dest, "merged.xml");
}

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

Davis King's avatar
Davis King committed
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
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
string to_png_name (const string& filename)
{
    string::size_type pos = filename.find_last_of(".");
    if (pos == string::npos)
        throw dlib::error("invalid filename: " + filename);
    return filename.substr(0,pos) + ".png";
}

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

void flip_dataset(const command_line_parser& parser)
{
#ifdef DLIB_PNG_SUPPORT
    image_dataset_metadata::dataset metadata;
    const string datasource = parser.option("flip").argument();
    load_image_dataset_metadata(metadata,datasource);

    // Set the current directory to be the one that contains the
    // metadata file. We do this because the file might contain
    // file paths which are relative to this folder.
    set_current_dir(get_parent_directory(file(datasource)));

    const string metadata_filename = get_parent_directory(file(datasource)).full_name() +
        directory::get_separator() + "flipped_" + file(datasource).name();


    array2d<rgb_pixel> img, temp;
    for (unsigned long i = 0; i < metadata.images.size(); ++i)
    {
        file f(metadata.images[i].filename);
        const string filename = get_parent_directory(f).full_name() + directory::get_separator() + "flipped_" + to_png_name(f.name());

        load_image(img, metadata.images[i].filename);
        flip_image_left_right(img, temp);
        save_png(temp, filename);

        for (unsigned long j = 0; j < metadata.images[i].boxes.size(); ++j)
        {
            metadata.images[i].boxes[j].rect = impl::flip_rect_left_right(metadata.images[i].boxes[j].rect, get_rect(img));

            // flip all the object parts
            std::map<std::string,point>::iterator k;
            for (k = metadata.images[i].boxes[j].parts.begin(); k != metadata.images[i].boxes[j].parts.end(); ++k)
            {
                k->second = impl::flip_rect_left_right(rectangle(k->second,k->second), get_rect(img)).tl_corner();
            }
        }

        metadata.images[i].filename = filename;
    }

    save_image_dataset_metadata(metadata, metadata_filename);
#else
    throw dlib::error("imglab must be compiled with libpng if you want to use the --flip option.");
#endif
}

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
370
371
372
373
374
375
376
377
378
379
380
381
382
// ----------------------------------------------------------------------------------------


int tile_dataset(const command_line_parser& parser)
{
#if defined(DLIB_PNG_SUPPORT) && defined(DLIB_JPEG_SUPPORT)
    if (parser.number_of_arguments() != 1)
    {
        cerr << "The --tile option requires you to give one XML file on the command line." << endl;
        return EXIT_FAILURE;
    }

    string out_image = parser.option("tile").argument();
    string ext = right_substr(out_image,".");
    if (ext != "png" && ext != "jpg")
    {
        cerr << "The output image file must have either .png or .jpg extension." << endl;
        return EXIT_FAILURE;
    }

    const unsigned long chip_size = get_option(parser, "size", 7000);

    dlib::image_dataset_metadata::dataset data;
    load_image_dataset_metadata(data, parser[0]);
    locally_change_current_dir chdir(get_parent_directory(file(parser[0])));
    dlib::array<array2d<rgb_pixel> > images;
    console_progress_indicator pbar(data.images.size());
    for (unsigned long i = 0; i < data.images.size(); ++i)
    {
        pbar.print_status(i);
        array2d<rgb_pixel> img;
        load_image(img, data.images[i].filename);

        // figure out what chips we want to take from this image
        std::vector<chip_details> dets;
        for (unsigned long j = 0; j < data.images[i].boxes.size(); ++j)
        {
            if (data.images[i].boxes[j].ignore)
                continue;

            rectangle rect = data.images[i].boxes[j].rect;
            dets.push_back(chip_details(rect, chip_size));
        }
        // Now grab all those chips at once.
        dlib::array<array2d<rgb_pixel> > chips;
        extract_image_chips(img, dets, chips);
        // and put the chips into the output.
        for (unsigned long j = 0; j < chips.size(); ++j)
            images.push_back(chips[j]);
    }

    chdir.revert();

    if (ext == "png")
        save_png(tile_images(images), out_image);
    else
        save_jpeg(tile_images(images), out_image);

    return EXIT_SUCCESS;
#else
    throw dlib::error("imglab must be compiled with libpng and libjpeg if you want to use the --tile option.");
#endif
}


Davis King's avatar
Davis King committed
383
384
// ----------------------------------------------------------------------------------------

385
386
387
388
389
int main(int argc, char** argv)
{
    try
    {

Davis King's avatar
Davis King committed
390
        command_line_parser parser;
391
392

        parser.add_option("h","Displays this information.");
393
394
395
        parser.add_option("v","Display version.");

        parser.set_group_name("Creating XML files");
396
        parser.add_option("c","Create an XML file named <arg> listing a set of images.",1);
397
        parser.add_option("r","Search directories recursively for images.");
398
399
400
401
        parser.add_option("convert","Convert foreign image Annotations from <arg> format to the imglab format. "
                          "Supported formats: pascal-xml, pascal-v1, idl.",1);

        parser.set_group_name("Viewing/Editing XML files");
402
403
        parser.add_option("tile","Chip out all the objects and save them as one big image called <arg>.",1);
        parser.add_option("size","When using --tile, make each object contain about <arg> pixels (default 7000).",1);
404
        parser.add_option("l","List all the labels in the given XML file.");
405
        parser.add_option("stats","List detailed statistics on the object labels in the given XML file.");
Davis King's avatar
Davis King committed
406
        parser.add_option("rename", "Rename all labels of <arg1> to <arg2>.",2);
407
        parser.add_option("parts","The display will allow image parts to be labeled.  The set of allowable parts "
Davis King's avatar
Davis King committed
408
                          "is defined by <arg> which should be a space separated list of parts.",1);
409
        parser.add_option("rmdiff","Remove boxes marked as difficult.");
410
411
412
413
414
415
        parser.add_option("shuffle","Randomly shuffle the order of the images listed in file <arg>.");
        parser.add_option("seed", "When using --shuffle, set the random seed to the string <arg>.",1);
        parser.add_option("split", "Split the contents of an XML file into two separate files.  One containing the "
            "images with objects labeled <arg> and another file with all the other images.  Additionally, the file "
            "containing the <arg> labeled objects will not contain any other labels other than <arg>. "
            "That is, the images in the first file are stripped of all labels other than the <arg> labels.",1);
416
417
418
419
        parser.add_option("add", "Add the image metadata from <arg1> into <arg2>.  If any of the image "
                                 "tags are in both files then the ones in <arg2> are deleted and replaced with the "
                                 "image tags from <arg1>.  The results are saved into merged.xml and neither <arg1> or "
                                 "<arg2> files are modified.",2);
Davis King's avatar
Davis King committed
420
421
        parser.add_option("flip", "Read an XML image dataset from the <arg> XML file and output a left-right flipped "
                                  "version of the dataset and an accompanying flipped XML file named flipped_<arg>.",1);
422
423
424

        parser.parse(argc, argv);

425
426
        const char* singles[] = {"h","c","r","l","convert","parts","rmdiff","seed", "shuffle", "split", "add", 
                                 "flip", "tile", "size"};
427
        parser.check_one_time_options(singles);
428
429
        const char* c_sub_ops[] = {"r", "convert"};
        parser.check_sub_options("c", c_sub_ops);
430
        parser.check_sub_option("shuffle", "seed");
431
        parser.check_sub_option("tile", "size");
Davis King's avatar
Davis King committed
432
        parser.check_incompatible_options("c", "l");
433
        parser.check_incompatible_options("c", "rmdiff");
434
        parser.check_incompatible_options("c", "add");
Davis King's avatar
Davis King committed
435
        parser.check_incompatible_options("c", "flip");
Davis King's avatar
Davis King committed
436
        parser.check_incompatible_options("c", "rename");
437
        parser.check_incompatible_options("c", "parts");
438
        parser.check_incompatible_options("c", "tile");
Davis King's avatar
Davis King committed
439
        parser.check_incompatible_options("l", "rename");
440
        parser.check_incompatible_options("l", "add");
441
        parser.check_incompatible_options("l", "parts");
Davis King's avatar
Davis King committed
442
443
        parser.check_incompatible_options("l", "flip");
        parser.check_incompatible_options("add", "flip");
444
445
446
        parser.check_incompatible_options("add", "tile");
        parser.check_incompatible_options("flip", "tile");
        parser.check_incompatible_options("shuffle", "tile");
447
448
        parser.check_incompatible_options("convert", "l");
        parser.check_incompatible_options("convert", "rename");
449
        parser.check_incompatible_options("convert", "parts");
450
        parser.check_incompatible_options("rmdiff", "rename");
451
        const char* convert_args[] = {"pascal-xml","pascal-v1","idl"};
452
        parser.check_option_arg_range("convert", convert_args);
453
454
455

        if (parser.option("h"))
        {
456
            cout << "Usage: imglab [options] <image files/directories or XML file>\n";
457
            parser.print_options(cout);
Davis King's avatar
Davis King committed
458
            cout << endl << endl;
459
460
461
            return EXIT_SUCCESS;
        }

462
463
464
465
466
467
        if (parser.option("add"))
        {
            merge_metadata_files(parser);
            return EXIT_SUCCESS;
        }

Davis King's avatar
Davis King committed
468
469
470
471
472
473
        if (parser.option("flip"))
        {
            flip_dataset(parser);
            return EXIT_SUCCESS;
        }

Davis King's avatar
Davis King committed
474
475
476
477
478
479
480
481
482
        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;
        }

483
484
485
486
487
        if (parser.option("tile"))
        {
            return tile_dataset(parser);
        }

488
489
        if (parser.option("c"))
        {
490
491
            if (parser.option("convert"))
            {
Davis King's avatar
Davis King committed
492
493
                if (parser.option("convert").argument() == "pascal-xml")
                    convert_pascal_xml(parser);
494
495
                else if (parser.option("convert").argument() == "pascal-v1")
                    convert_pascal_v1(parser);
496
497
                else if (parser.option("convert").argument() == "idl")
                    convert_idl(parser);
498
499
500
501
502
            }
            else
            {
                create_new_dataset(parser);
            }
503
504
            return EXIT_SUCCESS;
        }
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
        
        if (parser.option("rmdiff"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --rmdiff 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 < data.images.size(); ++i)
            {
                std::vector<dlib::image_dataset_metadata::box> boxes;
                for (unsigned long j = 0; j < data.images[i].boxes.size(); ++j)
                {
                    if (!data.images[i].boxes[j].difficult)
                        boxes.push_back(data.images[i].boxes[j]);
                }
                data.images[i].boxes = boxes;
            }
            save_image_dataset_metadata(data, parser[0]);
            return EXIT_SUCCESS;
        }
529

530
531
532
533
534
535
536
537
538
539
540
541
542
543
        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;
        }

544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
        if (parser.option("split"))
        {
            return split_dataset(parser);
        }

        if (parser.option("shuffle"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The -shuffle 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]);
            const string default_seed = cast_to_string(time(0));
            const string seed = get_option(parser, "seed", default_seed);
            dlib::rand rnd(seed);
            randomize_samples(data.images, rnd);
            save_image_dataset_metadata(data, parser[0]);
            return EXIT_SUCCESS;
        }

567
568
569
570
571
572
573
574
575
576
577
578
579
580
        if (parser.option("stats"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --stats 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_label_stats(data);
            return EXIT_SUCCESS;
        }

Davis King's avatar
Davis King committed
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
        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;
        }

599
600
        if (parser.number_of_arguments() == 1)
        {
601
            metadata_editor editor(parser[0]);
602
603
604
605
606
607
608
609
            if (parser.option("parts"))
            {
                std::vector<string> parts = split(parser.option("parts").argument());
                for (unsigned long i = 0; i < parts.size(); ++i)
                {
                    editor.add_labelable_part_name(parts[i]);
                }
            }
610
            editor.wait_until_closed();
611
        }
612
613
614
    }
    catch (exception& e)
    {
615
        cerr << e.what() << endl;
616
617
618
619
        return EXIT_FAILURE;
    }
}

Davis King's avatar
Davis King committed
620
621
// ----------------------------------------------------------------------------------------