(tf_images, tf_detections, trt_images, trt_detections, output_dir, annotations_path, labels_path)
| 63 | |
| 64 | |
| 65 | def compare_images(tf_images, tf_detections, trt_images, trt_detections, output_dir, annotations_path, labels_path): |
| 66 | labels = [] |
| 67 | if labels_path and os.path.exists(labels_path): |
| 68 | with open(labels_path) as f: |
| 69 | for i, label in enumerate(f): |
| 70 | labels.append(label.strip()) |
| 71 | |
| 72 | annotations = parse_annotations(annotations_path) |
| 73 | |
| 74 | count = 1 |
| 75 | for tf_img, tf_det, trt_img, trt_det in zip(tf_images, tf_detections, trt_images, trt_detections): |
| 76 | vis = [] |
| 77 | names = [] |
| 78 | colors = [] |
| 79 | |
| 80 | vis.append(visualize_detections(tf_img, None, tf_det, labels)) |
| 81 | names.append("TensorFlow") |
| 82 | colors.append("DarkOrange") |
| 83 | |
| 84 | vis.append(visualize_detections(trt_img, None, trt_det, labels)) |
| 85 | names.append("TensorRT") |
| 86 | colors.append("YellowGreen") |
| 87 | |
| 88 | if annotations: |
| 89 | img_id = os.path.splitext(os.path.basename(trt_img))[0] |
| 90 | if img_id.isnumeric(): |
| 91 | img_id = int(img_id) |
| 92 | if img_id in annotations.keys(): |
| 93 | vis.append(visualize_detections(trt_img, None, annotations[img_id], labels)) |
| 94 | names.append("Ground Truth") |
| 95 | colors.append("RoyalBlue") |
| 96 | else: |
| 97 | print("Image {} does not have a COCO annotation, skipping ground truth visualization".format(trt_img)) |
| 98 | |
| 99 | basename = os.path.splitext(os.path.basename(tf_img))[0] |
| 100 | output_path = os.path.join(output_dir, "{}.compare.png".format(basename)) |
| 101 | os.makedirs(output_dir, exist_ok=True) |
| 102 | concat_visualizations(vis, names, colors, output_path) |
| 103 | |
| 104 | print("Processing {} / {} images (Visualization)".format(count, len(tf_images)), end="\r") |
| 105 | count += 1 |
| 106 | print() |
| 107 | |
| 108 | |
| 109 | def main(args): |
no test coverage detected