91 lines
3.0 KiB
Python
91 lines
3.0 KiB
Python
|
import tempfile
|
||
|
|
import unittest
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
from cockpit_agent.grounding.adapter import GroundingResult
|
||
|
|
from cockpit_agent.perception.roi import crop_grounding_roi, grounding_to_source
|
||
|
|
|
||
|
|
|
||
|
|
def grounding(
|
||
|
|
*,
|
||
|
|
bbox_pixel: tuple[int, int, int, int],
|
||
|
|
width: int,
|
||
|
|
height: int,
|
||
|
|
query: str = "semantic function",
|
||
|
|
) -> GroundingResult:
|
||
|
|
x1, y1, x2, y2 = bbox_pixel
|
||
|
|
return GroundingResult(
|
||
|
|
semantic_query=query,
|
||
|
|
bbox_relative=(
|
||
|
|
x1 / width * 1000,
|
||
|
|
y1 / height * 1000,
|
||
|
|
x2 / width * 1000,
|
||
|
|
y2 / height * 1000,
|
||
|
|
),
|
||
|
|
bbox_pixel=bbox_pixel,
|
||
|
|
center_pixel=(round((x1 + x2) / 2), round((y1 + y2) / 2)),
|
||
|
|
image_width=width,
|
||
|
|
image_height=height,
|
||
|
|
raw_output="model output",
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class ROITest(unittest.TestCase):
|
||
|
|
def test_function_bbox_is_padded_and_cropped(self) -> None:
|
||
|
|
with tempfile.TemporaryDirectory() as directory:
|
||
|
|
root = Path(directory)
|
||
|
|
source = root / "source.png"
|
||
|
|
output = root / "roi.jpg"
|
||
|
|
Image.new("RGB", (200, 100), color="white").save(source)
|
||
|
|
result = crop_grounding_roi(
|
||
|
|
image_path=str(source),
|
||
|
|
function_grounding=grounding(
|
||
|
|
bbox_pixel=(50, 20, 150, 80),
|
||
|
|
width=200,
|
||
|
|
height=100,
|
||
|
|
),
|
||
|
|
output_path=str(output),
|
||
|
|
padding_ratio=0.10,
|
||
|
|
)
|
||
|
|
|
||
|
|
with Image.open(output) as roi_image:
|
||
|
|
roi_size = roi_image.size
|
||
|
|
|
||
|
|
self.assertEqual(result.bbox_pixel, (40, 14, 160, 86))
|
||
|
|
self.assertEqual((result.offset_x, result.offset_y), (40, 14))
|
||
|
|
self.assertEqual((result.width, result.height), (120, 72))
|
||
|
|
self.assertEqual(roi_size, (120, 72))
|
||
|
|
|
||
|
|
def test_roi_grounding_is_converted_to_source_coordinates(self) -> None:
|
||
|
|
with tempfile.TemporaryDirectory() as directory:
|
||
|
|
root = Path(directory)
|
||
|
|
source = root / "source.png"
|
||
|
|
Image.new("RGB", (200, 100), color="white").save(source)
|
||
|
|
roi = crop_grounding_roi(
|
||
|
|
image_path=str(source),
|
||
|
|
function_grounding=grounding(
|
||
|
|
bbox_pixel=(50, 20, 150, 80),
|
||
|
|
width=200,
|
||
|
|
height=100,
|
||
|
|
),
|
||
|
|
output_path=str(root / "roi.jpg"),
|
||
|
|
padding_ratio=0.10,
|
||
|
|
)
|
||
|
|
roi_grounding = grounding(
|
||
|
|
bbox_pixel=(10, 6, 30, 26),
|
||
|
|
width=roi.width,
|
||
|
|
height=roi.height,
|
||
|
|
query="semantic subcontrol",
|
||
|
|
)
|
||
|
|
converted = grounding_to_source(roi_grounding, roi)
|
||
|
|
|
||
|
|
self.assertEqual(converted.bbox_pixel, (50, 20, 70, 40))
|
||
|
|
self.assertEqual(converted.center_pixel, (60, 30))
|
||
|
|
self.assertEqual((converted.image_width, converted.image_height), (200, 100))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|