feature: support native YOLO .pt models while ensuring compatibility with Torchvision models (#495)

* feature: support native YOLO .pt models with centralized output normalization

* feature: support native YOLO .pt models while ensuring compatibility with Torchvision models

* Removed dtype extraction and subsequent casting in DetectionModelWrapper

Author: kashtennyson

perceptionmetrics/models/torch_detection.py

@@ -268,27 +268,51 @@ def __init__(
             try:
                 model = torch.jit.load(model, map_location=self.device)
                 model_type = "compiled"
-            except RuntimeError:
+            except Exception:
                 try:
-                    model = torch.load(
-                        model, map_location=self.device, weights_only=False
-                    )
+                    loaded = torch.load(model, map_location=self.device, weights_only=False)
+                    # Handle Ultralytics/YOLO-style dict checkpoints
+                    if isinstance(loaded, dict):
+                        candidate = loaded.get("ema") or loaded.get("model")
+                        if candidate is None or not hasattr(candidate, "forward"):
+                            raise ValueError(
+                                """
+                                The loaded .pt file is a dictionary but doesn't contain a valid model under keys 'model' or 'ema'. Please export to TorchScript for better compatibility.
+                                """
+                            )
+                        model = candidate
+                    else:
+                        model = loaded
                     model_type = "native"
-                except Exception as e:
-                    raise ValueError(
-                        f"Failed to load model. "
-                        f"Ensure it is a valid PyTorch or TorchScript model. Error : {e}"
-                    )
-
-        elif isinstance(model, torch.nn.Module):
-            model_fname = None
-            model_type = "native"
-        else:
-            raise ValueError("Model must be a filename or a torch.nn.Module")
+                # Fallback for missing Ultralytics dependency
+                except (ModuleNotFoundError, AttributeError) as e:
+                    raise ImportError(
+                        f"Failed to load native .pt model. This often happens if the 'ultralytics' "
+                        f"library is missing or incompatible. \nOriginal error: {e}\n"
+                        f"SUGGESTION: 'pip install ultralytics' or export your model to TorchScript."
+                    ) from e
 
         # Init parent class
         super().__init__(model, model_type, model_cfg, ontology_fname, model_fname)
-        self.model = self.model.to(self.device).eval()
+        
+        # Define the Wrapper with DType Alignment
+        class DetectionModelWrapper(torch.nn.Module):
+            def __init__(self, model):
+                super().__init__()
+                self.inner_model = model
+            
+            # Handle input precision, tuple extraction, and output precision
+            def forward(self, x):
+                out = self.inner_model(x)
+                # Only unwrap TUPLES (native YOLO .pt)
+                # Avoid unwrapping LISTS (torchvision Mask R-CNN)
+                if isinstance(out, tuple) and len(out) > 0:
+                    out = out[0]
+                if isinstance(out, torch.Tensor):
+                    return out.float()
+                return out
+
+        self.model = DetectionModelWrapper(self.model).to(self.device).float().eval()
 
         # Load post-processing functions for specific model formats
         self.model_format = self.model_cfg.get("model_format", "torchvision")