(args)
| 26 | |
| 27 | |
| 28 | def main(args): |
| 29 | automl_path = os.path.realpath(args.automl_path) |
| 30 | sys.path.insert(1, os.path.join(automl_path, "efficientdet")) |
| 31 | try: |
| 32 | import coco_metric |
| 33 | except ImportError: |
| 34 | print("Could not import the 'coco_metric' module from AutoML. Searching in: {}".format(automl_path)) |
| 35 | print("Please clone the repository https://github.com/google/automl and provide its path with --automl_path.") |
| 36 | sys.exit(1) |
| 37 | |
| 38 | trt_infer = TensorRTInfer(args.engine) |
| 39 | batcher = ImageBatcher(args.input, *trt_infer.input_spec()) |
| 40 | evaluator = coco_metric.EvaluationMetric(filename=args.annotations) |
| 41 | for batch, images, scales in batcher.get_batch(): |
| 42 | print("Processing Image {} / {}".format(batcher.image_index, batcher.num_images), end="\r") |
| 43 | detections = trt_infer.process(batch, scales, args.nms_threshold) |
| 44 | coco_det = np.zeros((len(images), max([len(d) for d in detections]), 7)) |
| 45 | coco_det[:, :, -1] = -1 |
| 46 | for i in range(len(images)): |
| 47 | for n in range(len(detections[i])): |
| 48 | source_id = int(os.path.splitext(os.path.basename(images[i]))[0]) |
| 49 | det = detections[i][n] |
| 50 | coco_det[i][n] = [ |
| 51 | source_id, |
| 52 | det["xmin"], |
| 53 | det["ymin"], |
| 54 | det["xmax"] - det["xmin"], |
| 55 | det["ymax"] - det["ymin"], |
| 56 | det["score"], |
| 57 | det["class"] + 1, # The COCO evaluator expects class 0 to be background, so offset by 1 |
| 58 | ] |
| 59 | evaluator.update_state(None, coco_det) |
| 60 | print() |
| 61 | evaluator.result(100) |
| 62 | |
| 63 | |
| 64 | if __name__ == "__main__": |
no test coverage detected