def test_download_pretrained_model():
    """Test download_pretrained_model()."""
    test_base_models = load_base_models_json()
    TEST_TRAINING_NAME = "luke_burnt"
    TEST_CHOSEN_MODEL = "ssd_mobilenet_v2"
    test_filepaths = set_filenames(test_base_models, TEST_TRAINING_NAME,
                                   TEST_CHOSEN_MODEL)
    download_pretrained_model(test_filepaths)
    assert os.path.isfile("/tf/models/research/deploy/" +
                          test_filepaths["pretrained_checkpoint"])
def test_download_base_training_config():
    """Test download_base_training_config()."""
    # pylint: disable=line-too-long
    test_base_models = load_base_models_json()
    TEST_TRAINING_NAME = "luke_burnt"
    TEST_CHOSEN_MODEL = "ssd_mobilenet_v2"
    test_filepaths = set_filenames(test_base_models, TEST_TRAINING_NAME,
                                   TEST_CHOSEN_MODEL)
    download_base_training_config(test_filepaths)
    assert os.path.isfile("/tf/models/research/deploy/" +
                          test_filepaths["base_pipeline_file"])
def test_set_filenames():
    """Test set_filenames()"""
    test_base_models = load_base_models_json()
    TEST_TRAINING_NAME = "luke_burnt"
    TEST_CHOSEN_MODEL = "ssd_mobilenet_v2"
    test_filepaths = set_filenames(test_base_models, TEST_TRAINING_NAME,
                                   TEST_CHOSEN_MODEL)
    assert test_filepaths[
        "val_export_dir"] == "/tf/dataset-export/luke_burnt/val/"
    assert (test_filepaths["pretrained_checkpoint"] ==
            "ssd_mobilenet_v2_320x320_coco17_tpu-8.tar.gz")
    assert (test_filepaths["pipeline_file"] ==
            "/tf/dataset-export/luke_burnt/pipeline_file.config")
def test_export_voxel51_dataset_to_tfrecords():
    """Test export_voxel51_dataset_to_tfrecords()."""
    # needed for set up
    create_voxel51_dataset("test")
    test_base_models = load_base_models_json()
    TEST_TRAINING_NAME = "luke_burnt"
    TEST_CHOSEN_MODEL = "ssd_mobilenet_v2"
    test_filepaths = set_filenames(test_base_models, TEST_TRAINING_NAME,
                                   TEST_CHOSEN_MODEL)

    export_voxel51_dataset_to_tfrecords("test",
                                        test_filepaths,
                                        label_field="detections",
                                        training_tag="train")
def test_save_mapping_to_file():
    """Test save_mapping_to_file()."""
    test_base_models = load_base_models_json()
    TEST_TRAINING_NAME = "luke_burnt"
    TEST_CHOSEN_MODEL = "ssd_mobilenet_v2"
    test_filepaths = set_filenames(test_base_models, TEST_TRAINING_NAME,
                                   TEST_CHOSEN_MODEL)
    test_mapping = str(
        'item {\n  name: "friend"\n  id: 1\n}\nitem {\n  name: "stopper"\n  id: 2\n}\n'
    )

    save_mapping_to_file(test_mapping, test_filepaths)

    with open(test_filepaths["label_map_file"], "r") as test_file:
        assert (
            test_file.read() ==
            'item {\n  name: "friend"\n  id: 1\n}\nitem {\n  name: "stopper"\n  id: 2\n}\n'
        )