ykcheungab
94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from orisys.frame_processors import UVUnrollFrameProcessor
|
|
from orisys.unroll import UVDomain
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
UV_MAP_PATH = REPO_ROOT / "unroll_map_uv.npz"
|
|
|
|
|
|
@pytest.fixture
|
|
def uv_domain():
|
|
if not UV_MAP_PATH.is_file():
|
|
pytest.skip(f"UV remap cache not found: {UV_MAP_PATH}")
|
|
return UVDomain.load(str(UV_MAP_PATH))
|
|
|
|
|
|
def test_uv_unroll_processor_output_shape(uv_domain):
|
|
uv_remap = uv_domain.remap
|
|
src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
|
|
src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))
|
|
|
|
processor = UVUnrollFrameProcessor(
|
|
uv_remap,
|
|
source_width=src_w,
|
|
source_height=src_h,
|
|
flip_source_vertical=True,
|
|
border_value=(0, 0, 0),
|
|
)
|
|
|
|
frame = np.random.randint(0, 255, (src_h, src_w, 3), dtype=np.uint8)
|
|
unrolled = processor.process(frame)
|
|
|
|
assert unrolled.shape == (uv_remap.out_height, uv_remap.out_width, 3)
|
|
assert processor.valid_mask.shape == (uv_remap.out_height, uv_remap.out_width)
|
|
assert unrolled.dtype == np.uint8
|
|
|
|
|
|
def test_uv_unroll_processor_prepare_source(uv_domain):
|
|
uv_remap = uv_domain.remap
|
|
src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
|
|
src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))
|
|
|
|
processor = UVUnrollFrameProcessor(
|
|
uv_remap,
|
|
source_width=src_w,
|
|
source_height=src_h,
|
|
flip_source_vertical=False,
|
|
)
|
|
frame = np.ones((src_h, src_w, 3), dtype=np.uint8) * 128
|
|
prepared = processor.prepare_source(frame)
|
|
assert prepared.shape == (src_h, src_w, 3)
|
|
|
|
|
|
def test_uv_unroll_processor_flatten_prepared_matches_process(uv_domain):
|
|
uv_remap = uv_domain.remap
|
|
src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
|
|
src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))
|
|
|
|
processor = UVUnrollFrameProcessor(
|
|
uv_remap,
|
|
source_width=src_w,
|
|
source_height=src_h,
|
|
flip_source_vertical=True,
|
|
border_value=(0, 0, 0),
|
|
)
|
|
frame = np.random.randint(0, 255, (src_h, src_w, 3), dtype=np.uint8)
|
|
prepared = processor.prepare_source(frame)
|
|
assert np.array_equal(processor.flatten_prepared(prepared), processor.process(frame))
|
|
|
|
|
|
def test_uv_unroll_processor_gray_fast_path_matches_bgr_result(uv_domain):
|
|
uv_remap = uv_domain.remap
|
|
src_h = int(uv_remap.metadata.get("source_height", uv_remap.out_height))
|
|
src_w = int(uv_remap.metadata.get("source_width", uv_remap.out_width))
|
|
|
|
processor = UVUnrollFrameProcessor(
|
|
uv_remap,
|
|
source_width=src_w,
|
|
source_height=src_h,
|
|
flip_source_vertical=True,
|
|
border_value=(0, 0, 0),
|
|
)
|
|
frame = np.random.randint(0, 255, (src_h, src_w, 3), dtype=np.uint8)
|
|
prepared_gray = processor.prepare_source_gray(frame)
|
|
gray_unrolled = processor.flatten_prepared_gray(prepared_gray)
|
|
bgr_unrolled = processor.process(frame)
|
|
assert gray_unrolled.shape == (uv_remap.out_height, uv_remap.out_width)
|
|
assert np.allclose(gray_unrolled, cv2.cvtColor(bgr_unrolled, cv2.COLOR_BGR2GRAY), atol=1)
|