|
|
@@ -323,9 +323,12 @@ class DetectMultiBackend(nn.Module): |
|
|
|
super().__init__() |
|
|
|
w = str(weights[0] if isinstance(weights, list) else weights) |
|
|
|
pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = self.model_type(w) # get backend |
|
|
|
stride, names = 32, [f'class{i}' for i in range(1000)] # assign defaults |
|
|
|
w = attempt_download(w) # download if not local |
|
|
|
fp16 &= (pt or jit or onnx or engine) and device.type != 'cpu' # FP16 |
|
|
|
stride, names = 32, [f'class{i}' for i in range(1000)] # assign defaults |
|
|
|
if data: # assign class names (optional) |
|
|
|
with open(data, errors='ignore') as f: |
|
|
|
names = yaml.safe_load(f)['names'] |
|
|
|
|
|
|
|
if pt: # PyTorch |
|
|
|
model = attempt_load(weights if isinstance(weights, list) else w, map_location=device) |
|
|
@@ -365,7 +368,9 @@ class DetectMultiBackend(nn.Module): |
|
|
|
network = ie.read_model(model=w, weights=Path(w).with_suffix('.bin')) |
|
|
|
executable_network = ie.compile_model(model=network, device_name="CPU") |
|
|
|
output_layer = next(iter(executable_network.outputs)) |
|
|
|
self._load_metadata(w.parent / 'meta.yaml') # load metadata |
|
|
|
meta = w.with_suffix('.yaml') |
|
|
|
if meta.exists(): |
|
|
|
stride, names = self._load_metadata(meta) # load metadata |
|
|
|
elif engine: # TensorRT |
|
|
|
LOGGER.info(f'Loading {w} for TensorRT inference...') |
|
|
|
import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download |
|
|
@@ -431,11 +436,7 @@ class DetectMultiBackend(nn.Module): |
|
|
|
output_details = interpreter.get_output_details() # outputs |
|
|
|
elif tfjs: |
|
|
|
raise Exception('ERROR: YOLOv5 TF.js inference is not supported') |
|
|
|
|
|
|
|
self.__dict__.update(locals()) # assign all variables to self |
|
|
|
if not hasattr(self, 'names') and data: # assign class names (optional) |
|
|
|
with open(data, errors='ignore') as f: |
|
|
|
names = yaml.safe_load(f)['names'] |
|
|
|
|
|
|
|
def forward(self, im, augment=False, visualize=False, val=False): |
|
|
|
# YOLOv5 MultiBackend inference |
|
|
@@ -495,13 +496,6 @@ class DetectMultiBackend(nn.Module): |
|
|
|
y = torch.tensor(y, device=self.device) |
|
|
|
return (y, []) if val else y |
|
|
|
|
|
|
|
def _load_metadata(self, f='path/to/meta.yaml'): |
|
|
|
# Load metadata from meta.yaml if it exists |
|
|
|
if Path(f).is_file(): |
|
|
|
with open(f, errors='ignore') as f: |
|
|
|
for k, v in yaml.safe_load(f).items(): |
|
|
|
setattr(self, k, v) # assign stride, names |
|
|
|
|
|
|
|
def warmup(self, imgsz=(1, 3, 640, 640)): |
|
|
|
# Warmup model by running inference once |
|
|
|
warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb |
|
|
@@ -522,6 +516,13 @@ class DetectMultiBackend(nn.Module): |
|
|
|
tflite &= not edgetpu # *.tflite |
|
|
|
return pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs |
|
|
|
|
|
|
|
@staticmethod |
|
|
|
def _load_metadata(f='path/to/meta.yaml'): |
|
|
|
# Load metadata from meta.yaml if it exists |
|
|
|
with open(f, errors='ignore') as f: |
|
|
|
d = yaml.safe_load(f) |
|
|
|
return d['stride'], d['names'] # assign stride, names |
|
|
|
|
|
|
|
|
|
|
|
class AutoShape(nn.Module): |
|
|
|
# YOLOv5 input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS |