tests/test_protocol.py

import unittest

from grounding_demo.protocol import parse_gam, parse_reference_boxes


class ProtocolTest(unittest.TestCase):
    def test_bbox(self):
        text = "<|object_ref_start|>car<|object_ref_end|><|box_start|><1><2><30><40>,<50><60><70><80><|box_end|>"
        result = parse_gam(text, "grounding")
        self.assertTrue(result.valid)
        self.assertEqual(len(result.primitives), 2)

    def test_point(self):
        text = "<|object_ref_start|>button<|object_ref_end|><|box_start|><500><600><|box_end|>"
        self.assertTrue(parse_gam(text, "gui").valid)

    def test_mask_point_tasks_parse_points_and_reject_boxes(self):
        point = "<|object_ref_start|>person<|object_ref_end|><|box_start|><500><600><|box_end|>"
        box = "<|object_ref_start|>person<|object_ref_end|><|box_start|><100><200><500><600><|box_end|>"
        for route in ("referring_point_in_mask", "dense_point_in_mask"):
            with self.subTest(route=route):
                self.assertTrue(parse_gam(point, route).valid)
                self.assertFalse(parse_gam(box, route).valid)

    def test_rejects_wrong_arity(self):
        text = "<|object_ref_start|>car<|object_ref_end|><|box_start|><1><2><3><|box_end|>"
        self.assertFalse(parse_gam(text, "grounding").valid)

    def test_visual_prompt_refs(self):
        prompt = "Given reference boxes <|box_start|><1><2><30><40><|box_end|> indicating objects"
        self.assertEqual(len(parse_reference_boxes(prompt)), 1)


if __name__ == "__main__":
    unittest.main()