Adamfan/objectdetection

1import numpy as np2import onnx3from onnx import shape_inference4try:5    import onnx_graphsurgeon as gs6except Exception as e:7    print('Import onnx_graphsurgeon failure: %s' % e)8 9import logging10 11LOGGER = logging.getLogger(__name__)12 13class RegisterNMS(object):14    def __init__(15        self,16        onnx_model_path: str,17        precision: str = "fp32",18    ):19 20        self.graph = gs.import_onnx(onnx.load(onnx_model_path))21        assert self.graph22        LOGGER.info("ONNX graph created successfully")23        # Fold constants via ONNX-GS that PyTorch2ONNX may have missed24        self.graph.fold_constants()25        self.precision = precision26        self.batch_size = 127    def infer(self):28        """29        Sanitize the graph by cleaning any unconnected nodes, do a topological resort,30        and fold constant inputs values. When possible, run shape inference on the31        ONNX graph to determine tensor shapes.32        """33        for _ in range(3):34            count_before = len(self.graph.nodes)35 36            self.graph.cleanup().toposort()37            try:38                for node in self.graph.nodes:39                    for o in node.outputs:40                        o.shape = None41                model = gs.export_onnx(self.graph)42                model = shape_inference.infer_shapes(model)43                self.graph = gs.import_onnx(model)44            except Exception as e:45                LOGGER.info(f"Shape inference could not be performed at this time:\n{e}")46            try:47                self.graph.fold_constants(fold_shapes=True)48            except TypeError as e:49                LOGGER.error(50                    "This version of ONNX GraphSurgeon does not support folding shapes, "51                    f"please upgrade your onnx_graphsurgeon module. Error:\n{e}"52                )53                raise54 55            count_after = len(self.graph.nodes)56            if count_before == count_after:57                # No new folding occurred in this iteration, so we can stop for now.58                break59 60    def save(self, output_path):61        """62        Save the ONNX model to the given location.63        Args:64            output_path: Path pointing to the location where to write65                out the updated ONNX model.66        """67        self.graph.cleanup().toposort()68        model = gs.export_onnx(self.graph)69        onnx.save(model, output_path)70        LOGGER.info(f"Saved ONNX model to {output_path}")71 72    def register_nms(73        self,74        *,75        score_thresh: float = 0.25,76        nms_thresh: float = 0.45,77        detections_per_img: int = 100,78    ):79        """80        Register the ``EfficientNMS_TRT`` plugin node.81        NMS expects these shapes for its input tensors:82            - box_net: [batch_size, number_boxes, 4]83            - class_net: [batch_size, number_boxes, number_labels]84        Args:85            score_thresh (float): The scalar threshold for score (low scoring boxes are removed).86            nms_thresh (float): The scalar threshold for IOU (new boxes that have high IOU87                overlap with previously selected boxes are removed).88            detections_per_img (int): Number of best detections to keep after NMS.89        """90 91        self.infer()92        # Find the concat node at the end of the network93        op_inputs = self.graph.outputs94        op = "EfficientNMS_TRT"95        attrs = {96            "plugin_version": "1",97            "background_class": -1,  # no background class98            "max_output_boxes": detections_per_img,99            "score_threshold": score_thresh,100            "iou_threshold": nms_thresh,101            "score_activation": False,102            "box_coding": 0,103        }104 105        if self.precision == "fp32":106            dtype_output = np.float32107        elif self.precision == "fp16":108            dtype_output = np.float16109        else:110            raise NotImplementedError(f"Currently not supports precision: {self.precision}")111 112        # NMS Outputs113        output_num_detections = gs.Variable(114            name="num_dets",115            dtype=np.int32,116            shape=[self.batch_size, 1],117        )  # A scalar indicating the number of valid detections per batch image.118        output_boxes = gs.Variable(119            name="det_boxes",120            dtype=dtype_output,121            shape=[self.batch_size, detections_per_img, 4],122        )123        output_scores = gs.Variable(124            name="det_scores",125            dtype=dtype_output,126            shape=[self.batch_size, detections_per_img],127        )128        output_labels = gs.Variable(129            name="det_classes",130            dtype=np.int32,131            shape=[self.batch_size, detections_per_img],132        )133 134        op_outputs = [output_num_detections, output_boxes, output_scores, output_labels]135 136        # Create the NMS Plugin node with the selected inputs. The outputs of the node will also137        # become the final outputs of the graph.138        self.graph.layer(op=op, name="batched_nms", inputs=op_inputs, outputs=op_outputs, attrs=attrs)139        LOGGER.info(f"Created NMS plugin '{op}' with attributes: {attrs}")140 141        self.graph.outputs = op_outputs142 143        self.infer()144 145    def save(self, output_path):146        """147        Save the ONNX model to the given location.148        Args:149            output_path: Path pointing to the location where to write150                out the updated ONNX model.151        """152        self.graph.cleanup().toposort()153        model = gs.export_onnx(self.graph)154        onnx.save(model, output_path)155        LOGGER.info(f"Saved ONNX model to {output_path}")156