Initial commit

fbshipit-source-id: da6be2f26e3a1202f4bffde8cb980e2dcb851294

sam3/perflib/tests/tests.py

@@ -0,0 +1,59 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+import os
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from sam3.perflib.masks_ops import masks_to_boxes
+
+
+class TestMasksToBoxes:
+    def test_masks_box(self):
+        def masks_box_check(masks, expected, atol=1e-4):
+            out = masks_to_boxes(masks, [1 for _ in range(masks.shape[0])])
+            assert out.dtype == torch.float
+            print("out: ", out)
+            print("expected: ", expected)
+            torch.testing.assert_close(
+                out, expected, rtol=0.0, check_dtype=True, atol=atol
+            )
+
+        # Check for int type boxes.
+        def _get_image():
+            assets_directory = os.path.join(
+                os.path.dirname(os.path.abspath(__file__)), "assets"
+            )
+            mask_path = os.path.join(assets_directory, "masks.tiff")
+            image = Image.open(mask_path)
+            return image
+
+        def _create_masks(image, masks):
+            for index in range(image.n_frames):
+                image.seek(index)
+                frame = np.array(image)
+                masks[index] = torch.tensor(frame)
+
+            return masks
+
+        expected = torch.tensor(
+            [
+                [127, 2, 165, 40],
+                [2, 50, 44, 92],
+                [56, 63, 98, 100],
+                [139, 68, 175, 104],
+                [160, 112, 198, 145],
+                [49, 138, 99, 182],
+                [108, 148, 152, 213],
+            ],
+            dtype=torch.float,
+        )
+
+        image = _get_image()
+        for dtype in [torch.float16, torch.float32, torch.float64]:
+            masks = torch.zeros(
+                (image.n_frames, image.height, image.width), dtype=dtype
+            )
+            masks = _create_masks(image, masks)
+            masks_box_check(masks, expected)