Fix YOLOv8 .pt inference crashes and improve output handling

[YUVAN0907]

Fix YOLOv8 .pt inference crashes and improve output handling

Fixes #525

---

What this PR does

This PR fixes inference crashes when using native YOLOv8 .pt models in the UI and improves output handling for better compatibility.

---

Problem

While testing YOLO models in the UI, inference failed even though the model loaded successfully.

Two common errors were observed:

'Tensor' object has no attribute 'boxes'

too many values to unpack (expected 2)

---

Before Fix (Error Screenshots)

• Error: 'Tensor' object has no attribute 'boxes'
• Error: too many values to unpack

---

Root Cause

postprocess_detection() assumed the model always returns a torch.Tensor, but YOLOv8 (Ultralytics) models can return:

• Results object (with .boxes)
• list or tuple
• raw torch.Tensor

Only the tensor case was handled previously, which caused crashes.

Additionally:

• Bounding boxes were returned in resized (model input) space instead of original image space
• Some .pt models failed to load properly
• Class name mapping was not robust

---

Changes

1. YOLO output handling (yolo.py)

• Added support for:
  • Results object (output.boxes)
  • list / tuple outputs
• Preserved existing tensor logic
• Added max_detections support

---

2. Model loading improvements (torch_detection.py)

• Added safe loading for YOLOv8 .pt models
• Supports dict checkpoints (model, state_dict)

---

3. Improved prediction output

• Bounding boxes scaled back to original image size
• Class indices mapped to class names

---

4. Robust ontology handling

Supports:

• dict with "classes"
• list format
• generic mappings

---

After Fix (Working Output)

• Model loads and runs inference successfully
• Bounding boxes displayed correctly

---

🧪 Testing

Tested via Streamlit UI:

Model type	Before	After
YOLOv8 .pt (Results output)	❌ Crash	✅ Works
YOLOv8 .pt (list/tuple output)	❌ Crash	✅ Works
TorchScript models	✅ Works	✅ Still works

---

Files changed

• perceptionmetrics/models/utils/yolo.py
• perceptionmetrics/models/torch_detection.py

---

Notes

• No breaking changes
• Existing functionality preserved
• Changes are backward compatible

---

Let me know if you'd like me to add tests for these cases or refine the implementation further.

@dpascualhe

[YUVAN0907]

Hi @dpascualhe , I've updated this PR and ensured there is no merge conflicts. This PR fixes YOLOv8 .pt inference crashes and improves output handling.
I've also tested existing models, and they continue to work as expected, so no existing functionality is affected

I'd really appreciate your feedback when you have time. Happy to refine it further.
Thanks!

[kashtennyson]

Hi @YUVAN0907,

This issue was already pointed out by me (#449). You can take a brief look at it and its related PR (#495). Thanks!

[YUVAN0907]

Hi @YUVAN0907,

This issue was already pointed out by me (#449). You can take a brief look at it and its related PR (#495). Thanks!

Hi @kashtennyson,

Thank you for pointing this out and for your work on #495 — I’ve reviewed your PR.

From what I understand, your changes mainly handle model loading and normalization (e.g., checkpoint dictionaries, tuple outputs, dtype alignment, etc.), ensuring the model behaves consistently at inference time.

In contrast, my PR focuses on a different stage of the pipeline — specifically postprocessing and UI inference handling. In particular, it:

• Adds support for Ultralytics "Results" objects ("output.boxes")
• Handles list/tuple outputs during postprocessing
• Fixes bounding box scaling back to the original image space
• Improves class mapping and ontology robustness

So while both PRs are related to YOLOv8 compatibility, they address different layers of the pipeline (model loading vs. postprocessing).

I believe the two approaches are complementary. I’d be happy to align my implementation with your wrapper logic if needed, or refactor parts to better integrate with your changes.

If you’re open to it, maybe we can collaborate to consolidate both fixes and present a unified solution for review by @dpascualhe.

Thanks again, and looking forward to your thoughts!

Best regards,
Yuvan

[YUVAN0907]

Hi @YUVAN0907,

This issue was already pointed out by me (#449). You can take a brief look at it and its related PR (#495). Thanks!

Just to add — my changes span both "torch_detection.py" and "yolo.py", whereas your PR primarily focuses on "torch_detection.py".

Because of this, I think our approaches cover different parts of the pipeline, and combining them could lead to a more complete and robust solution for YOLOv8 support.

Happy to collaborate and align the implementations if you're open to it.

[kashtennyson]

Yes @YUVAN0907, we can definitely collaborate! Currently, I am waiting for a feedback on my pull request from the maintainers. Once I receive that I will get back to you and we can collaborate on the final solution since I am already planning on to add a global support for models in .pt format across all perception tasks, and I aim to add this feature cleanly while respecting the architecture of the codebase.

[YUVAN0907]

That sounds great, I appreciate the direction you're taking.

@dpascualhe, since both PRs address different parts of YOLOv8 ".pt" support (model loading and postprocessing), combining them could provide a more complete and robust solution.

Happy to collaborate and align our implementations based on your feedback so we can integrate this cleanly into the codebase.

Best regards,
Yuvan😄

[dpascualhe]

Hi @YUVAN0907 ! We are about to merge @kashtennyson changes in this regard, plus your PR does too many things at once. I'd recommend open up a new issue if you find any relevant changes to do beyond support for native YOLO models and max_detections (#497 ) handling which are already being addressed. We will close this PR for now, but thanks for your interest!

[YUVAN0907]

Hi @dpascualhe,

Thank you for the feedback 😄

I understand that the PR covers multiple concerns, and I appreciate the clarification regarding the scope already being handled in #497 and related changes.

I’ll break down the remaining relevant improvements (such as postprocessing robustness, bounding box scaling, and ontology handling) into smaller, focused issues/PRs where appropriate.

Thanks again for the guidance, and I’ll align future contributions more closely with the project’s structure.

Best regards,
Yuvan