feat(pipeline): add batch processing for scene graph construction

This commit is contained in:
2026-03-28 17:32:15 +08:00
parent 3c9a6f6eaf
commit f604c85a79
7 changed files with 252 additions and 44 deletions

View File

@@ -3,7 +3,7 @@ from unittest.mock import Mock
import torch
from PIL import Image
from utils.image import segment_image
from utils.image import segment_image, segment_image_dataset
def test_segment_image_passes_pil_image_to_mask_generator() -> None:
@@ -80,3 +80,36 @@ def test_segment_image_filters_tensor_masks_by_min_area() -> None:
assert len(result) == 1
assert result[0]["area"] == 4
def test_segment_image_dataset_returns_per_image_masks_in_order() -> None:
first_masks = {
"masks": torch.tensor(
[[[1, 1, 0], [1, 1, 0], [0, 0, 0]]],
dtype=torch.float32,
)
}
second_masks = {
"masks": torch.tensor(
[[[1, 1, 1], [1, 1, 1], [1, 1, 1]]],
dtype=torch.float32,
)
}
mock_generator = Mock(side_effect=[first_masks, second_masks])
images = [
Image.new("RGB", (3, 3), color=(0, 0, 0)),
Image.new("RGB", (3, 3), color=(0, 0, 0)),
]
result = segment_image_dataset(
mock_generator,
images,
min_area=2,
max_masks=5,
points_per_batch=16,
)
assert len(result) == 2
assert result[0][0]["area"] == 4
assert result[1][0]["area"] == 9
assert mock_generator.call_count == 2