# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license from __future__ import annotations import platform from pathlib import Path import numpy as np import torch from ultralytics.utils import LOGGER from .base import BaseBackend class TensorFlowBackend(BaseBackend): """Google TensorFlow inference backend supporting multiple serialization formats. Loads and runs inference with Google TensorFlow models in SavedModel, GraphDef (.pb), and Edge TPU formats. Handles quantized model dequantization and task-specific output formatting. """ def __init__(self, weight: str | Path, device: torch.device | str, fp16: bool = False, format: str = "saved_model"): """Initialize the Google TensorFlow backend. Args: weight (str | Path): Path to the SavedModel directory, .pb file, or Edge TPU .tflite file. device (torch.device | str): Device to run inference on, e.g. 'tpu:0' to select an Edge TPU. fp16 (bool): Whether to use FP16 half-precision inference. format (str): Model format, one of "saved_model", "pb", or "edgetpu". """ assert format in {"saved_model", "pb", "edgetpu"}, f"Unsupported TensorFlow format: {format}." self.format = format super().__init__(weight, device, fp16) def load_model(self, weight: str | Path) -> None: """Load a Google TensorFlow model in SavedModel, GraphDef, or Edge TPU format. Args: weight (str | Path): Path to the model file or directory. """ if self.format in {"saved_model", "pb"}: import tensorflow as tf if self.format == "saved_model": LOGGER.info(f"Loading {weight} for TensorFlow SavedModel inference...") self.model = tf.saved_model.load(weight) self.apply_metadata(self.read_metadata(weight)) elif self.format == "pb": LOGGER.info(f"Loading {weight} for TensorFlow GraphDef inference...") from ultralytics.utils.export.tensorflow import gd_outputs def wrap_frozen_graph(gd, inputs, outputs): """Wrap a TensorFlow frozen graph for inference by pruning to specified input/output nodes.""" x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), []) ge = x.graph.as_graph_element return x.prune(tf.nest.map_structure(ge, inputs), tf.nest.map_structure(ge, outputs)) gd = tf.Graph().as_graph_def() with open(weight, "rb") as f: gd.ParseFromString(f.read()) self.frozen_func = wrap_frozen_graph(gd, inputs="x:0", outputs=gd_outputs(gd)) self.apply_metadata(self.read_metadata(weight)) else: # edgetpu try: from tflite_runtime.interpreter import Interpreter, load_delegate self.tf = None except ImportError: import tensorflow as tf self.tf = tf Interpreter, load_delegate = tf.lite.Interpreter, tf.lite.experimental.load_delegate device = self.device[3:] if str(self.device).startswith("tpu") else ":0" LOGGER.info(f"Loading {weight} on device {device[1:]} for TensorFlow Lite Edge TPU inference...") delegate = {"Linux": "libedgetpu.so.1", "Darwin": "libedgetpu.1.dylib", "Windows": "edgetpu.dll"}[ platform.system() ] self.interpreter = Interpreter( model_path=str(weight), experimental_delegates=[load_delegate(delegate, options={"device": device})], ) self.device = torch.device("cpu") # Edge TPU runs on CPU from PyTorch's perspective self.interpreter.allocate_tensors() self.input_details = self.interpreter.get_input_details() self.output_details = self.interpreter.get_output_details() self.apply_metadata(self.read_metadata(weight)) def forward(self, im: torch.Tensor) -> list[np.ndarray]: """Run Google TensorFlow inference with format-specific execution and output post-processing. Args: im (torch.Tensor): Input image tensor in BHWC format (converted from BCHW by AutoBackend). Returns: (list[np.ndarray]): Model predictions as a list of numpy arrays. """ im = im.cpu().numpy() if self.format == "saved_model": y = self.model.serving_default(im) if not isinstance(y, list): y = [y] elif self.format == "pb": import tensorflow as tf y = self.frozen_func(x=tf.constant(im)) else: h, w = im.shape[1:3] details = self.input_details[0] is_int = details["dtype"] in {np.int8, np.int16} if is_int: scale, zero_point = details["quantization"] im = (im / scale + zero_point).astype(details["dtype"]) self.interpreter.set_tensor(details["index"], im) self.interpreter.invoke() y = [] for output in self.output_details: x = self.interpreter.get_tensor(output["index"]) if self.task == "semantic" and x.ndim == 3: # Baked argmax class map [B, H, W] of integer class IDs, not boxes or quantized logits: # skip dequantization and xywh denormalization, which would corrupt and overflow the indices. y.append(x) continue if is_int: scale, zero_point = output["quantization"] x = (x.astype(np.float32) - zero_point) * scale if x.ndim == 3: # Denormalize xywh by image size if x.shape[-1] == 6 or self.end2end: x[:, :, [0, 2]] *= w x[:, :, [1, 3]] *= h if self.task == "pose": x[:, :, 6::3] *= w x[:, :, 7::3] *= h else: x[:, [0, 2]] *= w x[:, [1, 3]] *= h if self.task == "pose": # keypoints follow the box (4) and class-score (nc) channels kpt_start = 4 + len(self.names) x[:, kpt_start::3] *= w x[:, kpt_start + 1 :: 3] *= h y.append(x) if self.task == "segment": # segment with (det, proto) output order reversed if len(y[1].shape) != 4: y = list(reversed(y)) # should be y = (1, 116, 8400), (1, 160, 160, 32) if y[1].shape[-1] == 6: # end-to-end model y = [y[1]] else: y[1] = np.transpose(y[1], (0, 3, 1, 2)) # should be y = (1, 116, 8400), (1, 32, 160, 160) elif self.task in {"semantic", "depth"} and len(y) == 1 and y[0].ndim == 4: y[0] = np.transpose(y[0], (0, 3, 1, 2)) # NHWC → NCHW for semantic segmentation and depth anything logits return [x if isinstance(x, np.ndarray) else x.numpy() for x in y]