Fix cv2 lint errors

This commit is contained in:
Max Tretikov 2024-06-13 23:40:44 -06:00
parent f9c3779d4e
commit 5e89252ab6
2 changed files with 8 additions and 5 deletions

View File

@ -1,5 +1,6 @@
import os.path
from contextlib import contextmanager
from typing import Iterator
import cv2
from PIL import Image
@ -8,11 +9,11 @@ from . import node_helpers
def _open_exr(exr_path) -> Image.Image:
return Image.fromarray(cv2.imread(exr_path, cv2.IMREAD_COLOR))
return Image.fromarray(cv2.imread(exr_path, cv2.IMREAD_COLOR)) # pylint: disable=no-member
@contextmanager
def open_image(file_path: str) -> Image.Image:
def open_image(file_path: str) -> Iterator[Image.Image]:
_, ext = os.path.splitext(file_path)
if ext == ".exr":
yield _open_exr(file_path)

View File

@ -36,6 +36,8 @@ from torch import Tensor
from .component_model.images_types import RgbMaskTuple
read_exr = lambda fp: cv.imread(fp, cv.IMREAD_UNCHANGED).astype(np.float32) # pylint: disable=no-member
def mut_srgb_to_linear(np_array) -> None:
less = np_array <= 0.0404482362771082
np_array[less] = np_array[less] / 12.92
@ -49,7 +51,7 @@ def mut_linear_to_srgb(np_array) -> None:
def load_exr(file_path: str, srgb: bool) -> RgbMaskTuple:
image = cv.imread(file_path, cv.IMREAD_UNCHANGED).astype(np.float32)
image = read_exr(file_path)
rgb = np.flip(image[:, :, :3], 2).copy()
if srgb:
mut_linear_to_srgb(rgb)
@ -64,7 +66,7 @@ def load_exr(file_path: str, srgb: bool) -> RgbMaskTuple:
def load_exr_latent(file_path: str) -> Tuple[Tensor]:
image = cv.imread(file_path, cv.IMREAD_UNCHANGED).astype(np.float32)
image = read_exr(file_path)
image = image[:, :, np.array([2, 1, 0, 3])]
image = torch.unsqueeze(torch.from_numpy(image), 0)
image = torch.movedim(image, -1, 1)
@ -83,4 +85,4 @@ def save_exr(images: Tensor, filepaths_batched: Sequence[str], colorspace="linea
bgr[:, :, :, 3] = np.clip(1 - linear[:, :, :, 3], 0, 1) # invert alpha
for i in range(len(linear.shape[0])):
cv.imwrite(filepaths_batched[i], bgr[i])
cv.imwrite(filepaths_batched[i], bgr[i]) # pylint: disable=no-member