SAM2: image pipeline

tripy/examples/segment-anything-model-v2/README.md

@@ -22,6 +22,16 @@ This is an implementation of SAM2 model ([original repository](https://github.co
     python3 image_demo.py
     ```
 
+    <!-- Tripy: TEST: EXPECTED_STDOUT Start -->
+    <!--
+    ```
+    Generating image embedding took {137.81±10%} ms per run (averaged over 100 runs, with 5 warmup runs)
+    Predicting masks took {37.78±10%} ms per run (averaged over 100 runs, with 5 warmup runs)
+    Scores for each prediction: {0.78759766±5%} {0.640625±5%} {0.05099487±5%}
+    ```
+     -->
+    <!-- Tripy: TEST: EXPECTED_STDOUT End -->
+
 ### Video segmentation pipeline
 
 TBD

tripy/examples/segment-anything-model-v2/image_demo.py

@@ -143,6 +143,8 @@ def process_predictions(
             plt.savefig(os.path.join(save_path, f"mask_{i}_score_{score:.3f}.png"), bbox_inches="tight", pad_inches=0)
         plt.close(fig)
 
+    print(f"Scores for each prediction: {' '.join(map(str, scores))}")
+
     return {
         "masks": np.array(processed_masks),
         "scores": scores,
@@ -189,17 +191,16 @@ def time_function(func, num_warmup=5, num_runs=100, description=""):
         # Warmup runs
         for _ in range(num_warmup):
             func()
-
-        tp.default_stream().synchronize()
-        torch.cuda.synchronize()
+            tp.default_stream().synchronize()
+            torch.cuda.synchronize()
 
         # Actual timing
         start = time.perf_counter()
         for _ in range(num_runs):
             output = func()
+            tp.default_stream().synchronize()
+            torch.cuda.synchronize()
 
-        tp.default_stream().synchronize()
-        torch.cuda.synchronize()
         end = time.perf_counter()
 
         avg_time_ms = (end - start) * 1000 / num_runs
@@ -244,7 +245,6 @@ def predict_masks():
         input_labels=input_label,
         save_path=save_path,
     )
-
     return results
 
 

tripy/tests/test_examples.py

@@ -80,12 +80,43 @@ def __str__(self):
         return os.path.relpath(self.path, EXAMPLES_ROOT)
 
 
-EXAMPLES = [Example(["nanogpt"])]
+EXAMPLES = [Example(["nanogpt"]), Example(["segment-anything-model-v2"])]
 
 
 @pytest.mark.l1
 @pytest.mark.parametrize("example", EXAMPLES, ids=lambda case: str(case))
 def test_examples(example, sandboxed_install_run):
+
+    def test_with_tolerance(number, value, tolerance):
+        tolerance = float(tolerance) / 100
+        lower = float(number) * (1 - tolerance)
+        upper = float(number) * (1 + tolerance)
+        try:
+            return lower <= float(value) <= upper
+        except ValueError:
+            return False
+
+    def process_tolerances(expected_output):
+        specs = []
+        placeholder_regex = r"{(\d+\.?\d*)±(\d+)%}"
+        pattern = expected_output
+
+        # Replace tolerance patterns with more flexible capture group
+        matches = list(re.finditer(placeholder_regex, pattern))
+        for match in matches:
+            specs.append((match.group(1), match.group(2)))
+            pattern = pattern.replace(match.group(0), r"(\d+\.?\d*)", 1)
+
+        # Escape parentheses but not our capture group
+        pattern = pattern.replace("(", r"\(")
+        pattern = pattern.replace(")", r"\)")
+        pattern = pattern.replace(r"\(\d+\.?\d*\)", r"(\d+\.?\d*)")
+
+        # Make whitespace flexible
+        pattern = pattern.replace(" ", r"\s+")
+
+        return pattern.strip(), specs
+
     with open(example.readme, "r", encoding="utf-8") as f:
         contents = f.read()
         # Check that the README has all the expected sections.
@@ -101,9 +132,20 @@ def test_examples(example, sandboxed_install_run):
 
             code = str(block)
             if block.has_marker("test: expected_stdout"):
-                print("Checking command output against expected output: ", end="")
                 out = statuses[-1].stdout.strip()
-                matched = re.match(dedent(code).strip(), out)
+                expected = dedent(code).strip()
+                pattern, specs = process_tolerances(expected)
+
+                match = re.search(pattern, out)
+                if match and specs:
+                    # Check if captured numbers are within tolerance
+                    matched = all(
+                        test_with_tolerance(expected, actual, tolerance)
+                        for (expected, tolerance), actual in zip(specs, match.groups())
+                    )
+                else:
+                    matched = bool(match)
+
                 print("matched!" if matched else "did not match!")
                 print(f"==== STDOUT ====\n{out}")
                 assert matched