52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
|
import unittest
|
||
|
|
||
|
from PIL import Image
|
||
|
|
||
|
from onnx_web.chain.blend_denoise_fastnlmeans import BlendDenoiseFastNLMeansStage
|
||
|
from onnx_web.chain.result import ImageMetadata, StageResult
|
||
|
from tests.helpers import test_params, test_size
|
||
|
|
||
|
|
||
|
class TestBlendDenoiseFastNLMeansStage(unittest.TestCase):
|
||
|
def test_run(self):
|
||
|
# Create a dummy image
|
||
|
size = test_size()
|
||
|
image = Image.new("RGB", (size.width, size.height), color="white")
|
||
|
|
||
|
# Create a dummy StageResult object
|
||
|
sources = StageResult.from_images(
|
||
|
[image],
|
||
|
metadata=[
|
||
|
ImageMetadata(
|
||
|
test_params(),
|
||
|
size,
|
||
|
)
|
||
|
],
|
||
|
)
|
||
|
|
||
|
# Create an instance of BlendDenoiseLocalStdStage
|
||
|
stage = BlendDenoiseFastNLMeansStage()
|
||
|
|
||
|
# Call the run method with dummy parameters
|
||
|
result = stage.run(
|
||
|
_worker=None,
|
||
|
_server=None,
|
||
|
_stage=None,
|
||
|
_params=None,
|
||
|
sources=sources,
|
||
|
strength=5,
|
||
|
range=4,
|
||
|
stage_source=None,
|
||
|
callback=None,
|
||
|
)
|
||
|
|
||
|
# Assert that the result is an instance of StageResult
|
||
|
self.assertIsInstance(result, StageResult)
|
||
|
|
||
|
# Assert that the result contains the denoised image
|
||
|
self.assertEqual(len(result), 1)
|
||
|
self.assertEqual(result.size(), size)
|
||
|
|
||
|
# Assert that the metadata is preserved
|
||
|
self.assertEqual(result.metadata, sources.metadata)
|