File size: 6,110 Bytes
3afd6d6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | import os
import urllib.request
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
import folder_paths
def _center_crop_nchw(tensor: torch.Tensor, target_w: int, target_h: int) -> torch.Tensor:
_, _, h, w = tensor.shape
crop_top = max(0, (h - target_h) // 2)
crop_left = max(0, (w - target_w) // 2)
return tensor[:, :, crop_top:crop_top + target_h, crop_left:crop_left + target_w]
def _gaussian_kernel1d(sigma: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
radius = max(1, int(sigma * 3.0))
x = torch.arange(-radius, radius + 1, device=device, dtype=dtype)
kernel = torch.exp(-(x * x) / (2.0 * sigma * sigma))
return kernel / kernel.sum()
def _gaussian_blur_nchw(tensor: torch.Tensor, sigma: float) -> torch.Tensor:
if sigma <= 0:
return tensor
_, channels, _, _ = tensor.shape
kernel = _gaussian_kernel1d(sigma, tensor.device, tensor.dtype)
radius = kernel.shape[0] // 2
horizontal = kernel.view(1, 1, 1, -1).repeat(channels, 1, 1, 1)
vertical = kernel.view(1, 1, -1, 1).repeat(channels, 1, 1, 1)
blurred = F.pad(tensor, (radius, radius, 0, 0), mode="replicate")
blurred = F.conv2d(blurred, horizontal, groups=channels)
blurred = F.pad(blurred, (0, 0, radius, radius), mode="replicate")
return F.conv2d(blurred, vertical, groups=channels)
def _make_pillarbox_background(tensor: torch.Tensor, target_w: int, target_h: int) -> torch.Tensor:
_, _, h, w = tensor.shape
scale = max(target_w / w, target_h / h)
bg_w = max(target_w, int(round(w * scale)))
bg_h = max(target_h, int(round(h * scale)))
background = F.interpolate(tensor, size=(bg_h, bg_w), mode="bilinear", align_corners=False)
background = _center_crop_nchw(background, target_w, target_h)
sigma = 0.006 * min(target_h, target_w)
background = _gaussian_blur_nchw(background, sigma)
if background.shape[1] >= 3:
rgb = background[:, :3]
gray = rgb.mean(dim=1, keepdim=True)
background = background.clone()
background[:, :3] = (rgb * 0.8 + gray * 0.2) * 0.35
else:
background = background * 0.35
return background
def load_image_tensor(source_type: str, file_path: str | None, local_path: str | None, url: str | None) -> torch.Tensor | None:
"""Load a single image → float32 [1,H,W,C] tensor, or None on failure."""
from PIL import Image # type: ignore[import]
import io as _io
if source_type == "url" and url:
try:
with urllib.request.urlopen(url, timeout=15) as resp: # noqa: S310
pil_img = Image.open(_io.BytesIO(resp.read())).convert("RGB")
except Exception:
return None
else:
img_path: str | None = None
if source_type == "output" and file_path:
# output files use subfolder/filename format, need to join with output directory
output_dir = folder_paths.get_output_directory()
img_path = os.path.join(output_dir, file_path)
elif source_type == "input" and file_path:
img_path = folder_paths.get_annotated_filepath(file_path)
elif source_type == "local" and local_path:
img_path = local_path
if not img_path or not os.path.isfile(img_path):
return None
try:
pil_img = Image.open(img_path).convert("RGB")
except Exception:
return None
arr = np.array(pil_img).astype(np.float32) / 255.0
return torch.from_numpy(arr).unsqueeze(0) # [1,H,W,C]
def resize_image(tensor: torch.Tensor, target_w: int, target_h: int, method: str) -> torch.Tensor:
"""Resize [1,H,W,C] tensor to (target_h, target_w)."""
_, h, w, c = tensor.shape
t = tensor.permute(0, 3, 1, 2).float() # [1,C,H,W]
if method == "stretch":
out = F.interpolate(t, size=(target_h, target_w), mode="bilinear", align_corners=False)
elif method in ("resize", "pad", "pad (white)", "pad_edge", "pad_edge_pixel", "pillarbox_blur"):
scale = min(target_w / w, target_h / h)
new_w = max(1, int(round(w * scale)))
new_h = max(1, int(round(h * scale)))
resized = F.interpolate(t, size=(new_h, new_w), mode="bilinear", align_corners=False)
pad_top = (target_h - new_h) // 2
pad_left = (target_w - new_w) // 2
pad_bottom = target_h - new_h - pad_top
pad_right = target_w - new_w - pad_left
if method == "pad_edge_pixel":
out = F.pad(resized, (pad_left, pad_right, pad_top, pad_bottom), mode="replicate")
elif method == "pad_edge":
edge_pixels = torch.cat(
[
resized[:, :, 0, :],
resized[:, :, -1, :],
resized[:, :, :, 0],
resized[:, :, :, -1],
],
dim=2,
)
edge_color = edge_pixels.mean(dim=2).view(resized.shape[0], c, 1, 1)
out = edge_color.expand(-1, -1, target_h, target_w).clone()
out[:, :, pad_top:pad_top + new_h, pad_left:pad_left + new_w] = resized
elif method == "pillarbox_blur":
out = _make_pillarbox_background(t, target_w, target_h)
out[:, :, pad_top:pad_top + new_h, pad_left:pad_left + new_w] = resized
else:
fill_value = 1.0 if method == "pad (white)" else 0.0
out = torch.full((tensor.shape[0], c, target_h, target_w), fill_value, dtype=t.dtype, device=t.device)
out[:, :, pad_top:pad_top + new_h, pad_left:pad_left + new_w] = resized
elif method == "crop":
scale = max(target_w / w, target_h / h)
new_w = max(target_w, int(round(w * scale)))
new_h = max(target_h, int(round(h * scale)))
resized = F.interpolate(t, size=(new_h, new_w), mode="bilinear", align_corners=False)
out = _center_crop_nchw(resized, target_w, target_h)
else:
out = F.interpolate(t, size=(target_h, target_w), mode="bilinear", align_corners=False)
return out.permute(0, 2, 3, 1) # [1,H,W,C]
|