fix/update: corrected torch.py and added testing suite for utility functions

perceptionmetrics/utils/torch.py

@@ -16,6 +16,7 @@ def data_to_device(
     :type device: torch.device
     :return: Data moved to device
     :rtype: Union[tuple, list]
+    :raises TypeError: If data is not a tensor, list, or tuple
     """
     if isinstance(data, (tuple, list)):
         return type(data)(
@@ -25,7 +26,9 @@ def data_to_device(
     elif torch.is_tensor(data):
         return data.to(device)
     else:
-        return data
+        raise TypeError(
+            f"data_to_device expected torch.Tensor, list, or tuple, but got {type(data)}"
+        )
 
 
 def get_data_shape(data: Union[tuple, list]) -> Union[tuple, list]:
@@ -35,6 +38,7 @@ def get_data_shape(data: Union[tuple, list]) -> Union[tuple, list]:
     :type data: Union[tuple, list]
     :return: Data shape
     :rtype: Union[tuple, list]
+    :raises TypeError: If data is not a tensor, list, or tuple
     """
     if isinstance(data, (tuple, list)):
         return type(data)(
@@ -43,7 +47,9 @@ def get_data_shape(data: Union[tuple, list]) -> Union[tuple, list]:
     elif torch.is_tensor(data):
         return tuple(data.shape)
     else:
-        return tuple(data.shape)
+        raise TypeError(
+            f"get_data_shape expected torch.Tensor, list, or tuple, but got {type(data)}"
+        )
 
 
 def unsqueeze_data(data: Union[tuple, list], dim: int = 0) -> Union[tuple, list]:
@@ -55,6 +61,7 @@ def unsqueeze_data(data: Union[tuple, list], dim: int = 0) -> Union[tuple, list]
     :type dim: int, optional
     :return: Unsqueezed data
     :rtype: Union[tuple, list]
+    :raises TypeError: If data is not a tensor, list, or tuple
     """
     if isinstance(data, (tuple, list)):
         return type(data)(
@@ -64,7 +71,9 @@ def unsqueeze_data(data: Union[tuple, list], dim: int = 0) -> Union[tuple, list]
     elif torch.is_tensor(data):
         return data.unsqueeze(dim)
     else:
-        return data
+        raise TypeError(
+            f"unsqueeze_data expected torch.Tensor, list, or tuple, but got {type(data)}"
+        )
 
 
 def get_device_info():

tests/utils/test_image.py

@@ -0,0 +1,114 @@
+import numpy as np
+from PIL import Image
+from unittest.mock import MagicMock, patch
+import pytest
+import supervision as sv
+from perceptionmetrics.utils.image import draw_detections
+
+
+def test_draw_detections_calls_supervision():
+    # Setup
+    img = Image.new("RGB", (100, 100))
+    boxes = np.array([[0, 0, 10, 10]])
+    class_ids = np.array([0])
+    class_names = ["cat"]
+    scores = np.array([0.9])
+
+    mock_box_annotator = MagicMock()
+    mock_box_annotator.annotate.return_value = np.zeros((100, 100, 3), dtype=np.uint8)
+
+    with patch("perceptionmetrics.utils.image.sv.Detections") as mock_detections, patch(
+        "perceptionmetrics.utils.image.sv.BoxAnnotator", return_value=mock_box_annotator
+    ), patch("perceptionmetrics.utils.image.sv.Color"), patch(
+        "perceptionmetrics.utils.image.sv.ColorPalette"
+    ):
+
+        draw_detections(img, boxes, class_ids, class_names, scores)
+
+        assert mock_detections.called
+        args, kwargs = mock_detections.call_args
+        assert np.array_equal(kwargs["xyxy"], boxes)
+        assert np.array_equal(kwargs["class_id"], class_ids)
+        assert np.array_equal(kwargs["confidence"], scores)
+
+        assert mock_box_annotator.annotate.called
+
+
+def test_draw_detections_label_construction():
+    # Setup
+    img = Image.new("RGB", (100, 100))
+    boxes = np.array([[0, 0, 10, 10], [10, 10, 20, 20]])
+    class_ids = np.array([0, 1])
+    class_names = ["cat", "dog"]
+    scores = np.array([0.9, 0.8])
+
+    # First BoxAnnotator (older style) will fail on annotate
+    mock_box_annotator_old = MagicMock()
+    mock_box_annotator_old.annotate.side_effect = TypeError(
+        "Simulation of mismatching arguments"
+    )
+
+    # Second BoxAnnotator (modern style) will succeed
+    mock_box_annotator_new = MagicMock()
+    mock_box_annotator_new.annotate.return_value = np.zeros(
+        (100, 100, 3), dtype=np.uint8
+    )
+
+    mock_label_annotator = MagicMock()
+    mock_label_annotator.annotate.return_value = np.zeros((100, 100, 3), dtype=np.uint8)
+
+    with patch(
+        "perceptionmetrics.utils.image.sv.BoxAnnotator",
+        side_effect=[mock_box_annotator_old, mock_box_annotator_new],
+    ), patch(
+        "perceptionmetrics.utils.image.sv.LabelAnnotator",
+        return_value=mock_label_annotator,
+    ), patch(
+        "perceptionmetrics.utils.image.sv.Color"
+    ), patch(
+        "perceptionmetrics.utils.image.sv.ColorPalette"
+    ):
+
+        draw_detections(img, boxes, class_ids, class_names, scores)
+
+        # Check labels passed to LabelAnnotator
+        assert mock_label_annotator.annotate.called
+        args, kwargs = mock_label_annotator.annotate.call_args
+        labels = kwargs.get("labels")
+        assert labels == ["cat: 0.90", "dog: 0.80"]
+
+
+def test_draw_detections_missing_names():
+    img = Image.new("RGB", (100, 100))
+    boxes = np.array([[0, 0, 10, 10]])
+    class_ids = np.array([5])
+    class_names = []  # No name for ID 5
+
+    mock_box_annotator_old = MagicMock()
+    mock_box_annotator_old.annotate.side_effect = TypeError()
+    mock_box_annotator_new = MagicMock()
+    mock_box_annotator_new.annotate.return_value = np.zeros(
+        (100, 100, 3), dtype=np.uint8
+    )
+
+    mock_label_annotator = MagicMock()
+    mock_label_annotator.annotate.return_value = np.zeros((100, 100, 3), dtype=np.uint8)
+
+    with patch(
+        "perceptionmetrics.utils.image.sv.BoxAnnotator",
+        side_effect=[mock_box_annotator_old, mock_box_annotator_new],
+    ), patch(
+        "perceptionmetrics.utils.image.sv.LabelAnnotator",
+        return_value=mock_label_annotator,
+    ), patch(
+        "perceptionmetrics.utils.image.sv.Color"
+    ), patch(
+        "perceptionmetrics.utils.image.sv.ColorPalette"
+    ):
+
+        draw_detections(img, boxes, class_ids, class_names)
+
+        assert mock_label_annotator.annotate.called
+        args, kwargs = mock_label_annotator.annotate.call_args
+        labels = kwargs.get("labels")
+        assert labels == ["5"]

tests/utils/test_torch.py

@@ -0,0 +1,76 @@
+import torch
+import pytest
+from perceptionmetrics.utils.torch import data_to_device, get_data_shape, unsqueeze_data
+
+
+def test_data_to_device():
+    # Setup
+    device = torch.device("cpu")
+    t1 = torch.randn(2, 2)
+    t2 = torch.randn(3, 3)
+    data = [t1, (t2,)]
+
+    # Execute
+    moved_data = data_to_device(data, device)
+
+    # Verify
+    assert torch.equal(moved_data[0], t1.to(device))
+    assert torch.equal(moved_data[1][0], t2.to(device))
+    assert isinstance(moved_data, list)
+    assert isinstance(moved_data[1], tuple)
+
+
+def test_get_data_shape():
+    # Setup
+    t1 = torch.randn(2, 3)
+    t2 = torch.randn(4, 5, 6)
+    data = (t1, [t2])
+
+    # Execute
+    shapes = get_data_shape(data)
+
+    # Verify
+    assert shapes == ((2, 3), [(4, 5, 6)])
+    assert isinstance(shapes, tuple)
+    assert isinstance(shapes[1], list)
+
+
+def test_unsqueeze_data():
+    # Setup
+    t1 = torch.randn(2, 2)
+    t2 = torch.randn(3, 3)
+    data = [t1, (t2,)]
+
+    # Execute
+    unsqueezed = unsqueeze_data(data, dim=0)
+
+    # Verify
+    assert unsqueezed[0].shape == (1, 2, 2)
+    assert unsqueezed[1][0].shape == (1, 3, 3)
+    assert isinstance(unsqueezed, list)
+    assert isinstance(unsqueezed[1], tuple)
+
+
+def test_torch_raises_type_error():
+    # Ensure non-tensors raise TypeError
+    data = "string"
+    device = torch.device("cpu")
+
+    with pytest.raises(TypeError, match="expected torch.Tensor"):
+        data_to_device(data, device)
+
+    with pytest.raises(TypeError, match="expected torch.Tensor"):
+        get_data_shape(data)
+
+    with pytest.raises(TypeError, match="expected torch.Tensor"):
+        unsqueeze_data(data)
+
+def test_torch_raises_type_error_nested():
+    # Test nested invalid types
+    data = [torch.randn(2), "invalid"]
+
+    with pytest.raises(TypeError, match="expected torch.Tensor"):
+        data_to_device(data, torch.device("cpu"))
+
+    with pytest.raises(TypeError, match="expected torch.Tensor"):
+        get_data_shape(data)