ultralytic export

This commit is contained in:
Siwat Sirichai 2025-11-11 01:28:19 +07:00
parent bf7b68edb1
commit fdaeb9981c
14 changed files with 2241 additions and 507 deletions

View file

@ -101,14 +101,15 @@ class YOLOv8Utils:
# Get output tensor (first and only output)
output_name = list(outputs.keys())[0]
output = outputs[output_name] # (1, 84, 8400)
output = outputs[output_name] # (1, 4+num_classes, 8400)
# Transpose to (1, 8400, 84) for easier processing
output = output.transpose(1, 2).squeeze(0) # (8400, 84)
# Transpose to (1, 8400, 4+num_classes) for easier processing
output = output.transpose(1, 2).squeeze(0) # (8400, 4+num_classes)
# Split bbox coordinates and class scores (vectorized)
# Format: [cx, cy, w, h, class_scores...]
bboxes = output[:, :4] # (8400, 4) - (cx, cy, w, h)
class_scores = output[:, 4:] # (8400, 80)
class_scores = output[:, 4:] # (8400, num_classes) - dynamically sized
# Get max class score and corresponding class ID for all anchors (vectorized)
max_scores, class_ids = torch.max(class_scores, dim=1) # (8400,), (8400,)