From 18d54cfc0c799dd2a984e9aa187bb9f8b3b9e0ca Mon Sep 17 00:00:00 2001 From: Dowon Date: Thu, 5 Feb 2026 19:56:23 +0900 Subject: [PATCH] test: mask_to_pil test --- tests/test_ultralytics.py | 36 +++++++++++++++++++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/tests/test_ultralytics.py b/tests/test_ultralytics.py index e75171d..6d62fa0 100644 --- a/tests/test_ultralytics.py +++ b/tests/test_ultralytics.py @@ -1,8 +1,10 @@ +import numpy as np import pytest +import torch from huggingface_hub import hf_hub_download from PIL import Image -from adetailer.ultralytics import ultralytics_predict +from adetailer.ultralytics import mask_to_pil, ultralytics_predict @pytest.mark.parametrize( @@ -60,3 +62,35 @@ def test_yolo_world(sample_image2: Image.Image, klass: str): assert len(result.masks) > 0 assert len(result.confidences) > 0 assert len(result.bboxes) == len(result.masks) == len(result.confidences) + + +class TestMaskToPil: + def test_mask_to_pil_float32(self): + mask = torch.tensor([[[0.0, 1.0], [0.0, 1.0]]], dtype=torch.float32) + imgs = mask_to_pil(mask, shape=(2, 2)) + + assert len(imgs) == 1 + img = imgs[0] + assert isinstance(img, Image.Image) + + arr = np.array(img) + assert arr.shape == (2, 2) + assert arr.dtype == np.uint8 + + expected = np.array([[0, 255], [0, 255]], dtype=np.uint8) + np.testing.assert_array_equal(arr, expected) + + def test_mask_to_pil_uint8(self): + mask = torch.tensor([[[0, 1], [0, 1]]], dtype=torch.uint8) + imgs = mask_to_pil(mask, shape=(2, 2)) + + assert len(imgs) == 1 + img = imgs[0] + assert isinstance(img, Image.Image) + + arr = np.array(img) + assert arr.shape == (2, 2) + assert arr.dtype == np.uint8 + + expected = np.array([[0, 255], [0, 255]], dtype=np.uint8) + np.testing.assert_array_equal(arr, expected)