wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-24 13:20:19 +08:00
Code Issues Actions 12 Packages Projects Releases Wiki Activity

Fix cv2 lint errors

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
comfy/images.py
View File

@ -1,5 +1,6 @@
import os.path import os.path
from contextlib import contextmanager from contextlib import contextmanager
from typing import Iterator
import cv2 import cv2
from PIL import Image from PIL import Image
@ -8,11 +9,11 @@ from . import node_helpers
def _open_exr(exr_path) -> Image.Image: 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 @contextmanager
def open_image(file_path: str) -> Image.Image: def open_image(file_path: str) -> Iterator[Image.Image]:
_, ext = os.path.splitext(file_path) _, ext = os.path.splitext(file_path)
if ext == ".exr": if ext == ".exr":
yield _open_exr(file_path) yield _open_exr(file_path)

8
comfy/open_exr.py
View File

@ -36,6 +36,8 @@ from torch import Tensor
from .component_model.images_types import RgbMaskTuple 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: def mut_srgb_to_linear(np_array) -> None:
less = np_array <= 0.0404482362771082 less = np_array <= 0.0404482362771082
np_array[less] = np_array[less] / 12.92 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: 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() rgb = np.flip(image[:, :, :3], 2).copy()
if srgb: if srgb:
mut_linear_to_srgb(rgb) 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]: 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 = image[:, :, np.array([2, 1, 0, 3])]
image = torch.unsqueeze(torch.from_numpy(image), 0) image = torch.unsqueeze(torch.from_numpy(image), 0)
image = torch.movedim(image, -1, 1) 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 bgr[:, :, :, 3] = np.clip(1 - linear[:, :, :, 3], 0, 1) # invert alpha
for i in range(len(linear.shape[0])): 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