| 504 | self.__dict__.update(locals()) # assign all variables to self |
| 505 | |
| 506 | def forward(self, im, augment=False, visualize=False): |
| 507 | # YOLOv5 MultiBackend inference |
| 508 | b, ch, h, w = im.shape # batch, channel, height, width |
| 509 | if self.fp16 and im.dtype != torch.float16: |
| 510 | im = im.half() # to FP16 |
| 511 | if self.nhwc: |
| 512 | im = im.permute(0, 2, 3, 1) # torch BCHW to numpy BHWC shape(1,320,192,3) |
| 513 | |
| 514 | if self.pt: # PyTorch |
| 515 | y = self.model(im, augment=augment, visualize=visualize) if augment or visualize else self.model(im) |
| 516 | elif self.jit: # TorchScript |
| 517 | y = self.model(im) |
| 518 | elif self.dnn: # ONNX OpenCV DNN |
| 519 | im = im.cpu().numpy() # torch to numpy |
| 520 | self.net.setInput(im) |
| 521 | y = self.net.forward() |
| 522 | elif self.onnx: # ONNX Runtime |
| 523 | im = im.cpu().numpy() # torch to numpy |
| 524 | y = self.session.run(self.output_names, {self.session.get_inputs()[0].name: im}) |
| 525 | elif self.xml: # OpenVINO |
| 526 | im = im.cpu().numpy() # FP32 |
| 527 | y = list(self.executable_network([im]).values()) |
| 528 | elif self.engine: # TensorRT |
| 529 | if self.dynamic and im.shape != self.bindings['images'].shape: |
| 530 | i = self.model.get_binding_index('images') |
| 531 | self.context.set_binding_shape(i, im.shape) # reshape if dynamic |
| 532 | self.bindings['images'] = self.bindings['images']._replace(shape=im.shape) |
| 533 | for name in self.output_names: |
| 534 | i = self.model.get_binding_index(name) |
| 535 | self.bindings[name].data.resize_(tuple(self.context.get_binding_shape(i))) |
| 536 | s = self.bindings['images'].shape |
| 537 | assert im.shape == s, f"input size {im.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}" |
| 538 | self.binding_addrs['images'] = int(im.data_ptr()) |
| 539 | self.context.execute_v2(list(self.binding_addrs.values())) |
| 540 | y = [self.bindings[x].data for x in sorted(self.output_names)] |
| 541 | elif self.coreml: # CoreML |
| 542 | im = im.cpu().numpy() |
| 543 | im = Image.fromarray((im[0] * 255).astype('uint8')) |
| 544 | # im = im.resize((192, 320), Image.ANTIALIAS) |
| 545 | y = self.model.predict({'image': im}) # coordinates are xywh normalized |
| 546 | if 'confidence' in y: |
| 547 | box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels |
| 548 | conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float) |
| 549 | y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1) |
| 550 | else: |
| 551 | y = list(reversed(y.values())) # reversed for segmentation models (pred, proto) |
| 552 | elif self.paddle: # PaddlePaddle |
| 553 | im = im.cpu().numpy().astype(np.float32) |
| 554 | self.input_handle.copy_from_cpu(im) |
| 555 | self.predictor.run() |
| 556 | y = [self.predictor.get_output_handle(x).copy_to_cpu() for x in self.output_names] |
| 557 | elif self.triton: # NVIDIA Triton Inference Server |
| 558 | y = self.model(im) |
| 559 | else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU) |
| 560 | im = im.cpu().numpy() |
| 561 | if self.saved_model: # SavedModel |
| 562 | y = self.model(im, training=False) if self.keras else self.model(im) |
| 563 | elif self.pb: # GraphDef |