(args)
| 166 | |
| 167 | |
| 168 | def main(args): |
| 169 | if args.output: |
| 170 | output_dir = os.path.realpath(args.output) |
| 171 | os.makedirs(output_dir, exist_ok=True) |
| 172 | |
| 173 | labels = [] |
| 174 | if args.labels: |
| 175 | with open(args.labels) as f: |
| 176 | for i, label in enumerate(f): |
| 177 | labels.append(label.strip()) |
| 178 | |
| 179 | trt_infer = TensorRTInfer(args.engine) |
| 180 | if args.input: |
| 181 | print("Inferring data in {}".format(args.input)) |
| 182 | batcher = ImageBatcher(args.input, *trt_infer.input_spec()) |
| 183 | for batch, images, scales in batcher.get_batch(): |
| 184 | print("Processing Image {} / {}".format(batcher.image_index, batcher.num_images), end="\r") |
| 185 | detections = trt_infer.process(batch, scales, args.nms_threshold) |
| 186 | if args.output: |
| 187 | for i in range(len(images)): |
| 188 | basename = os.path.splitext(os.path.basename(images[i]))[0] |
| 189 | # Image Visualizations |
| 190 | output_path = os.path.join(output_dir, "{}.png".format(basename)) |
| 191 | visualize_detections(images[i], output_path, detections[i], labels) |
| 192 | # Text Results |
| 193 | output_results = "" |
| 194 | for d in detections[i]: |
| 195 | line = [d['xmin'], d['ymin'], d['xmax'], d['ymax'], d['score'], d['class']] |
| 196 | output_results += "\t".join([str(f) for f in line]) + "\n" |
| 197 | with open(os.path.join(output_dir, "{}.txt".format(basename)), "w") as f: |
| 198 | f.write(output_results) |
| 199 | else: |
| 200 | print("No input provided, running in benchmark mode") |
| 201 | spec = trt_infer.input_spec() |
| 202 | batch = 255 * np.random.rand(*spec[0]).astype(spec[1]) |
| 203 | iterations = 200 |
| 204 | times = [] |
| 205 | for i in range(20): # GPU warmup iterations |
| 206 | trt_infer.infer(batch) |
| 207 | for i in range(iterations): |
| 208 | start = time.time() |
| 209 | trt_infer.infer(batch) |
| 210 | times.append(time.time() - start) |
| 211 | print("Iteration {} / {}".format(i + 1, iterations), end="\r") |
| 212 | print("Benchmark results include time for H2D and D2H memory copies") |
| 213 | print("Average Latency: {:.3f} ms".format( |
| 214 | 1000 * np.average(times))) |
| 215 | print("Average Throughput: {:.1f} ips".format( |
| 216 | trt_infer.batch_size / np.average(times))) |
| 217 | |
| 218 | print() |
| 219 | print("Finished Processing") |
| 220 | |
| 221 | |
| 222 | if __name__ == "__main__": |
no test coverage detected