(args)
| 56 | return cfg |
| 57 | |
| 58 | def main(args): |
| 59 | # Set up Detectron 2 config and build evaluator. |
| 60 | cfg = setup(args.det2_config, args.det2_weights) |
| 61 | dataset_name = cfg.DATASETS.TEST[0] |
| 62 | evaluator = build_evaluator(dataset_name) |
| 63 | evaluator.reset() |
| 64 | |
| 65 | trt_infer = TensorRTInfer(args.engine) |
| 66 | batcher = ImageBatcher(args.input, *trt_infer.input_spec(), config_file=args.det2_config) |
| 67 | |
| 68 | for batch, images, scales in batcher.get_batch(): |
| 69 | print("Processing Image {} / {}".format(batcher.image_index, batcher.num_images), end="\r") |
| 70 | detections = trt_infer.infer(batch, scales, args.nms_threshold) |
| 71 | for i in range(len(images)): |
| 72 | # Get inference image resolution. |
| 73 | infer_im = Image.open(images[i]) |
| 74 | im_width, im_height = infer_im.size |
| 75 | pred_boxes = [] |
| 76 | scores = [] |
| 77 | pred_classes = [] |
| 78 | # Number of detections. |
| 79 | num_instances = len(detections[i]) |
| 80 | # Reserve numpy array to hold all mask predictions per image. |
| 81 | pred_masks = np.empty((num_instances, 28, 28), dtype=np.float32) |
| 82 | # Image ID, required for Detectron 2 evaluations. |
| 83 | source_id = int(os.path.splitext(os.path.basename(images[i]))[0]) |
| 84 | # Loop over every single detection. |
| 85 | for n in range(num_instances): |
| 86 | det = detections[i][n] |
| 87 | # Append box coordinates data. |
| 88 | pred_boxes.append([det['ymin'], det['xmin'], det['ymax'], det['xmax']]) |
| 89 | # Append score. |
| 90 | scores.append(det['score']) |
| 91 | # Append class. |
| 92 | pred_classes.append(det['class']) |
| 93 | # Append mask. |
| 94 | pred_masks[n] = det['mask'] |
| 95 | # Create new Instances object required for Detectron 2 evalutions and add: |
| 96 | # boxes, scores, pred_classes, pred_masks. |
| 97 | image_shape = (im_height, im_width) |
| 98 | instances = Instances(image_shape) |
| 99 | instances.pred_boxes = Boxes(pred_boxes) |
| 100 | instances.scores = torch.tensor(scores) |
| 101 | instances.pred_classes = torch.tensor(pred_classes) |
| 102 | roi_masks = ROIMasks(torch.tensor(pred_masks)) |
| 103 | instances.pred_masks = roi_masks.to_bitmasks(instances.pred_boxes, im_height, im_width, args.iou_threshold).tensor |
| 104 | # Process evaluations per image. |
| 105 | image_dict = [{'instances': instances}] |
| 106 | input_dict = [{'image_id': source_id}] |
| 107 | evaluator.process(input_dict, image_dict) |
| 108 | |
| 109 | # Final evaluations, generation of mAP accuracy performance. |
| 110 | evaluator.evaluate() |
| 111 | |
| 112 | |
| 113 | if __name__ == "__main__": |
no test coverage detected