(args)
| 182 | |
| 183 | |
| 184 | def main(args): |
| 185 | output_dir = os.path.realpath(args.output) |
| 186 | os.makedirs(output_dir, exist_ok=True) |
| 187 | |
| 188 | labels = [] |
| 189 | if args.labels: |
| 190 | with open(args.labels) as f: |
| 191 | for i, label in enumerate(f): |
| 192 | labels.append(label.strip()) |
| 193 | |
| 194 | trt_infer = TensorRTInfer(args.engine, args.preprocessor, args.detection_type, args.iou_threshold) |
| 195 | batcher = ImageBatcher(args.input, *trt_infer.input_spec(), preprocessor=args.preprocessor) |
| 196 | for batch, images, scales in batcher.get_batch(): |
| 197 | print("Processing Image {} / {}".format(batcher.image_index, batcher.num_images), end="\r") |
| 198 | detections = trt_infer.infer(batch, scales, args.nms_threshold) |
| 199 | for i in range(len(images)): |
| 200 | basename = os.path.splitext(os.path.basename(images[i]))[0] |
| 201 | # Image Visualizations |
| 202 | output_path = os.path.join(output_dir, "{}.png".format(basename)) |
| 203 | visualize_detections(images[i], output_path, detections[i], labels) |
| 204 | # Text Results |
| 205 | output_results = "" |
| 206 | for d in detections[i]: |
| 207 | line = [d['xmin'], d['ymin'], d['xmax'], d['ymax'], d['score'], d['class']] |
| 208 | output_results += "\t".join([str(f) for f in line]) + "\n" |
| 209 | with open(os.path.join(args.output, "{}.txt".format(basename)), "w") as f: |
| 210 | f.write(output_results) |
| 211 | print() |
| 212 | print("Finished Processing") |
| 213 | |
| 214 | |
| 215 | if __name__ == "__main__": |
no test coverage detected