data_loader.py 文件源码

def load_data(directory, data_processor_module):
    file_list = os.listdir(directory)
    data_set = DataSet(data_processor_module)

    for i, f in enumerate(file_list):
        full_filename = os.path.join(directory, f)
        img = cv2.imread(full_filename)
        if img is None:
            print("WARNING: File: '{}' could not be loaded".format(full_filename))
            continue
        # The files will be of the type:
        # CLASS-source-frame-itemnumber.jpg
        label = f.split("-")[0].lower()
        if label == "nao_part" or label == "nothing" and random.random() > 0.2:
            # Because we have too many nao_parts, remove a lot of them
            continue
        # if label != "nao_part":
        data_set.add_image(img, label)

        if i % 100:
            print("{}/{} - {:5.2f}%".format(i, len(file_list), i*100./len(file_list)), end="\r")

    print("                        ", end="\r")

    if not data_set.loaded_images:
        raise ValueError("Could not load any images.")

    data_set.process()

    return data_set