Chiebot-Mirror
/
YOLOv8-TensorRT
mirror of https://github.com/triple-Mu/YOLOv8-TensorRT.git
9
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix boxes split bug

Browse Source 
Fix boxes split bug
pull/32/head
tripleMu 2 years ago committed by GitHub
parent 3c0339127a 5692d10a49
commit 133d1625d9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 2
      models/common.py

2
models/common.py
Unescape View File

@ -153,7 +153,7 @@ class PostSeg(nn.Module):
self.shape = shape self.shape = shape
x = [i.view(b, self.no, -1) for i in res] x = [i.view(b, self.no, -1) for i in res]
y = torch.cat(x, 2) y = torch.cat(x, 2)
boxes, scores = y[:, :, ...], y[:, b_reg_num:, ...].sigmoid() boxes, scores = y[:, :b_reg_num, ...], y[:, b_reg_num:, ...].sigmoid()
boxes = boxes.view(b, 4, self.reg_max, -1).permute(0, 1, 3, 2) boxes = boxes.view(b, 4, self.reg_max, -1).permute(0, 1, 3, 2)
boxes = boxes.softmax(-1) @ torch.arange(self.reg_max).to(boxes) boxes = boxes.softmax(-1) @ torch.arange(self.reg_max).to(boxes)
boxes0, boxes1 = -boxes[:, :2, ...], boxes[:, 2:, ...] boxes0, boxes1 = -boxes[:, :2, ...], boxes[:, 2:, ...]

Write Preview
Loading…
Cancel
Save