Glenn Jocher
GitHub
3 years ago
1 changed files with 2 additions and 2 deletions
+ 2
- 2
utils/general.py
View File
| detections with shape: nx6 (x1, y1, x2, y2, conf, cls) | detections with shape: nx6 (x1, y1, x2, y2, conf, cls) | ||||
| """ | """ | ||||
| nc = prediction[0].shape[1] - 5 # number of classes | |||||
| nc = prediction.shape[2] - 5 # number of classes | |||||
| xc = prediction[..., 4] > conf_thres # candidates | xc = prediction[..., 4] > conf_thres # candidates | ||||
| # Settings | # Settings | ||||
| merge = False # use merge-NMS | merge = False # use merge-NMS | ||||
| t = time.time() | t = time.time() | ||||
| output = [torch.zeros(0, 6)] * prediction.shape[0] | |||||
| output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0] | |||||
| for xi, x in enumerate(prediction): # image index, image inference | for xi, x in enumerate(prediction): # image index, image inference | ||||
| # Apply constraints | # Apply constraints | ||||
| # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height | # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height |
Loading…