62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
import os
|
||
from typing import Tuple
|
||
|
||
import numpy as np
|
||
import torch
|
||
from wand.image import Image
|
||
|
||
from folder_paths import get_input_directory
|
||
|
||
|
||
MAX_COLOR_VALUE = 255.0 # Normalize pixel values to [0,1]
|
||
|
||
|
||
class DDSFileImport:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
input_dir = get_input_directory()
|
||
files = sorted(
|
||
f for f in os.listdir(input_dir)
|
||
if os.path.isfile(os.path.join(input_dir, f)) and f.lower().endswith(".dds")
|
||
)
|
||
return {"required": {"image": (files, {"image_upload": True})}}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
CATEGORY = "image"
|
||
FUNCTION = "import_dds"
|
||
|
||
def import_dds(self, image: str) -> Tuple[torch.Tensor,]:
|
||
base_dir = get_input_directory()
|
||
# Resolve the absolute path and guard against traversal
|
||
safe_path = os.path.abspath(os.path.join(base_dir, image))
|
||
if os.path.commonpath([safe_path, base_dir]) != base_dir:
|
||
raise ValueError("Invalid file path")
|
||
|
||
try:
|
||
with Image(filename=safe_path) as img:
|
||
# Ensure the file is a DDS texture
|
||
if img.format.upper() != "DDS":
|
||
raise ValueError(f"File {image} is not a DDS image")
|
||
|
||
img.alpha_channel = 'activate'
|
||
img.format = 'RGBA'
|
||
|
||
data = np.array(img, dtype=np.float32) / MAX_COLOR_VALUE
|
||
|
||
# Strip alpha when it’s fully opaque or fully transparent
|
||
if data.shape[2] == 4:
|
||
alpha = data[..., 3]
|
||
if np.all(alpha == 0) or np.all(alpha == 1):
|
||
data = data[..., :3]
|
||
|
||
tensor = torch.from_numpy(data)[None, ...] # [1,H,W,C]
|
||
except Exception as exc: # pragma: no cover
|
||
raise RuntimeError(f"Failed to load DDS file {image}") from exc
|
||
|
||
return (tensor,)
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"Load DDS File": DDSFileImport
|
||
}
|
||
|