tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
Code

Simplify Pillow encoder compatibility

Browse Source
This commit is contained in:
aswanth-07 2026-08-28 17:40:40 +05:30
parent 02105c8c63
commit 0caa292600
2 changed files with 4 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
nerfstudio/data/utils/data_utils.py
View File

@ -20,20 +20,9 @@ from typing import IO, List, Tuple, Union
import cv2
import numpy as np
import torch
from packaging.version import Version
from PIL import Image, __version__ as PILLOW_VERSION
from PIL import Image
from PIL.Image import Image as PILImage
_PILLOW_VERSION = Version(PILLOW_VERSION)
def _set_encoder_image(encoder, im: PILImage) -> None:
"""Set the image on a Pillow raw encoder across supported Pillow versions."""
if _PILLOW_VERSION >= Version("12.2.0"):
encoder.setimage(im.im, (0, 0) + im.size)
else:
encoder.setimage(im.im)
def pil_to_numpy(im: PILImage) -> np.ndarray:
"""Converts a PIL Image object to a NumPy array.
@ -49,7 +38,7 @@ def pil_to_numpy(im: PILImage) -> np.ndarray:
# Unpack data
e = Image._getencoder(im.mode, "raw", im.mode)
_set_encoder_image(e, im)
e.setimage(im.im, (0, 0) + im.size)
# NumPy buffer for the result
shape, typestr = Image._conv_type_shape(im)

21
tests/utils/test_data_utils.py
View File

@ -12,28 +12,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest.mock import Mock
import numpy as np
import pytest
from packaging.version import Version
from PIL import Image
from nerfstudio.data.utils import data_utils
@pytest.mark.parametrize("pillow_version, uses_extents", [("12.2.0", True), ("12.1.0", False)])
def test_set_encoder_image(monkeypatch, pillow_version, uses_extents):
image = Image.new("L", (2, 2))
encoder = Mock()
monkeypatch.setattr(data_utils, "_PILLOW_VERSION", Version(pillow_version))
data_utils._set_encoder_image(encoder, image)
if uses_extents:
encoder.setimage.assert_called_once_with(image.im, (0, 0) + image.size)
else:
encoder.setimage.assert_called_once_with(image.im)
from nerfstudio.data.utils.data_utils import pil_to_numpy
@pytest.mark.parametrize(
@ -49,7 +32,7 @@ def test_set_encoder_image(monkeypatch, pillow_version, uses_extents):
def test_pil_to_numpy(values):
image = Image.fromarray(values)
result = data_utils.pil_to_numpy(image)
result = pil_to_numpy(image)
np.testing.assert_array_equal(result, values)
assert result.flags.writeable