ultralytics 8.0.37 add TFLite metadata in AutoBackend (#953)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com> Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com> Co-authored-by: Aarni Koskela <akx@iki.fi>
This commit is contained in:
parent
20fe708f31
commit
bdc6cd4d8b
18 changed files with 86 additions and 46 deletions
|
|
@ -1,7 +1,9 @@
|
|||
# Ultralytics YOLO 🚀, GPL-3.0 license
|
||||
|
||||
import ast
|
||||
import contextlib
|
||||
import json
|
||||
import platform
|
||||
import zipfile
|
||||
from collections import OrderedDict, namedtuple
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
|
@ -207,6 +209,12 @@ class AutoBackend(nn.Module):
|
|||
interpreter.allocate_tensors() # allocate
|
||||
input_details = interpreter.get_input_details() # inputs
|
||||
output_details = interpreter.get_output_details() # outputs
|
||||
# load metadata
|
||||
with contextlib.suppress(zipfile.BadZipFile):
|
||||
with zipfile.ZipFile(w, "r") as model:
|
||||
meta_file = model.namelist()[0]
|
||||
meta = ast.literal_eval(model.read(meta_file).decode("utf-8"))
|
||||
stride, names = int(meta['stride']), meta['names']
|
||||
elif tfjs: # TF.js
|
||||
raise NotImplementedError('ERROR: YOLOv8 TF.js inference is not supported')
|
||||
elif paddle: # PaddlePaddle
|
||||
|
|
@ -214,7 +222,7 @@ class AutoBackend(nn.Module):
|
|||
check_requirements('paddlepaddle-gpu' if cuda else 'paddlepaddle')
|
||||
import paddle.inference as pdi
|
||||
if not Path(w).is_file(): # if not *.pdmodel
|
||||
w = next(Path(w).rglob('*.pdmodel')) # get *.xml file from *_openvino_model dir
|
||||
w = next(Path(w).rglob('*.pdmodel')) # get *.pdmodel file from *_paddle_model dir
|
||||
weights = Path(w).with_suffix('.pdiparams')
|
||||
config = pdi.Config(str(w), str(weights))
|
||||
if cuda:
|
||||
|
|
@ -328,6 +336,9 @@ class AutoBackend(nn.Module):
|
|||
scale, zero_point = output['quantization']
|
||||
x = (x.astype(np.float32) - zero_point) * scale # re-scale
|
||||
y.append(x)
|
||||
# TF segment fixes: export is reversed vs ONNX export and protos are transposed
|
||||
if len(self.output_details) == 2: # segment
|
||||
y = [y[1], np.transpose(y[0], (0, 3, 1, 2))]
|
||||
y = [x if isinstance(x, np.ndarray) else x.numpy() for x in y]
|
||||
y[0][..., :4] *= [w, h, w, h] # xywh normalized to pixels
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue