使用python推理tflight(包括int8)模型

#coding:utf-8
import cv2
import numpy as np

import tensorflow as tf
Interpreter = tf.lite.Interpreter

def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
  # Resize and pad image while meeting stride-multiple constraints
  shape = im.shape[:2]  # current shape [height, width]
  if isinstance(new_shape, int):
    new_shape = (new_shape, new_shape)

  # Scale ratio (new / old)
  r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
  if not scaleup:  # only scale down, do not scale up (for better val mAP)
    r = min(r, 1.0)

  # Compute padding
  ratio = r, r  # width, height ratios
  new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
  dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding

  dw /= 2  # divide padding into 2 sides
  dh /= 2

  if shape[::-1] != new_unpad:  # resize
    im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
  top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
  left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
  im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
  return im, ratio, (dw, dh)

def nms_pi(dets, thresh):
  """
  refer to:
  https://github.com/facebookresearch/Detectron/blob/main/detectron/utils/cython_nms.pyx
  Apply classic DPM-style greedy NMS.
  """
  if dets.shape[0] == 0:
    return dets[[], :]
  scores = dets[:, 0]
  x1 = dets[:, 1]
  y1 = dets[:, 2]
  x2 = dets[:, 3]
  y2 = dets[:, 4]

  areas = (x2 - x1 + 1) * (y2 - y1 + 1)
  order = scores.argsort()[::-1]

  ndets = dets.shape[0]
  suppressed = np.zeros((ndets), dtype=np.int32)

  for _i in range(ndets):
    i = order[_i]
    if suppressed[i] == 1:
      continue
    ix1 = x1[i]
    iy1 = y1[i]
    ix2 = x2[i]
    iy2 = y2[i]
    iarea = areas[i]
    for _j in range(_i + 1, ndets):
      j = order[_j]
      if suppressed[j] == 1:
        continue
      xx1 = max(ix1, x1[j])
      yy1 = max(iy1, y1[j])
      xx2 = min(ix2, x2[j])
      yy2 = min(iy2, y2[j])
      w = max(0.0, xx2 - xx1 + 1)
      h = max(0.0, yy2 - yy1 + 1)
      inter = w * h
      ovr = inter / (iarea + areas[j] - inter)

      if ovr >= thresh:
        suppressed[j] = 1
  keep = np.where(suppressed == 0)[0]
  dets = dets[keep, :]
  ### score,xmin,ymin,xmax,ymax,?
  return dets

def post_detections(outs_probablity, image_shape, dw, dh, ratio):
  '''

  :param outs_probablity:
  :param image_shape:
  :param input_width:
  :param input_height:
  :param is_letterbox:
  :param dw:
  :param dh:
  :param ratio:
  :return: res: bbox_score,xmin, ymin, xmax, ymax,class_socre
  '''
  # print('outs_probablity \n',outs_probablity[:3]) ## yolo  xc,yc,w,h,bbx_score,class_score
  height, width, _ = image_shape

  detections = np.zeros(outs_probablity.shape)
  # detections[:, 0] = outs_probablity[:,4]
  detections[:, 0] = outs_probablity[:, 4]

  detections[:, 1] = outs_probablity[:, 0] - outs_probablity[:, 2] / 2.
  detections[:, 2] = outs_probablity[:, 1] - outs_probablity[:, 3] / 2.
  detections[:, 3] = outs_probablity[:, 0] + outs_probablity[:, 2] / 2.
  detections[:, 4] = outs_probablity[:, 1] + outs_probablity[:, 3] / 2.
  detections[:, 5] = outs_probablity[:, 5]

  detections = detections[detections[:, 0] > 0.25]
  detections = detections[detections[:, 5] > 0.25]

  detections[:, 0] = detections[:, 0] * detections[:, 5]
  nms_thres = 0.25
  # nms_thres = 0.5
  # nms_thres = 0.45
  # nms_thres = 0.15
  res = nms_pi(detections, nms_thres)
  print(res)

  res[:, 1] = (res[:, 1] - dw) / ratio[0]
  res[:, 2] = (res[:, 2] - dh) / ratio[0]
  res[:, 3] = (res[:, 3] - dw) / ratio[0]
  res[:, 4] = (res[:, 4] - dh) / ratio[0]

  return res  ## score,xmin,ymin,xmax,ymax(与输入图像的尺寸一致)

class HumanDetect(object):

  def __init__(self, model_name):
    # Initialize model
    interpreter = Interpreter(model_path=model_name, num_threads=4)
    interpreter.allocate_tensors()

    self.input_details = interpreter.get_input_details()  # inputs
    self.output_details = interpreter.get_output_details()  # outputs

    self._input_index = self.input_details[0]['index']
    self._output_index = self.output_details[0]['index']


    self._input_height = self.input_details[0]['shape'][1]
    self._input_width = self.input_details[0]['shape'][2]

    self._interpreter = interpreter
    self._crop_region = None

  def detect(self,img_path,dst_path):
    image = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), -1)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image_height, image_width, _ = image.shape ## [h,w,c]
    input_image,ratio, (dw, dh) = letterbox(image,(self._input_height,self._input_width))
    input_image = input_image.astype(np.float32)  # uint8 to fp16/32
    input_image /= 255.  # 0 - 255 to 0.0 - 1.0
    input_image = np.expand_dims(input_image, axis=0) ## [1,h,w,c]

    input = self.input_details[0]
    int8 = input["dtype"] == np.uint8  # is TFLite quantized uint8 model
    print('int8',int8)
    if int8:
      scale, zero_point = input["quantization"]
      input_image = (input_image / scale + zero_point).astype(np.uint8)  # de-scale

    self._interpreter.set_tensor(self._input_index,input_image)
    self._interpreter.invoke()

    y = []
    for output in self.output_details:
      x = self._interpreter.get_tensor(output["index"])
      print('x',x.shape)
      if False:
        txt_path = 'data/images/detect_output_0.txt'
        txt_path = 'data/images/detect_output_1.txt'
        with open(txt_path) as f:
          data = f.readlines()
        res = []
        for line in data:
          line = line.strip().split()
          line = list(map(int, line))
          res.append(line)
        res = np.array(res)
        x = np.reshape(res,(1500,6))
        x = np.expand_dims(x, axis=0)
        print('x.shape',x.shape)

      if int8:
        scale, zero_point = output["quantization"]
        x = (x.astype(np.float32) - zero_point) * scale  # re-scale
      y.append(x)
    y = [x if isinstance(x, np.ndarray) else x.numpy() for x in y]

    y[0][..., :4] *= [self._input_width, self._input_height, self._input_width, self._input_height]
    res = post_detections(y[0][0], image.shape, dw=dw, dh=dh, ratio = ratio)
    res = np.array(res)
    image = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), -1)
    res = res[:,1:]
    for box in res:
      org_xmin = int(box[0])
      org_ymin = int(box[1])
      org_xmax = int(box[2])
      org_ymax = int(box[3])
      cv2.rectangle(image, (org_xmin, org_ymin), (org_xmax, org_ymax), (0, 255, 0), thickness=2)
    cv2.imencode('.jpg', image)[1].tofile(dst_path)
    cv2.imshow('f',image)
    cv2.waitKey(0)