text: from __future__ import print_function
  import argparse
  import collections
  import os
  import re
  import logging
  import ctypes
  import numpy as np
  import io
  
  import cv2
  import tensorrt as trt
  import pycuda.driver as cuda
  from yolo_visualization import BBoxVisualization
  
  try:
      ctypes.cdll.LoadLibrary('./yolo/plugins/libyolo_layer.so')
  except OSError as e:
      print('ERROR: failed to load ./yolo/plugins/libyolo_layer.so.  '
                       'Did you forget to do a "make" in the "./yolo/plugins/" '
                       'subdirectory?')
      exit()
  
  class YoloModel:
      def __init__(self):
          self.default_model_dir = './yolo'
          self.default_model = 'yolov4-tiny-416'
          self.default_labels = 'coco.names'
          self.parser = argparse.ArgumentParser()
          self.parser.add_argument('--labels', help='label file path',
                              default=os.path.join(self.default_model_dir, self.default_labels))
          self.parser.add_argument('--top_k', type=int, default=10,
                              help='number of categories with highest score to display')
          self.parser.add_argument('--threshold', type=float, default=0.70,
                              help='classifier score threshold')
          """Parse input arguments."""
          desc = ('Run the TensorRT optimized object detecion model on an input '
                  'video and save BBoxed overlaid output as another video.')
          self.parser.add_argument(
              '-c', '--category_num', type=int, default=80,
              help='number of object categories [80]')
          self.parser.add_argument(
              '-m', '--model', type=str,
              default=self.default_model,
              help=('[yolov3|yolov3-tiny|yolov3-spp|yolov4|yolov4-tiny]-'
                    '[{dimension}], where dimension could be a single '
                    'number (e.g. 288, 416, 608) or WxH (e.g. 416x256)'))
          self.args = self.parser.parse_args([])
  
          print('Loading {} with {} labels.'.format(self.args.model, self.args.labels))
  
          self.cls_dict = self.load_labels(self.args.labels)
          self.vis = BBoxVisualization(self.cls_dict)
          self.h, self.w = get_input_shape(self.args.model)
          self.trt_yolo = TrtYOLO(self.default_model_dir,self.args.model, (self.h, self.w), self.args.category_num)
  
          self.labels = self.load_labels(self.args.labels)
          print('labels {}'.format(self.labels))
      def draw_bboxes(self, img, objects):
          gun_count = 0
          """Draw detected bounding boxes on the original image."""
          for obj in objects:
              #logging.debug('obj is {}'.format(obj))
              cl = int(obj['class_id'])
              if cl == 1:
                  gun_count += 1
              cf = obj['confidence']
              x_min, y_min, x_max, y_max = int(obj['xmin']), int(obj['ymin']), int(obj['xmax']), int(obj['ymax'])
              color = self.vis.colors[cl]
              cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color, 2)
              txt_loc = (max(x_min+2, 0), max(y_min+2, 0))
              txt = '{} {:.2f}'.format(cl, cf)
              img = draw_boxed_text(img, txt, txt_loc, color)
          return (img,gun_count)
      def load_labels(self, path):
          p = re.compile(r'\s*(\d+)(.+)')
          with io.open(path, 'r', encoding='utf8') as f:
              #print(f.readlines())
              lines = (p.match(line).groups() for line in f.readlines())
              return {int(num): text.strip() for num, text in lines}
  def get_input_shape_by_model_name(model_name):
      """Infer input tensor shape of the model based on name.
  
      For example,
          'yolov4-tiny-416': (416, 416)
          'yolov4-custom-416x256': (256, 416)
  
      # Args
          model_name
      """
      yolo_dim = model_name.split('-')[-1]
      if not yolo_dim or not yolo_dim[0].isdigit():
          raise ValueError('bad dim spec in model name (%s)!' % model_name)
      if 'x' in yolo_dim:
          dim_split = yolo_dim.split('x')
          if len(dim_split) != 2:
              raise ValueError('bad yolo_dim (%s)!' % yolo_dim)
          w, h = int(dim_split[0]), int(dim_split[1])
      else:
          h = w = int(yolo_dim)
      if h % 32 != 0 or w % 32 != 0:
          raise ValueError('yolo_dim not multiples of 32 (%s)!' % yolo_dim)
      print('yolo input shape is {}x{}'.format(h,w))
      return h, w
  def get_input_shape(model_name):
      """Infer input tensor shape fixed by 416.
      # Args
          model_name
      """
      return 416, 416
  
  def _preprocess_yolo(img, input_shape, letter_box=False):
      """Preprocess an image before TRT YOLO inferencing.
  
      # Args
          img: int8 numpy array of shape (img_h, img_w, 3)
          input_shape: a tuple of (H, W)
          letter_box: boolean, specifies whether to keep aspect ratio and
                      create a "letterboxed" image for inference
  
      # Returns
          preprocessed img: float32 numpy array of shape (3, H, W)
      """
      if letter_box:
          img_h, img_w, _ = img.shape
          new_h, new_w = input_shape[0], input_shape[1]
          offset_h, offset_w = 0, 0
          if (new_w / img_w) <= (new_h / img_h):
              new_h = int(img_h * new_w / img_w)
              offset_h = (input_shape[0] - new_h) // 2
          else:
              new_w = int(img_w * new_h / img_h)
              offset_w = (input_shape[1] - new_w) // 2
          resized = cv2.resize(img, (new_w, new_h))
          img = np.full((input_shape[0], input_shape[1], 3), 127, dtype=np.uint8)
          img[offset_h:(offset_h + new_h), offset_w:(offset_w + new_w), :] = resized
      else:
          img = cv2.resize(img, (input_shape[1], input_shape[0]))
  
      img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
      img = img.transpose((2, 0, 1)).astype(np.float32)
      img /= 255.0
      return img
  
  
  def _nms_boxes(detections, nms_threshold):
      """Apply the Non-Maximum Suppression (NMS) algorithm on the bounding
      boxes with their confidence scores and return an array with the
      indexes of the bounding boxes we want to keep.
  
      # Args
          detections: Nx7 numpy arrays of
                      [[x, y, w, h, box_confidence, class_id, class_prob],
                       ......]
      """
      x_coord = detections[:, 0]
      y_coord = detections[:, 1]
      width = detections[:, 2]
      height = detections[:, 3]
      box_confidences = detections[:, 4] * detections[:, 6]
  
      areas = width * height
      ordered = box_confidences.argsort()[::-1]
  
      keep = list()
      while ordered.size > 0:
          # Index of the current element:
          i = ordered[0]
          keep.append(i)
          xx1 = np.maximum(x_coord[i], x_coord[ordered[1:]])
          yy1 = np.maximum(y_coord[i], y_coord[ordered[1:]])
          xx2 = np.minimum(x_coord[i] + width[i], x_coord[ordered[1:]] + width[ordered[1:]])
          yy2 = np.minimum(y_coord[i] + height[i], y_coord[ordered[1:]] + height[ordered[1:]])
  
          width1 = np.maximum(0.0, xx2 - xx1 + 1)
          height1 = np.maximum(0.0, yy2 - yy1 + 1)
          intersection = width1 * height1
          union = (areas[i] + areas[ordered[1:]] - intersection)
          iou = intersection / union
          indexes = np.where(iou <= nms_threshold)[0]
          ordered = ordered[indexes + 1]
  
      keep = np.array(keep)
      return keep
  
  
  def _postprocess_yolo(trt_outputs, img_w, img_h, conf_th, nms_threshold,
                        input_shape, letter_box=False):
      """Postprocess TensorRT outputs.
  
      # Args
          trt_outputs: a list of 2 or 3 tensors, where each tensor
                      contains a multiple of 7 float32 numbers in
                      the order of [x, y, w, h, box_confidence, class_id, class_prob]
          conf_th: confidence threshold
          letter_box: boolean, referring to _preprocess_yolo()
  
      # Returns
          boxes, scores, classes (after NMS)
      """
      # filter low-conf detections and concatenate results of all yolo layers
      detections = []
      for o in trt_outputs:
          dets = o.reshape((-1, 7))
          dets = dets[dets[:, 4] * dets[:, 6] >= conf_th]
          detections.append(dets)
      detections = np.concatenate(detections, axis=0)
  
      if len(detections) == 0:
          boxes = np.zeros((0, 4), dtype=np.int)
          scores = np.zeros((0,), dtype=np.float32)
          classes = np.zeros((0,), dtype=np.float32)
      else:
          box_scores = detections[:, 4] * detections[:, 6]
  
          # scale x, y, w, h from [0, 1] to pixel values
          old_h, old_w = img_h, img_w
          offset_h, offset_w = 0, 0
          if letter_box:
              if (img_w / input_shape[1]) >= (img_h / input_shape[0]):
                  old_h = int(input_shape[0] * img_w / input_shape[1])
                  offset_h = (old_h - img_h) // 2
              else:
                  old_w = int(input_shape[1] * img_h / input_shape[0])
                  offset_w = (old_w - img_w) // 2
          detections[:, 0:4] *= np.array(
              [old_w, old_h, old_w, old_h], dtype=np.float32)
  
          # NMS
          nms_detections = np.zeros((0, 7), dtype=detections.dtype)
          for class_id in set(detections[:, 5]):
              idxs = np.where(detections[:, 5] == class_id)
              cls_detections = detections[idxs]
              keep = _nms_boxes(cls_detections, nms_threshold)
              nms_detections = np.concatenate(
                  [nms_detections, cls_detections[keep]], axis=0)
  
          xx = nms_detections[:, 0].reshape(-1, 1)
          yy = nms_detections[:, 1].reshape(-1, 1)
          if letter_box:
              xx = xx - offset_w
              yy = yy - offset_h
          ww = nms_detections[:, 2].reshape(-1, 1)
          hh = nms_detections[:, 3].reshape(-1, 1)
          boxes = np.concatenate([xx, yy, xx+ww, yy+hh], axis=1) + 0.5
          boxes = boxes.astype(np.int)
          scores = nms_detections[:, 4] * nms_detections[:, 6]
          classes = nms_detections[:, 5]
      return boxes, scores, classes
  
  
  class HostDeviceMem(object):
      """Simple helper data class that's a little nicer to use than a 2-tuple."""
      def __init__(self, host_mem, device_mem):
          self.host = host_mem
          self.device = device_mem
  
      def __str__(self):
          return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
  
      def __repr__(self):
          return self.__str__()
  
  
  def allocate_buffers(engine):
      """Allocates all host/device in/out buffers required for an engine."""
      inputs = []
      outputs = []
      bindings = []
      output_idx = 0
      stream = cuda.Stream()
      assert 3 <= len(engine) <= 4  # expect 1 input, plus 2 or 3 outpus
      for binding in engine:
          binding_dims = engine.get_binding_shape(binding)
          if len(binding_dims) == 4:
              # explicit batch case (TensorRT 7+)
              size = trt.volume(binding_dims)
          elif len(binding_dims) == 3:
              # implicit batch case (TensorRT 6 or older)
              size = trt.volume(binding_dims) * engine.max_batch_size
          else:
              raise ValueError('bad dims of binding %s: %s' % (binding, str(binding_dims)))
          dtype = trt.nptype(engine.get_binding_dtype(binding))
          # Allocate host and device buffers
          host_mem = cuda.pagelocked_empty(size, dtype)
          device_mem = cuda.mem_alloc(host_mem.nbytes)
          # Append the device buffer to device bindings.
          bindings.append(int(device_mem))
          # Append to the appropriate list.
          if engine.binding_is_input(binding):
              inputs.append(HostDeviceMem(host_mem, device_mem))
          else:
              # each grid has 3 anchors, each anchor generates a detection
              # output of 7 float32 values
              assert size % 7 == 0
              outputs.append(HostDeviceMem(host_mem, device_mem))
              output_idx += 1
      return inputs, outputs, bindings, stream
  
  
  def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
      """do_inference (for TensorRT 6.x or lower)
  
      This function is generalized for multiple inputs/outputs.
      Inputs and outputs are expected to be lists of HostDeviceMem objects.
      """
      # Transfer input data to the GPU.
      [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
      # Run inference.
      context.execute_async(batch_size=batch_size,
                            bindings=bindings,
                            stream_handle=stream.handle)
      # Transfer predictions back from the GPU.
      [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
      # Synchronize the stream
      stream.synchronize()
      # Return only the host outputs.
      return [out.host for out in outputs]
  
  
  def do_inference_v2(context, bindings, inputs, outputs, stream):
      """do_inference_v2 (for TensorRT 7.0+)
  
      This function is generalized for multiple inputs/outputs for full
      dimension networks.
      Inputs and outputs are expected to be lists of HostDeviceMem objects.
      """
      # Transfer input data to the GPU.
      [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
      # Run inference.
      context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
      # Transfer predictions back from the GPU.
      [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
      # Synchronize the stream
      stream.synchronize()
      # Return only the host outputs.
      return [out.host for out in outputs]
  
  
  class TrtYOLO(object):
      """TrtYOLO class encapsulates things needed to run TRT YOLO."""
  
      def _load_engine(self):
          TRTbin = '{}/{}.trt'.format(self.model_path,self.model)
          with open(TRTbin, 'rb') as f, trt.Runtime(self.trt_logger) as runtime:
              return runtime.deserialize_cuda_engine(f.read())
  
      def __init__(self,model_path ,model, input_shape, category_num=2,
                   cuda_ctx=None):
          """Initialize TensorRT plugins, engine and conetxt."""
          self.model_path = model_path
          self.model = model
          self.input_shape = input_shape
          self.category_num = category_num
          self.letter_box = False
          self.cuda_ctx = cuda_ctx
          import pycuda.autoinit  # This is needed for initializing CUDA driver
          if self.cuda_ctx:
              self.cuda_ctx.push()
  
          self.inference_fn = do_inference if trt.__version__[0] < '7' \
                                           else do_inference_v2
          self.trt_logger = trt.Logger(trt.Logger.INFO)
          self.engine = self._load_engine()
  
          try:
              self.context = self.engine.create_execution_context()
              self.inputs, self.outputs, self.bindings, self.stream = \
                  allocate_buffers(self.engine)
          except Exception as e:
              raise RuntimeError('fail to allocate CUDA resources')
          finally:
              if self.cuda_ctx:
                  self.cuda_ctx.pop()
  
      def __del__(self):
          """Free CUDA memories."""
          del self.outputs
          del self.inputs
          del self.stream
  
      def detect(self, img, conf_th=0.3, letter_box=None):
          """Detect objects in the input image."""
          letter_box = self.letter_box if letter_box is None else letter_box
          img_resized = _preprocess_yolo(img, self.input_shape, letter_box)
  
          # Set host input to the image. The do_inference() function
          # will copy the input to the GPU before executing.
          self.inputs[0].host = np.ascontiguousarray(img_resized)
          if self.cuda_ctx:
              self.cuda_ctx.push()
          trt_outputs = self.inference_fn(
              context=self.context,
              bindings=self.bindings,
              inputs=self.inputs,
              outputs=self.outputs,
              stream=self.stream)
          if self.cuda_ctx:
              self.cuda_ctx.pop()
  
          boxes, scores, classes = _postprocess_yolo(
              trt_outputs, img.shape[1], img.shape[0], conf_th,
              nms_threshold=0.25, input_shape=self.input_shape,
              letter_box=letter_box)
  
          # clip x1, y1, x2, y2 within original image
          boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, img.shape[1]-1)
          boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, img.shape[0]-1)
          return boxes, scores, classes
  