python YoloV4到TFLite模型给出了完全错误的预测

nfs0ujit  于 2023-04-04  发布在  Python
关注(0)|答案(2)|浏览(173)

我在自定义数据集上训练了我的yolov 4模型,它给了我大约90-100%的准确率,太棒了!但是我想将权重转换为TFLite,因为我需要在边缘设备上执行一些检测任务。但是在我将其转换为第一个TF权重,然后是TFLite权重之后,如下图所示,它输出了非常随机的预测。第一张图片来自Darknet检测,第二张来自TFLite检测。
使用以下命令进行暗网检测:
!./darknet detector test data/obj.data cfg/yolov4-obj.cfg /content/drive/MyDrive/yolov4/backup/yolov4-obj_last.weights /content/drive/MyDrive/yolov4/test/t2.webp -thresh 0.3

使用TFLite重量:python detect.py

我不知道是什么原因导致的。此外,我在将暗网权重转换为tflite权重时没有遇到任何错误。下面是脚本:

  1. save_model.py-保存自定义yolov 4 tf模型用于tflite转换
    python save_model.py --weights ./data/custom.weights --output ./checkpoints/custom-416 --input_size 416 --model yolov4 --framework tflite
import tensorflow as tf
from absl import app, flags, logging
from absl.flags import FLAGS
from core.yolov4 import YOLO, decode, filter_boxes
import core.utils as utils
from core.config import cfg

flags.DEFINE_string('weights', './data/yolov4.weights', 'path to weights file')
flags.DEFINE_string('output', './checkpoints/yolov4-416', 'path to output')
flags.DEFINE_boolean('tiny', False, 'is yolo-tiny or not')
flags.DEFINE_integer('input_size', 416, 'define input size of export model')
flags.DEFINE_float('score_thres', 0.2, 'define score threshold')
flags.DEFINE_string('framework', 'tf', 'define what framework do you want to convert (tf, trt, tflite)')
flags.DEFINE_string('model', 'yolov4', 'yolov3 or yolov4')

def save_tf():
  STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config()

  input_layer = tf.keras.layers.Input([FLAGS.input_size, FLAGS.input_size, 3])
  feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny)
  bbox_tensors = []
  prob_tensors = []
  if FLAGS.tiny:
    for i, fm in enumerate(feature_maps):
      if i == 0:
        output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE, FLAGS.framework)
      else:
        output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE, FLAGS.framework)
      bbox_tensors.append(output_tensors[0])
      prob_tensors.append(output_tensors[1])
  else:
    for i, fm in enumerate(feature_maps):
      if i == 0:
        output_tensors = decode(fm, FLAGS.input_size // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE, FLAGS.framework)
      elif i == 1:
        output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE, FLAGS.framework)
      else:
        output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE, FLAGS.framework)
      bbox_tensors.append(output_tensors[0])
      prob_tensors.append(output_tensors[1])
  pred_bbox = tf.concat(bbox_tensors, axis=1)
  pred_prob = tf.concat(prob_tensors, axis=1)
  if FLAGS.framework == 'tflite':
    pred = (pred_bbox, pred_prob)
  else:
    boxes, pred_conf = filter_boxes(pred_bbox, pred_prob, score_threshold=FLAGS.score_thres, input_shape=tf.constant([FLAGS.input_size, FLAGS.input_size]))
    pred = tf.concat([boxes, pred_conf], axis=-1)
  model = tf.keras.Model(input_layer, pred)
  utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.tiny)
  model.summary()
  model.save(FLAGS.output)

def main(_argv):
  save_tf()

if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass
  1. python convert2tflite.py-现在,将上面的yolov 4 tf模型转换为tflite权重。
import tensorflow as tf
from absl import logging
import numpy as np

def save_tflite(model_path, tf_weights_path, quantize_mode='float16'):
  converter = tf.lite.TFLiteConverter.from_saved_model(tf_weights_path)
  tflite_model = converter.convert()
  open(model_path, 'wb').write(tflite_model)

  logging.info("model saved to: {}".format(model_path))

def demo(model_path):
  interpreter = tf.lite.Interpreter(model_path=model_path)
  interpreter.allocate_tensors()
  logging.info('tflite model loaded')

  input_details = interpreter.get_input_details()
  print(input_details)
  output_details = interpreter.get_output_details()
  print(output_details)

  input_shape = input_details[0]['shape']

  input_data = np.array(np.random.random_sample(input_shape), dtype=np.float32)

  interpreter.set_tensor(input_details[0]['index'], input_data)
  interpreter.invoke()
  output_data = [interpreter.get_tensor(output_details[i]['index']) for i in range(len(output_details))]

  print(output_data)

def main():
  save_tflite(
    tf_weights_path='./checkpoints/anpr-416',
    model_path='./checkpoints/anpr-416.tflite'
  )
  demo(model_path='./checkpoints/anpr-416.tflite')

if __name__ == '__main__':
    try:
        main()
    except SystemExit:
        pass

如果有人能给予任何见解,我真的会很感激!谢谢!

pengsaosao

pengsaosao1#

好吧,我做错的是,我没有改变类的名称在obj.names文件,而转换到TF模型,这是位于core/config.py,所以在obj.names文件中所需的变化后,在core/config.py,我现在能够作出正确的检测!

hlswsv35

hlswsv352#

我在yolo v5和tflite上也遇到了同样的问题(yolov5使用PyTourch权重,而不是暗网),你能解释一下你是如何解决这个问题的吗?

相关问题