| 124 | |
| 125 | |
| 126 | def _parseArgs(): |
| 127 | parser = argparse.ArgumentParser(description="Downloader of TensorRT sample data files.") |
| 128 | parser.add_argument( |
| 129 | "-d", |
| 130 | "--data", |
| 131 | help="Specify the data directory, data will be downloaded to there. $TRT_DATA_DIR will be overwritten by this argument.", |
| 132 | ) |
| 133 | parser.add_argument( |
| 134 | "-f", |
| 135 | "--file", |
| 136 | help="Specify the path to the download.yml, default to `download.yml` in the working directory", |
| 137 | default="download.yml", |
| 138 | ) |
| 139 | parser.add_argument( |
| 140 | "-o", "--overwrite", help="Force to overwrite if MD5 check failed", action="store_true", default=False |
| 141 | ) |
| 142 | parser.add_argument( |
| 143 | "-v", |
| 144 | "--verify", |
| 145 | help="Verify if the data has been downloaded. Will not download if specified.", |
| 146 | action="store_true", |
| 147 | default=False, |
| 148 | ) |
| 149 | |
| 150 | args, _ = parser.parse_known_args() |
| 151 | data = os.environ.get("TRT_DATA_DIR", None) if args.data is None else args.data |
| 152 | if data is None: |
| 153 | raise ValueError("Data directory must be specified by either `-d $DATA` or environment variable $TRT_DATA_DIR.") |
| 154 | |
| 155 | return data, args |
| 156 | |
| 157 | |
| 158 | def verifyChecksum(data_dir, yaml_path): |