Create toy pascal dataset

create_toy_pascal.py

import numpy as np

OLD_CLASSES = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car',
               'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
               'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']

NEW_CLASSES = ['bicycle', 'bus', 'car', 'cat', 'cow', 'dog', 'horse', 'motorbike',
               'person', 'train']

OLD_INDICES = [OLD_CLASSES.index(x) for x in NEW_CLASSES]
AREA_LIMIT = 0.15

def filter_images(lst_file, new_file):
    with open(lst_file, 'r') as f:
        out_list = []
        count = 0
        for line in f.readlines():
            l = line.split()
            label = l[1:-1]
            path = l[-1]
            head_width = int(label[0])
            obj_width = int(label[1])
            objects = np.reshape(np.array(label[head_width:]).astype(float), (-1, obj_width))
            new_label = []
            keep = True
            for k, obj in enumerate(objects):
                if int(obj[0]) in OLD_INDICES:
                    width = obj[3] - obj[1]
                    height = obj[4] - obj[2]
                    if width * height < AREA_LIMIT:
                        new_label = []
                        keep = False
                        break
                    else:
                        obj[0] = OLD_INDICES.index(int(obj[0]))
                        new_label.append(obj)
            if keep and new_label:
                new_label = np.array(new_label)
                new_line = [count] + label[:head_width] + objects.ravel().tolist() + [path]
                new_line = '\t'.join([str(x) for x in new_line])
                out_list.append(new_line)
                count += 1
        print(new_file, 'count: ', count)

    if out_list:
        with open(new_file, 'w') as f:
            for line in out_list:
                f.write(line + '\n')

def save_syntxt(fname):
    with open(fname, 'w') as f:
        for c in NEW_CLASSES:
            f.write(c + '\n')


if __name__ == '__main__':
    filter_images('val.lst', 'val2.lst')
    filter_images('train.lst', 'train2.lst')
    save_syntxt('names.txt')