67 lines
1.8 KiB
Python
67 lines
1.8 KiB
Python
"""Shared pytest fixtures for ImageOps tests."""
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import pytest
|
|
|
|
|
|
@pytest.fixture
|
|
def rgb_image_1x64():
|
|
"""Single-frame 64x64 RGB test image with a gradient pattern."""
|
|
img = torch.zeros(1, 64, 64, 3, dtype=torch.float32)
|
|
# Horizontal red gradient, vertical green gradient
|
|
img[..., 0] = torch.linspace(0, 1, 64).unsqueeze(0)
|
|
img[..., 1] = torch.linspace(0, 1, 64).unsqueeze(1)
|
|
img[..., 2] = 0.25
|
|
return img
|
|
|
|
|
|
@pytest.fixture
|
|
def rgba_image_1x64():
|
|
"""Single-frame 64x64 RGBA test image with semi-transparent alpha."""
|
|
img = torch.zeros(1, 64, 64, 4, dtype=torch.float32)
|
|
img[..., 0] = torch.linspace(0, 1, 64).unsqueeze(0)
|
|
img[..., 1] = torch.linspace(0, 1, 64).unsqueeze(1)
|
|
img[..., 2] = 0.25
|
|
img[..., 3] = 0.8
|
|
return img
|
|
|
|
|
|
@pytest.fixture
|
|
def rgb_batch_4x64():
|
|
"""4-frame 64x64 RGB batch."""
|
|
return torch.rand(4, 64, 64, 3, dtype=torch.float32)
|
|
|
|
|
|
@pytest.fixture
|
|
def rgba_batch_4x64():
|
|
"""4-frame 64x64 RGBA batch."""
|
|
return torch.rand(4, 64, 64, 4, dtype=torch.float32)
|
|
|
|
|
|
@pytest.fixture
|
|
def mask_1x64():
|
|
"""Single-frame 64x64 mask with a radial gradient."""
|
|
ys = torch.linspace(-1, 1, 64)
|
|
xs = torch.linspace(-1, 1, 64)
|
|
gy, gx = torch.meshgrid(ys, xs, indexing="ij")
|
|
r = (gx * gx + gy * gy).sqrt().clamp(0, 1)
|
|
return (1.0 - r).unsqueeze(0) # [1, 64, 64]
|
|
|
|
|
|
@pytest.fixture
|
|
def mask_4x64():
|
|
"""4-frame 64x64 mask batch."""
|
|
return torch.rand(4, 64, 64, dtype=torch.float32)
|
|
|
|
|
|
@pytest.fixture
|
|
def hdr_image_1x64():
|
|
"""Single-frame 64x64 RGB image with HDR values (outside [0,1])."""
|
|
img = torch.zeros(1, 64, 64, 3, dtype=torch.float32)
|
|
img[0, :32, :, 0] = 2.0 # bright red
|
|
img[0, 32:, :, 1] = -0.25 # negative green
|
|
img[0, :, :32, 2] = 4.0 # very bright blue
|
|
img[0, :, 32:, 2] = 0.5
|
|
return img
|