xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 05:24:22 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

prediction augment (TTA) fix (#288)

Browse Source
This commit is contained in:
leonnil 2024-06-23 08:21:40 +00:00
parent 2c36ab0f10
commit aad320dd80
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
ultralytics/nn/tasks.py
View File

@ -324,7 +324,11 @@ class DetectionModel(BaseModel):
y = [] # outputs y = [] # outputs
for si, fi in zip(s, f): for si, fi in zip(s, f):
xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max())) xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
yi = super().predict(xi)[0] # forward yi = super().predict(xi) # forward
if isinstance(yi, dict):
yi = yi["one2one"] # yolov10 outputs
if isinstance(yi, (list, tuple)):
yi = yi[0]
yi = self._descale_pred(yi, fi, si, img_size) yi = self._descale_pred(yi, fi, si, img_size)
y.append(yi) y.append(yi)
y = self._clip_augmented(y) # clip augmented tails y = self._clip_augmented(y) # clip augmented tails