main.cpp 44.8 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"
8
#include "cluster.h"
9
#include "flip_dataset.h"
Davis King's avatar
Davis King committed
10
#include <dlib/cmd_line_parser.h>
Davis King's avatar
Davis King committed
11
#include <dlib/image_transforms.h>
12
#include <dlib/svm.h>
13
#include <dlib/console_progress_indicator.h>
Davis King's avatar
Davis King committed
14
#include <dlib/md5.h>
15

16
17
#include <iostream>
#include <fstream>
18
#include <string>
19
#include <set>
20

21
#include <dlib/dir_nav.h>
22
23


24
const char* VERSION = "1.12";
Davis King's avatar
Davis King committed
25
26


27
28
29
30

using namespace std;
using namespace dlib;

31
32
// ----------------------------------------------------------------------------------------

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

    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
56
            const string temp = strip_path(file(parser[i]), parent_dir);
Davis King's avatar
Davis King committed
57
58
59
60
61
62
63
            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], 
64
                                                                  match_endings(".png .PNG .jpeg .JPEG .jpg .JPG .bmp .BMP .dng .DNG .gif .GIF"),
Davis King's avatar
Davis King committed
65
66
67
68
69
                                                                  depth);
            sort(files.begin(), files.end());

            for (unsigned long j = 0; j < files.size(); ++j)
            {
Davis King's avatar
Davis King committed
70
                meta.images.push_back(image(strip_path(files[j], parent_dir)));
Davis King's avatar
Davis King committed
71
72
73
74
75
76
77
78
            }
        }
    }

    save_image_dataset_metadata(meta, filename);
}

// ----------------------------------------------------------------------------------------
79

80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
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)
    {
102
        auto&& temp = data.images[i];
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

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

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

130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
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;
        }
    }
}

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

154
155
156
157
158
void print_all_label_stats (
    const dlib::image_dataset_metadata::dataset& data
)
{
    std::map<std::string, running_stats<double> > area_stats, aspect_ratio;
159
    std::map<std::string, int> image_hits;
160
    std::set<std::string> labels;
161
    unsigned long num_unignored_boxes = 0;
162
163
    for (unsigned long i = 0; i < data.images.size(); ++i)
    {
164
        std::set<std::string> temp;
165
166
167
        for (unsigned long j = 0; j < data.images[i].boxes.size(); ++j)
        {
            labels.insert(data.images[i].boxes[j].label);
168
            temp.insert(data.images[i].boxes[j].label);
169
170
171
172

            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());
173
174
175

            if (!data.images[i].boxes[j].ignore)
                ++num_unignored_boxes;
176
        }
177
178
179
180

        // 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;
181
182
    }

183
    cout << "Number of images: "<< data.images.size() << endl;
184
185
    cout << "Number of different labels: "<< labels.size() << endl;
    cout << "Number of non-ignored boxes: " << num_unignored_boxes << endl << endl;
186
187
188

    for (std::set<std::string>::iterator i = labels.begin(); i != labels.end(); ++i)
    {
189
190
191
        if (i->size() == 0)
            cout << "Unlabeled Boxes:" << endl;
        else
192
            cout << "Label: "<< *i << endl;
193
194
195
196
197
198
199
200
201
        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;
202
203
204
205
206
    }
}

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

Davis King's avatar
Davis King committed
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
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;
        }
    }

}

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

226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
void ignore_labels (
    dlib::image_dataset_metadata::dataset& data,
    const std::string& label
)
{
    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 == label)
                data.images[i].boxes[j].ignore = true;
        }
    }
}

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

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
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");
}

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

271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
void rotate_dataset(const command_line_parser& parser)
{
    image_dataset_metadata::dataset metadata;
    const string datasource = parser[0];
    load_image_dataset_metadata(metadata,datasource);

    double angle = get_option(parser, "rotate", 0);

    // 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 file_prefix = "rotated_"+ cast_to_string(angle) + "_";
    const string metadata_filename = get_parent_directory(file(datasource)).full_name() +
        directory::get_separator() + file_prefix + file(datasource).name();


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

        load_image(img, metadata.images[i].filename);
        const point_transform_affine tran = rotate_image(img, temp, angle*pi/180);
297
298
299
300
301
302
303
304
305
        if (parser.option("jpg"))
        {
            filename = to_jpg_name(filename);
            save_jpeg(temp, filename,JPEG_QUALITY);
        }
        else
        {
            save_png(temp, filename);
        }
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335

        for (unsigned long j = 0; j < metadata.images[i].boxes.size(); ++j)
        {
            const rectangle rect = metadata.images[i].boxes[j].rect;
            rectangle newrect;
            newrect += tran(rect.tl_corner());
            newrect += tran(rect.tr_corner());
            newrect += tran(rect.bl_corner());
            newrect += tran(rect.br_corner());
            // now make newrect have the same area as the starting rect.
            double ratio = std::sqrt(rect.area()/(double)newrect.area());
            newrect = centered_rect(newrect, newrect.width()*ratio, newrect.height()*ratio);
            metadata.images[i].boxes[j].rect = newrect;

            // rotate 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 = tran(k->second); 
            }
        }

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

    save_image_dataset_metadata(metadata, metadata_filename);
}

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

Davis King's avatar
Davis King committed
336
337
338
339
340
341
342
343
int resample_dataset(const command_line_parser& parser)
{
    if (parser.number_of_arguments() != 1)
    {
        cerr << "The --resample option requires you to give one XML file on the command line." << endl;
        return EXIT_FAILURE;
    }

Davis King's avatar
Davis King committed
344
345
    const size_t obj_size = get_option(parser,"cropped-object-size",100*100); 
    const double margin_scale =  get_option(parser,"crop-size",2.5); // cropped image will be this times wider than the object.
346
    const unsigned long min_object_size = get_option(parser,"min-object-size",1);
347
    const bool one_object_per_image = parser.option("one-object-per-image");
Davis King's avatar
Davis King committed
348
349

    dlib::image_dataset_metadata::dataset data, resampled_data;
Davis King's avatar
Davis King committed
350
351
352
353
354
    std::ostringstream sout;
    sout << "\nThe --resample parameters which generated this dataset were:" << endl;
    sout << "   cropped-object-size: "<< obj_size << endl;
    sout << "   crop-size: "<< margin_scale << endl;
    sout << "   min-object-size: "<< min_object_size << endl;
355
356
    if (one_object_per_image)
        sout << "   one_object_per_image: true" << endl;
Davis King's avatar
Davis King committed
357
    resampled_data.comment = data.comment + sout.str();
Davis King's avatar
Davis King committed
358
359
360
361
    resampled_data.name = data.name + " RESAMPLED";

    load_image_dataset_metadata(data, parser[0]);
    locally_change_current_dir chdir(get_parent_directory(file(parser[0])));
362
    dlib::rand rnd;
Davis King's avatar
Davis King committed
363

364
365
366
    const size_t image_size = std::round(std::sqrt(obj_size*margin_scale*margin_scale));
    const chip_dims cdims(image_size, image_size);

Davis King's avatar
Davis King committed
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
    console_progress_indicator pbar(data.images.size());
    for (unsigned long i = 0; i < data.images.size(); ++i)
    {
        // don't even bother loading images that don't have objects.
        if (data.images[i].boxes.size() == 0)
            continue;

        pbar.print_status(i);
        array2d<rgb_pixel> img, chip;
        load_image(img, data.images[i].filename);


        // figure out what chips we want to take from this image
        for (unsigned long j = 0; j < data.images[i].boxes.size(); ++j)
        {
            const rectangle rect = data.images[i].boxes[j].rect;
383
            if (data.images[i].boxes[j].ignore || rect.area() < min_object_size)
Davis King's avatar
Davis King committed
384
385
                continue;

386
            const auto max_dim = std::max(rect.width(), rect.height());
387

388
            const double rand_scale_perturb = 1 - 0.3*(rnd.get_random_double()-0.5);
389
            const rectangle crop_rect = centered_rect(rect, max_dim*margin_scale*rand_scale_perturb, max_dim*margin_scale*rand_scale_perturb);
Davis King's avatar
Davis King committed
390

391
392
            const rectangle_transform tform = get_mapping_to_chip(chip_details(crop_rect, cdims));
            extract_image_chip(img, chip_details(crop_rect, cdims), chip);
Davis King's avatar
Davis King committed
393
394
395
396
397
398
399
400
401
402
403
404
405

            image_dataset_metadata::image dimg;
            // Now transform the boxes to the crop and also mark them as ignored if they
            // have already been cropped out or are outside the crop.
            for (size_t k = 0; k < data.images[i].boxes.size(); ++k)
            {
                image_dataset_metadata::box box = data.images[i].boxes[k];
                // ignore boxes outside the cropped image
                if (crop_rect.intersect(box.rect).area() == 0)
                    continue;

                // mark boxes we include in the crop as ignored.  Also mark boxes that
                // aren't totally within the crop as ignored.
406
                if (crop_rect.contains(grow_rect(box.rect,10)) && (!one_object_per_image || k==j))
Davis King's avatar
Davis King committed
407
408
409
410
                    data.images[i].boxes[k].ignore = true;
                else
                    box.ignore = true;

411
412
413
                if (box.rect.area() < min_object_size)
                    box.ignore = true;

Davis King's avatar
Davis King committed
414
                box.rect = tform(box.rect);
415
416
                for (auto&& p : box.parts)
                    p.second = tform.get_tform()(p.second);
Davis King's avatar
Davis King committed
417
418
                dimg.boxes.push_back(box);
            }
419
420
421
422
423
            // Put a 64bit hash of the image data into the name to make sure there are no
            // file name conflicts.
            std::ostringstream sout;
            sout << hex << murmur_hash3_128bit(&chip[0][0], chip.size()*sizeof(chip[0][0])).second;
            dimg.filename = data.images[i].filename + "_RESAMPLED_"+sout.str()+".png";
Davis King's avatar
Davis King committed
424

425
426
427
428
429
430
431
432
433
            if (parser.option("jpg"))
            {
                dimg.filename = to_jpg_name(dimg.filename);
                save_jpeg(chip,dimg.filename, JPEG_QUALITY);
            }
            else
            {
                save_png(chip,dimg.filename);
            }
Davis King's avatar
Davis King committed
434
435
436
437
438
439
440
441
442
443
            resampled_data.images.push_back(dimg);
        }
    }

    save_image_dataset_metadata(resampled_data, parser[0] + ".RESAMPLED.xml");

    return EXIT_SUCCESS;
}

// ----------------------------------------------------------------------------------------
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460

int tile_dataset(const command_line_parser& parser)
{
    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;
    }

461
    const unsigned long chip_size = get_option(parser, "size", 8000);
462
463
464
465
466
467
468
469

    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)
    {
Davis King's avatar
Davis King committed
470
471
472
473
        // don't even bother loading images that don't have objects.
        if (data.images[i].boxes.size() == 0)
            continue;

474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
        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;
}


Davis King's avatar
Davis King committed
507
508
// ----------------------------------------------------------------------------------------

509
510
511
512
513
int main(int argc, char** argv)
{
    try
    {

Davis King's avatar
Davis King committed
514
        command_line_parser parser;
515
516

        parser.add_option("h","Displays this information.");
517
518
519
        parser.add_option("v","Display version.");

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

525
        parser.set_group_name("Viewing XML files");
526
        parser.add_option("tile","Chip out all the objects and save them as one big image called <arg>.",1);
527
528
        parser.add_option("size","When using --tile or --cluster, make each extracted object contain "
                                 "about <arg> pixels (default 8000).",1);
529
        parser.add_option("l","List all the labels in the given XML file.");
530
        parser.add_option("stats","List detailed statistics on the object labels in the given XML file.");
Davis King's avatar
Davis King committed
531
        parser.add_option("files","List all the files in the given XML file.");
532

Davis King's avatar
Davis King committed
533
        parser.set_group_name("Editing/Transforming XML datasets");
Davis King's avatar
Davis King committed
534
        parser.add_option("rename", "Rename all labels of <arg1> to <arg2>.",2);
535
        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
536
                          "is defined by <arg> which should be a space separated list of parts.",1);
Davis King's avatar
Davis King committed
537
        parser.add_option("rmempty","Remove all images that don't contain non-ignored annotations and save the results to a new XML file.");
Davis King's avatar
Davis King committed
538
539
        parser.add_option("rmdupes","Remove duplicate images from the dataset.  This is done by comparing "
                                    "the md5 hash of each image file and removing duplicate images. " );
540
541
        parser.add_option("rmdiff","Set the ignored flag to true for boxes marked as difficult.");
        parser.add_option("rmtrunc","Set the ignored flag to true for boxes that are partially outside the image.");
542
        parser.add_option("sort-num-objects","Sort the images listed an XML file so images with many objects are listed first.");
543
        parser.add_option("sort","Alphabetically sort the images in an XML file.");
544
        parser.add_option("shuffle","Randomly shuffle the order of the images listed in an XML file.");
545
546
        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 "
547
            "images with objects labeled <arg> and another file with all the other images. ",1);
548
549
550
551
        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
552
553
        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);
554
555
        parser.add_option("rotate", "Read an XML image dataset and output a copy that is rotated counter clockwise by <arg> degrees. "
                                  "The output is saved to an XML file prefixed with rotated_<arg>.",1);
556
557
        parser.add_option("cluster", "Cluster all the objects in an XML file into <arg> different clusters and save "
                                     "the results as cluster_###.xml and cluster_###.jpg files.",1);
Davis King's avatar
Davis King committed
558
        parser.add_option("ignore", "Mark boxes labeled as <arg> as ignored.  The resulting XML file is output as a separate file and the original is not modified.",1);
559
        parser.add_option("rmlabel","Remove all boxes labeled <arg> and save the results to a new XML file.",1);
560
        parser.add_option("rm-other-labels","Remove all boxes not labeled <arg> and save the results to a new XML file.",1);
Davis King's avatar
Davis King committed
561
        parser.add_option("rmignore","Remove all boxes marked ignore and save the results to a new XML file.");
562
        parser.add_option("rm-if-overlaps","Remove all boxes labeled <arg> if they overlap any box not labeled <arg> and save the results to a new XML file.",1);
Davis King's avatar
Davis King committed
563
        parser.add_option("jpg", "When saving images to disk, write them as jpg files instead of png.");
Davis King's avatar
Davis King committed
564
565
566
567
568

        parser.set_group_name("Cropping sub images");
        parser.add_option("resample", "Crop out images that are centered on each object in the dataset. "
                                      "The output is a new XML dataset."); 
        parser.add_option("cropped-object-size", "When doing --resample, make the cropped objects contain about <arg> pixels (default 10000).",1);
569
        parser.add_option("min-object-size", "When doing --resample, skip objects that have fewer than <arg> pixels in them (default 1).",1);
Davis King's avatar
Davis King committed
570
        parser.add_option("crop-size", "When doing --resample, the entire cropped image will be <arg> times wider than the object (default 2.5).",1); 
571
        parser.add_option("one-object-per-image", "When doing --resample, only include one non-ignored object per image (i.e. the central object).");
Davis King's avatar
Davis King committed
572
573


574
575
576

        parser.parse(argc, argv);

Davis King's avatar
Davis King committed
577
        const char* singles[] = {"h","c","r","l","files","convert","parts","rmdiff", "rmtrunc", "rmdupes", "seed", "shuffle", "split", "add", 
Davis King's avatar
Davis King committed
578
                                 "flip", "rotate", "tile", "size", "cluster", "resample", "min-object-size", "rmempty",
579
580
                                 "crop-size", "cropped-object-size", "rmlabel", "rm-other-labels", "rm-if-overlaps", "sort-num-objects", 
                                 "one-object-per-image", "jpg", "rmignore", "sort"};
581
        parser.check_one_time_options(singles);
582
583
        const char* c_sub_ops[] = {"r", "convert"};
        parser.check_sub_options("c", c_sub_ops);
584
        parser.check_sub_option("shuffle", "seed");
585
        const char* resample_sub_ops[] = {"min-object-size", "crop-size", "cropped-object-size", "one-object-per-image"};
Davis King's avatar
Davis King committed
586
        parser.check_sub_options("resample", resample_sub_ops);
587
588
        const char* size_parent_ops[] = {"tile", "cluster"};
        parser.check_sub_options(size_parent_ops, "size");
Davis King's avatar
Davis King committed
589
        parser.check_incompatible_options("c", "l");
Davis King's avatar
Davis King committed
590
        parser.check_incompatible_options("c", "files");
591
        parser.check_incompatible_options("c", "rmdiff");
Davis King's avatar
Davis King committed
592
        parser.check_incompatible_options("c", "rmempty");
593
        parser.check_incompatible_options("c", "rmlabel");
594
        parser.check_incompatible_options("c", "rm-other-labels");
Davis King's avatar
Davis King committed
595
        parser.check_incompatible_options("c", "rmignore");
596
        parser.check_incompatible_options("c", "rm-if-overlaps");
Davis King's avatar
Davis King committed
597
        parser.check_incompatible_options("c", "rmdupes");
598
        parser.check_incompatible_options("c", "rmtrunc");
599
        parser.check_incompatible_options("c", "add");
Davis King's avatar
Davis King committed
600
        parser.check_incompatible_options("c", "flip");
601
        parser.check_incompatible_options("c", "rotate");
Davis King's avatar
Davis King committed
602
        parser.check_incompatible_options("c", "rename");
603
        parser.check_incompatible_options("c", "ignore");
604
        parser.check_incompatible_options("c", "parts");
605
        parser.check_incompatible_options("c", "tile");
606
        parser.check_incompatible_options("c", "cluster");
Davis King's avatar
Davis King committed
607
        parser.check_incompatible_options("c", "resample");
Davis King's avatar
Davis King committed
608
        parser.check_incompatible_options("l", "rename");
609
        parser.check_incompatible_options("l", "ignore");
610
        parser.check_incompatible_options("l", "add");
611
        parser.check_incompatible_options("l", "parts");
Davis King's avatar
Davis King committed
612
        parser.check_incompatible_options("l", "flip");
613
        parser.check_incompatible_options("l", "rotate");
Davis King's avatar
Davis King committed
614
615
616
617
618
619
        parser.check_incompatible_options("files", "rename");
        parser.check_incompatible_options("files", "ignore");
        parser.check_incompatible_options("files", "add");
        parser.check_incompatible_options("files", "parts");
        parser.check_incompatible_options("files", "flip");
        parser.check_incompatible_options("files", "rotate");
Davis King's avatar
Davis King committed
620
        parser.check_incompatible_options("add", "flip");
621
        parser.check_incompatible_options("add", "rotate");
622
623
        parser.check_incompatible_options("add", "tile");
        parser.check_incompatible_options("flip", "tile");
624
        parser.check_incompatible_options("rotate", "tile");
625
        parser.check_incompatible_options("cluster", "tile");
Davis King's avatar
Davis King committed
626
        parser.check_incompatible_options("resample", "tile");
627
        parser.check_incompatible_options("flip", "cluster");
628
        parser.check_incompatible_options("rotate", "cluster");
629
        parser.check_incompatible_options("add", "cluster");
Davis King's avatar
Davis King committed
630
631
632
        parser.check_incompatible_options("flip", "resample");
        parser.check_incompatible_options("rotate", "resample");
        parser.check_incompatible_options("add", "resample");
633
        parser.check_incompatible_options("shuffle", "tile");
634
        parser.check_incompatible_options("sort-num-objects", "tile");
635
        parser.check_incompatible_options("sort", "tile");
636
        parser.check_incompatible_options("convert", "l");
Davis King's avatar
Davis King committed
637
        parser.check_incompatible_options("convert", "files");
638
        parser.check_incompatible_options("convert", "rename");
639
        parser.check_incompatible_options("convert", "ignore");
640
        parser.check_incompatible_options("convert", "parts");
641
        parser.check_incompatible_options("convert", "cluster");
Davis King's avatar
Davis King committed
642
        parser.check_incompatible_options("convert", "resample");
643
        parser.check_incompatible_options("rmdiff", "rename");
644
        parser.check_incompatible_options("rmdiff", "ignore");
Davis King's avatar
Davis King committed
645
646
        parser.check_incompatible_options("rmempty", "ignore");
        parser.check_incompatible_options("rmempty", "rename");
647
648
        parser.check_incompatible_options("rmlabel", "ignore");
        parser.check_incompatible_options("rmlabel", "rename");
649
650
        parser.check_incompatible_options("rm-other-labels", "ignore");
        parser.check_incompatible_options("rm-other-labels", "rename");
Davis King's avatar
Davis King committed
651
652
        parser.check_incompatible_options("rmignore", "ignore");
        parser.check_incompatible_options("rmignore", "rename");
653
654
        parser.check_incompatible_options("rm-if-overlaps", "ignore");
        parser.check_incompatible_options("rm-if-overlaps", "rename");
Davis King's avatar
Davis King committed
655
        parser.check_incompatible_options("rmdupes", "rename");
656
        parser.check_incompatible_options("rmdupes", "ignore");
657
        parser.check_incompatible_options("rmtrunc", "rename");
658
        parser.check_incompatible_options("rmtrunc", "ignore");
659
        const char* convert_args[] = {"pascal-xml","pascal-v1","idl"};
660
        parser.check_option_arg_range("convert", convert_args);
661
        parser.check_option_arg_range("cluster", 2, 999);
Davis King's avatar
Davis King committed
662
        parser.check_option_arg_range("rotate", -360, 360);
663
        parser.check_option_arg_range("size", 10*10, 1000*1000);
664
        parser.check_option_arg_range("min-object-size", 1, 10000*10000);
Davis King's avatar
Davis King committed
665
666
        parser.check_option_arg_range("cropped-object-size", 4, 10000*10000);
        parser.check_option_arg_range("crop-size", 1.0, 100.0);
667
668
669

        if (parser.option("h"))
        {
670
            cout << "Usage: imglab [options] <image files/directories or XML file>\n";
671
            parser.print_options(cout);
Davis King's avatar
Davis King committed
672
            cout << endl << endl;
673
674
675
            return EXIT_SUCCESS;
        }

676
677
678
679
680
681
        if (parser.option("add"))
        {
            merge_metadata_files(parser);
            return EXIT_SUCCESS;
        }

Davis King's avatar
Davis King committed
682
683
684
685
686
687
        if (parser.option("flip"))
        {
            flip_dataset(parser);
            return EXIT_SUCCESS;
        }

688
689
690
691
692
693
        if (parser.option("rotate"))
        {
            rotate_dataset(parser);
            return EXIT_SUCCESS;
        }

Davis King's avatar
Davis King committed
694
695
696
697
698
699
700
701
702
        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;
        }

703
704
705
706
707
        if (parser.option("tile"))
        {
            return tile_dataset(parser);
        }

708
709
710
711
712
        if (parser.option("cluster"))
        {
            return cluster_dataset(parser);
        }

Davis King's avatar
Davis King committed
713
714
715
716
717
        if (parser.option("resample"))
        {
            return resample_dataset(parser);
        }

718
719
        if (parser.option("c"))
        {
720
721
            if (parser.option("convert"))
            {
Davis King's avatar
Davis King committed
722
723
                if (parser.option("convert").argument() == "pascal-xml")
                    convert_pascal_xml(parser);
724
725
                else if (parser.option("convert").argument() == "pascal-v1")
                    convert_pascal_v1(parser);
726
727
                else if (parser.option("convert").argument() == "idl")
                    convert_idl(parser);
728
729
730
731
732
            }
            else
            {
                create_new_dataset(parser);
            }
733
734
            return EXIT_SUCCESS;
        }
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
        
        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)
            {
                for (unsigned long j = 0; j < data.images[i].boxes.size(); ++j)
                {
750
751
752
753
754
755
756
757
                    if (data.images[i].boxes[j].difficult)
                        data.images[i].boxes[j].ignore = true;
                }
            }
            save_image_dataset_metadata(data, parser[0]);
            return EXIT_SUCCESS;
        }

Davis King's avatar
Davis King committed
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
        if (parser.option("rmempty"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --rmempty option requires you to give one XML file on the command line." << endl;
                return EXIT_FAILURE;
            }

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

            data2 = data;
            data2.images.clear();
            for (unsigned long i = 0; i < data.images.size(); ++i)
            {
                bool has_label = false;
                for (unsigned long j = 0; j < data.images[i].boxes.size(); ++j)
                {
                    if (!data.images[i].boxes[j].ignore)
                        has_label = true;
                }
                if (has_label)
                    data2.images.push_back(data.images[i]);
            }
            save_image_dataset_metadata(data2, parser[0] + ".rmempty.xml");
            return EXIT_SUCCESS;
        }

786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
        if (parser.option("rmlabel"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --rmlabel 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 auto label = parser.option("rmlabel").argument();

            for (auto&& img : data.images)
            {
                std::vector<dlib::image_dataset_metadata::box> boxes;
                for (auto&& b : img.boxes)
                {
                    if (b.label != label)
                        boxes.push_back(b);
                }
                img.boxes = boxes;
            }

            save_image_dataset_metadata(data, parser[0] + ".rmlabel-"+label+".xml");
            return EXIT_SUCCESS;
        }

814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
        if (parser.option("rm-other-labels"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --rm-other-labels 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 auto labels = parser.option("rm-other-labels").argument();
            // replace comma by dash to form the file name
            std::string strlabels = labels;
            std::replace(strlabels.begin(), strlabels.end(), ',', '-');
            std::vector<string> all_labels = split(labels, ",");
            for (auto&& img : data.images)
            {
                std::vector<dlib::image_dataset_metadata::box> boxes;
                for (auto&& b : img.boxes)
                {
                    if (std::find(all_labels.begin(), all_labels.end(), b.label) != all_labels.end())
                        boxes.push_back(b);
                }
                img.boxes = boxes;
            }

            save_image_dataset_metadata(data, parser[0] + ".rm-other-labels-"+ strlabels +".xml");
            return EXIT_SUCCESS;
        }

Davis King's avatar
Davis King committed
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
        if (parser.option("rmignore"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --rmignore 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 (auto&& img : data.images)
            {
                std::vector<dlib::image_dataset_metadata::box> boxes;
                for (auto&& b : img.boxes)
                {
                    if (!b.ignore)
                        boxes.push_back(b);
                }
                img.boxes = boxes;
            }

            save_image_dataset_metadata(data, parser[0] + ".rmignore.xml");
            return EXIT_SUCCESS;
        }

871
872
873
874
875
876
877
878
879
880
881
882
883
        if (parser.option("rm-if-overlaps"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --rm-if-overlaps 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 auto label = parser.option("rm-if-overlaps").argument();

884
            test_box_overlap overlaps(0.5);
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916

            for (auto&& img : data.images)
            {
                std::vector<dlib::image_dataset_metadata::box> boxes;
                for (auto&& b : img.boxes)
                {
                    if (b.label != label)
                    {
                        boxes.push_back(b);
                    }
                    else
                    {
                        bool has_overlap = false;
                        for (auto&& b2 : img.boxes)
                        {
                            if (b2.label != label && overlaps(b2.rect, b.rect))
                            {
                                has_overlap = true;
                                break;
                            }
                        }
                        if (!has_overlap)
                            boxes.push_back(b);
                    }
                }
                img.boxes = boxes;
            }

            save_image_dataset_metadata(data, parser[0] + ".rm-if-overlaps-"+label+".xml");
            return EXIT_SUCCESS;
        }

Davis King's avatar
Davis King committed
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
        if (parser.option("rmdupes"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --rmdupes option requires you to give one XML file on the command line." << endl;
                return EXIT_FAILURE;
            }

            dlib::image_dataset_metadata::dataset data, data_out;
            std::set<std::string> hashes;
            load_image_dataset_metadata(data, parser[0]);
            data_out = data;
            data_out.images.clear();

            for (unsigned long i = 0; i < data.images.size(); ++i)
            {
                ifstream fin(data.images[i].filename.c_str(), ios::binary);
                string hash = md5(fin);
                if (hashes.count(hash) == 0)
                {
                    hashes.insert(hash);
                    data_out.images.push_back(data.images[i]);
                }
            }
            save_image_dataset_metadata(data_out, parser[0]);
            return EXIT_SUCCESS;
        }

945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
        if (parser.option("rmtrunc"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --rmtrunc 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]);
            {
                locally_change_current_dir chdir(get_parent_directory(file(parser[0])));
                for (unsigned long i = 0; i < data.images.size(); ++i)
                {
                    array2d<unsigned char> img;
                    load_image(img, data.images[i].filename);
                    const rectangle area = get_rect(img);
                    for (unsigned long j = 0; j < data.images[i].boxes.size(); ++j)
                    {
                        if (!area.contains(data.images[i].boxes[j].rect))
                            data.images[i].boxes[j].ignore = true;
                    }
967
968
969
970
971
                }
            }
            save_image_dataset_metadata(data, parser[0]);
            return EXIT_SUCCESS;
        }
972

973
974
975
976
977
978
979
980
981
982
983
984
985
986
        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
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
        if (parser.option("files"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --files 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 (size_t i = 0; i < data.images.size(); ++i)
                cout << data.images[i].filename << "\n";
            return EXIT_SUCCESS;
        }

1002
1003
1004
1005
1006
1007
1008
1009
1010
        if (parser.option("split"))
        {
            return split_dataset(parser);
        }

        if (parser.option("shuffle"))
        {
            if (parser.number_of_arguments() != 1)
            {
1011
                cerr << "The --shuffle option requires you to give one XML file on the command line." << endl;
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
                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;
        }

1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
        if (parser.option("sort-num-objects"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --sort-num-objects 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]);
            std::sort(data.images.rbegin(),  data.images.rend(), 
                [](const image_dataset_metadata::image& a, const image_dataset_metadata::image& b) { return a.boxes.size() < b.boxes.size(); });
            save_image_dataset_metadata(data, parser[0]);
            return EXIT_SUCCESS;
        }

1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
        if (parser.option("sort"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --sort 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]);
            std::sort(data.images.begin(),  data.images.end(), 
                [](const image_dataset_metadata::image& a, const image_dataset_metadata::image& b) { return a.filename < b.filename; });
            save_image_dataset_metadata(data, parser[0]);
            return EXIT_SUCCESS;
        }

1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
        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
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
        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;
        }

1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
        if (parser.option("ignore"))
        {
            if (parser.number_of_arguments() != 1)
            {
                cerr << "The --ignore 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("ignore").count(); ++i)
            {
                ignore_labels(data, parser.option("ignore").argument());
            }
            save_image_dataset_metadata(data, parser[0]+".ignored.xml");
            return EXIT_SUCCESS;
        }

1107
1108
        if (parser.number_of_arguments() == 1)
        {
1109
            metadata_editor editor(parser[0]);
1110
1111
1112
1113
1114
1115
1116
1117
            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]);
                }
            }
1118
            editor.wait_until_closed();
1119
            return EXIT_SUCCESS;
1120
        }
Davis King's avatar
Davis King committed
1121
1122
1123

        cout << "Invalid command, give -h to see options." << endl;
        return EXIT_FAILURE;
1124
1125
1126
    }
    catch (exception& e)
    {
1127
        cerr << e.what() << endl;
1128
1129
1130
1131
        return EXIT_FAILURE;
    }
}

Davis King's avatar
Davis King committed
1132
1133
// ----------------------------------------------------------------------------------------