HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.0.229 add model.embed() method (#7098)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher 2023-12-22 15:32:06 +01:00 • committed by GitHub
parent 38eaf5e29f
commit 5b3e20379f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
11 changed files with 65 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

5
ultralytics/nn/autobackend.py
View file

@ -333,7 +333,7 @@ class AutoBackend(nn.Module):
self.__dict__.update(locals()) # assign all variables to self
def forward(self, im, augment=False, visualize=False):
def forward(self, im, augment=False, visualize=False, embed=None):
"""
Runs inference on the YOLOv8 MultiBackend model.
@ -341,6 +341,7 @@ class AutoBackend(nn.Module):
im (torch.Tensor): The image tensor to perform inference on.
augment (bool): whether to perform data augmentation during inference, defaults to False
visualize (bool): whether to visualize the output predictions, defaults to False
embed (list, optional): A list of feature vectors/embeddings to return.
Returns:
(tuple): Tuple containing the raw output tensor, and processed output for visualization (if visualize=True)
@ -352,7 +353,7 @@ class AutoBackend(nn.Module):
im = im.permute(0, 2, 3, 1) # torch BCHW to numpy BHWC shape(1,320,192,3)
if self.pt or self.nn_module: # PyTorch
y = self.model(im, augment=augment, visualize=visualize) if augment or visualize else self.model(im)
y = self.model(im, augment=augment, visualize=visualize, embed=embed)
elif self.jit: # TorchScript
y = self.model(im)
elif self.dnn: # ONNX OpenCV DNN