1067 lines
38 KiB
Python
1067 lines
38 KiB
Python
"""Input processing for webhook data."""
|
|
|
|
import asyncio
|
|
import base64
|
|
import io
|
|
import json
|
|
import logging
|
|
import re
|
|
from typing import Any, List, Optional, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageOps
|
|
|
|
try:
|
|
import aiohttp
|
|
except ImportError:
|
|
aiohttp = None
|
|
|
|
try:
|
|
import torchaudio
|
|
except ImportError:
|
|
torchaudio = None
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DATA_URL_PATTERN = re.compile(r'^data:([^;,]+)?(;base64)?,(.*)$', re.DOTALL)
|
|
HTTP_URL_PATTERN = re.compile(r'^https?://', re.IGNORECASE)
|
|
|
|
|
|
class WebhookInputProcessor:
|
|
"""Processes webhook input data into ComfyUI-compatible formats."""
|
|
|
|
def __init__(self, download_timeout: float = 30.0):
|
|
self.download_timeout = download_timeout
|
|
|
|
async def process_inputs(self, webhook_inputs: dict) -> dict:
|
|
"""
|
|
Process all inputs and return ComfyUI-compatible values.
|
|
|
|
Args:
|
|
webhook_inputs: Dictionary of input name -> value
|
|
|
|
Returns:
|
|
Dictionary of input name -> processed value
|
|
"""
|
|
processed = {}
|
|
|
|
for key, value in webhook_inputs.items():
|
|
try:
|
|
processed[key] = await self._process_value(value)
|
|
except Exception as e:
|
|
logger.error(f"Failed to process input '{key}': {e}")
|
|
# Keep original value on error
|
|
processed[key] = value
|
|
|
|
return processed
|
|
|
|
async def _process_value(self, value: Any) -> Any:
|
|
"""Process a single value based on its type."""
|
|
if isinstance(value, str):
|
|
return await self._process_string(value)
|
|
elif isinstance(value, list):
|
|
return await self._process_array(value)
|
|
elif isinstance(value, dict):
|
|
# JSON data - serialize to string
|
|
return json.dumps(value)
|
|
else:
|
|
# Numbers, booleans - pass through
|
|
return value
|
|
|
|
async def _process_string(self, value: str) -> Any:
|
|
"""Process a string value, detecting if it's base64/URL data."""
|
|
if self._is_base64_image(value):
|
|
return await self._decode_base64_image(value)
|
|
elif self._is_base64_audio(value):
|
|
return await self._decode_base64_audio(value)
|
|
elif self._is_url(value):
|
|
return await self._download_and_decode(value)
|
|
else:
|
|
# Plain string
|
|
return value
|
|
|
|
async def _process_array(self, items: list) -> Union[torch.Tensor, list]:
|
|
"""
|
|
Process array of items.
|
|
|
|
If all items are images, returns a batched tensor [N, H, W, C].
|
|
Otherwise returns a list of processed items.
|
|
"""
|
|
if not items:
|
|
return []
|
|
|
|
# Check if all items are image data
|
|
if all(self._is_image_data(item) for item in items):
|
|
tensors = []
|
|
for item in items:
|
|
if isinstance(item, str):
|
|
if self._is_base64_image(item):
|
|
tensor = await self._decode_base64_image(item)
|
|
elif self._is_url(item):
|
|
tensor = await self._download_and_decode(item)
|
|
else:
|
|
continue
|
|
tensors.append(tensor)
|
|
|
|
if tensors:
|
|
# Stack into batch tensor [N, H, W, C]
|
|
return self._batch_images(tensors)
|
|
return []
|
|
|
|
# Otherwise process each item individually
|
|
return [await self._process_value(item) for item in items]
|
|
|
|
def _batch_images(self, tensors: List[torch.Tensor]) -> torch.Tensor:
|
|
"""
|
|
Batch multiple image tensors, handling different sizes.
|
|
|
|
Args:
|
|
tensors: List of [1, H, W, C] tensors
|
|
|
|
Returns:
|
|
Batched tensor [N, H, W, C]
|
|
"""
|
|
if not tensors:
|
|
return torch.zeros(0, 64, 64, 3)
|
|
|
|
# Check if all images have the same size
|
|
shapes = [t.shape[1:3] for t in tensors]
|
|
if all(s == shapes[0] for s in shapes):
|
|
# All same size, simple cat
|
|
return torch.cat(tensors, dim=0)
|
|
|
|
# Different sizes - pad to largest
|
|
max_h = max(t.shape[1] for t in tensors)
|
|
max_w = max(t.shape[2] for t in tensors)
|
|
|
|
padded = []
|
|
for t in tensors:
|
|
h, w = t.shape[1], t.shape[2]
|
|
if h < max_h or w < max_w:
|
|
# Center pad
|
|
pad_h_before = (max_h - h) // 2
|
|
pad_h_after = max_h - h - pad_h_before
|
|
pad_w_before = (max_w - w) // 2
|
|
pad_w_after = max_w - w - pad_w_before
|
|
# F.pad expects (left, right, top, bottom) for 4D tensor
|
|
t = torch.nn.functional.pad(
|
|
t.permute(0, 3, 1, 2), # [B, C, H, W]
|
|
(pad_w_before, pad_w_after, pad_h_before, pad_h_after),
|
|
mode='constant',
|
|
value=0
|
|
).permute(0, 2, 3, 1) # Back to [B, H, W, C]
|
|
padded.append(t)
|
|
|
|
return torch.cat(padded, dim=0)
|
|
|
|
# Detection methods
|
|
|
|
def _is_base64_image(self, value: str) -> bool:
|
|
"""Check if string is a base64-encoded image."""
|
|
if value.startswith("data:image/"):
|
|
return True
|
|
# Check for raw base64 that might be an image (heuristic)
|
|
if len(value) > 100 and re.match(r'^[A-Za-z0-9+/]+=*$', value):
|
|
# Try to detect PNG/JPEG magic bytes
|
|
try:
|
|
decoded = base64.b64decode(value[:32])
|
|
# PNG magic: 89 50 4E 47
|
|
if decoded.startswith(b'\x89PNG'):
|
|
return True
|
|
# JPEG magic: FF D8 FF
|
|
if decoded.startswith(b'\xff\xd8\xff'):
|
|
return True
|
|
# WebP magic: RIFF....WEBP
|
|
if decoded.startswith(b'RIFF') and b'WEBP' in decoded[:16]:
|
|
return True
|
|
except:
|
|
pass
|
|
return False
|
|
|
|
def _is_base64_audio(self, value: str) -> bool:
|
|
"""Check if string is a base64-encoded audio."""
|
|
return value.startswith("data:audio/")
|
|
|
|
def _is_url(self, value: str) -> bool:
|
|
"""Check if string is an HTTP(S) URL."""
|
|
return bool(HTTP_URL_PATTERN.match(value))
|
|
|
|
def _is_image_data(self, value: Any) -> bool:
|
|
"""Check if value is image data (base64 or URL)."""
|
|
if not isinstance(value, str):
|
|
return False
|
|
return self._is_base64_image(value) or (
|
|
self._is_url(value) and
|
|
any(ext in value.lower() for ext in ['.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp'])
|
|
)
|
|
|
|
# Decoding methods
|
|
|
|
async def _decode_base64_image(self, data: str) -> torch.Tensor:
|
|
"""
|
|
Decode base64 image to tensor.
|
|
|
|
Supports:
|
|
- data:image/png;base64,...
|
|
- data:image/jpeg;base64,...
|
|
- data:image/webp;base64,...
|
|
- Raw base64 (no prefix)
|
|
|
|
Returns:
|
|
Tensor [1, H, W, C] with values 0-1
|
|
"""
|
|
# Strip data URL prefix if present
|
|
if data.startswith("data:"):
|
|
match = DATA_URL_PATTERN.match(data)
|
|
if match:
|
|
data = match.group(3)
|
|
else:
|
|
_, data = data.split(",", 1)
|
|
|
|
# Decode base64
|
|
image_bytes = base64.b64decode(data)
|
|
|
|
# Load with PIL
|
|
pil_image = Image.open(io.BytesIO(image_bytes))
|
|
pil_image = ImageOps.exif_transpose(pil_image)
|
|
|
|
# Convert to RGB if needed
|
|
if pil_image.mode == "I":
|
|
# 16-bit images
|
|
pil_image = pil_image.point(lambda i: i * (1 / 255))
|
|
if pil_image.mode != "RGB":
|
|
pil_image = pil_image.convert("RGB")
|
|
|
|
# Convert to tensor [1, H, W, C] with values 0-1
|
|
np_image = np.array(pil_image).astype(np.float32) / 255.0
|
|
tensor = torch.from_numpy(np_image).unsqueeze(0)
|
|
|
|
return tensor
|
|
|
|
async def _decode_base64_audio(self, data: str) -> dict:
|
|
"""
|
|
Decode base64 audio to ComfyUI audio format.
|
|
|
|
Returns:
|
|
{"waveform": tensor [1, channels, samples], "sample_rate": int}
|
|
"""
|
|
if torchaudio is None:
|
|
raise ImportError("torchaudio is required for audio processing")
|
|
|
|
# Strip data URL prefix if present
|
|
if data.startswith("data:"):
|
|
_, data = data.split(",", 1)
|
|
|
|
audio_bytes = base64.b64decode(data)
|
|
|
|
# Use torchaudio to load
|
|
buffer = io.BytesIO(audio_bytes)
|
|
waveform, sample_rate = torchaudio.load(buffer)
|
|
|
|
# Resample to 44100 if needed (common ComfyUI sample rate)
|
|
target_rate = 44100
|
|
if sample_rate != target_rate:
|
|
resampler = torchaudio.transforms.Resample(sample_rate, target_rate)
|
|
waveform = resampler(waveform)
|
|
sample_rate = target_rate
|
|
|
|
return {
|
|
"waveform": waveform.unsqueeze(0), # [1, channels, samples]
|
|
"sample_rate": sample_rate
|
|
}
|
|
|
|
async def _download_and_decode(self, url: str) -> Any:
|
|
"""
|
|
Download content from URL and decode based on content type.
|
|
|
|
Returns:
|
|
Decoded content (tensor for images, dict for audio, etc.)
|
|
"""
|
|
if aiohttp is None:
|
|
raise ImportError("aiohttp is required for URL downloads")
|
|
|
|
timeout = aiohttp.ClientTimeout(total=self.download_timeout)
|
|
|
|
async with aiohttp.ClientSession(timeout=timeout) as session:
|
|
async with session.get(url) as response:
|
|
response.raise_for_status()
|
|
content_type = response.headers.get("Content-Type", "").lower()
|
|
data = await response.read()
|
|
|
|
# Detect type from content-type header or URL
|
|
if "image" in content_type or self._looks_like_image_url(url):
|
|
return await self._decode_image_bytes(data)
|
|
elif "audio" in content_type or self._looks_like_audio_url(url):
|
|
return await self._decode_audio_bytes(data)
|
|
else:
|
|
# Return as string
|
|
return data.decode("utf-8", errors="replace")
|
|
|
|
async def _decode_image_bytes(self, data: bytes) -> torch.Tensor:
|
|
"""Decode image bytes to tensor."""
|
|
pil_image = Image.open(io.BytesIO(data))
|
|
pil_image = ImageOps.exif_transpose(pil_image)
|
|
|
|
if pil_image.mode == "I":
|
|
pil_image = pil_image.point(lambda i: i * (1 / 255))
|
|
if pil_image.mode != "RGB":
|
|
pil_image = pil_image.convert("RGB")
|
|
|
|
np_image = np.array(pil_image).astype(np.float32) / 255.0
|
|
tensor = torch.from_numpy(np_image).unsqueeze(0)
|
|
return tensor
|
|
|
|
async def _decode_audio_bytes(self, data: bytes) -> dict:
|
|
"""Decode audio bytes to ComfyUI format."""
|
|
if torchaudio is None:
|
|
raise ImportError("torchaudio is required for audio processing")
|
|
|
|
buffer = io.BytesIO(data)
|
|
waveform, sample_rate = torchaudio.load(buffer)
|
|
|
|
target_rate = 44100
|
|
if sample_rate != target_rate:
|
|
resampler = torchaudio.transforms.Resample(sample_rate, target_rate)
|
|
waveform = resampler(waveform)
|
|
sample_rate = target_rate
|
|
|
|
return {
|
|
"waveform": waveform.unsqueeze(0),
|
|
"sample_rate": sample_rate
|
|
}
|
|
|
|
def _looks_like_image_url(self, url: str) -> bool:
|
|
"""Check if URL looks like an image based on extension."""
|
|
lower = url.lower()
|
|
return any(ext in lower for ext in ['.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp'])
|
|
|
|
def _looks_like_audio_url(self, url: str) -> bool:
|
|
"""Check if URL looks like audio based on extension."""
|
|
lower = url.lower()
|
|
return any(ext in lower for ext in ['.mp3', '.wav', '.flac', '.ogg', '.opus', '.m4a'])
|
|
|
|
# Specialized extraction methods for nodes
|
|
|
|
async def process_image(self, value: Any) -> torch.Tensor:
|
|
"""
|
|
Process a single image input.
|
|
|
|
Args:
|
|
value: Base64 string, URL, or already a tensor
|
|
|
|
Returns:
|
|
Tensor [1, H, W, C]
|
|
"""
|
|
if isinstance(value, torch.Tensor):
|
|
return value
|
|
if isinstance(value, str):
|
|
if self._is_base64_image(value):
|
|
return await self._decode_base64_image(value)
|
|
elif self._is_url(value):
|
|
return await self._download_and_decode(value)
|
|
# Return empty tensor for invalid input
|
|
return torch.zeros(1, 64, 64, 3)
|
|
|
|
async def process_audio(self, value: Any) -> dict:
|
|
"""
|
|
Process a single audio input.
|
|
|
|
Args:
|
|
value: Base64 string, URL, or already an audio dict
|
|
|
|
Returns:
|
|
{"waveform": tensor, "sample_rate": int}
|
|
"""
|
|
if isinstance(value, dict) and "waveform" in value:
|
|
return value
|
|
if isinstance(value, str):
|
|
if self._is_base64_audio(value):
|
|
return await self._decode_base64_audio(value)
|
|
elif self._is_url(value):
|
|
return await self._download_and_decode(value)
|
|
# Return empty audio for invalid input
|
|
return {"waveform": torch.zeros(1, 1, 44100), "sample_rate": 44100}
|
|
|
|
async def process_mask(self, value: Any) -> torch.Tensor:
|
|
"""
|
|
Process a mask input (grayscale image).
|
|
|
|
Args:
|
|
value: Base64 string or URL of grayscale image
|
|
|
|
Returns:
|
|
Tensor [1, H, W] for MASK type
|
|
"""
|
|
if isinstance(value, torch.Tensor):
|
|
if len(value.shape) == 3:
|
|
return value
|
|
elif len(value.shape) == 4:
|
|
# Convert [B, H, W, C] to [B, H, W] using first channel
|
|
return value[:, :, :, 0]
|
|
|
|
if isinstance(value, str):
|
|
# Decode image
|
|
if self._is_base64_image(value):
|
|
tensor = await self._decode_base64_image(value)
|
|
elif self._is_url(value):
|
|
tensor = await self._download_and_decode(value)
|
|
else:
|
|
return torch.zeros(1, 64, 64)
|
|
|
|
# Convert to grayscale mask
|
|
if len(tensor.shape) == 4:
|
|
# Average RGB channels
|
|
mask = tensor.mean(dim=-1)
|
|
else:
|
|
mask = tensor
|
|
return mask
|
|
|
|
return torch.zeros(1, 64, 64)
|
|
|
|
async def process_video(self, value: Any) -> Any:
|
|
"""
|
|
Process video input.
|
|
|
|
Note: Video processing is more complex and may require
|
|
additional dependencies like decord or av.
|
|
|
|
Args:
|
|
value: Base64 string or URL
|
|
|
|
Returns:
|
|
Video data (format depends on ComfyUI video nodes)
|
|
"""
|
|
# Basic implementation - return None for now
|
|
# Full video support would require decord or av
|
|
logger.warning("Video input processing not fully implemented")
|
|
return None
|
|
|
|
async def process_latent(self, value: Any) -> dict:
|
|
"""
|
|
Process latent input (safetensors or torch file).
|
|
|
|
Args:
|
|
value: Base64 string (safetensors/pt) or URL
|
|
|
|
Returns:
|
|
{"samples": tensor, ...} latent dict
|
|
"""
|
|
if isinstance(value, dict) and "samples" in value:
|
|
return value
|
|
|
|
default_latent = {"samples": torch.zeros(1, 4, 64, 64)}
|
|
|
|
if isinstance(value, str):
|
|
try:
|
|
# Get data bytes
|
|
if value.startswith("data:"):
|
|
_, data_str = value.split(",", 1)
|
|
data = base64.b64decode(data_str)
|
|
elif self._is_url(value):
|
|
if aiohttp is None:
|
|
return default_latent
|
|
timeout = aiohttp.ClientTimeout(total=self.download_timeout)
|
|
async with aiohttp.ClientSession(timeout=timeout) as session:
|
|
async with session.get(value) as response:
|
|
response.raise_for_status()
|
|
data = await response.read()
|
|
else:
|
|
# Raw base64
|
|
data = base64.b64decode(value)
|
|
|
|
# Try safetensors first
|
|
try:
|
|
from safetensors.torch import load
|
|
buffer = io.BytesIO(data)
|
|
tensors = load(buffer.getvalue())
|
|
if "samples" in tensors:
|
|
return {"samples": tensors["samples"]}
|
|
# Return first tensor as samples
|
|
for key, tensor in tensors.items():
|
|
return {"samples": tensor}
|
|
except ImportError:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
# Try torch load
|
|
try:
|
|
buffer = io.BytesIO(data)
|
|
loaded = torch.load(buffer, weights_only=True)
|
|
if isinstance(loaded, dict) and "samples" in loaded:
|
|
return loaded
|
|
if isinstance(loaded, torch.Tensor):
|
|
return {"samples": loaded}
|
|
except Exception:
|
|
pass
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to process latent: {e}")
|
|
|
|
return default_latent
|
|
|
|
async def process_mesh(self, value: Any) -> dict:
|
|
"""
|
|
Process 3D mesh input (GLB, OBJ, STL).
|
|
|
|
Args:
|
|
value: Base64 string or URL
|
|
|
|
Returns:
|
|
{"vertices": tensor, "faces": tensor} mesh dict
|
|
"""
|
|
default_mesh = {
|
|
"vertices": torch.zeros(1, 3, 3),
|
|
"faces": torch.zeros(1, 1, 3, dtype=torch.long)
|
|
}
|
|
|
|
if isinstance(value, dict) and "vertices" in value:
|
|
return value
|
|
|
|
if isinstance(value, str):
|
|
try:
|
|
# Get data bytes
|
|
if value.startswith("data:"):
|
|
_, data_str = value.split(",", 1)
|
|
data = base64.b64decode(data_str)
|
|
elif self._is_url(value):
|
|
if aiohttp is None:
|
|
return default_mesh
|
|
timeout = aiohttp.ClientTimeout(total=self.download_timeout)
|
|
async with aiohttp.ClientSession(timeout=timeout) as session:
|
|
async with session.get(value) as response:
|
|
response.raise_for_status()
|
|
data = await response.read()
|
|
else:
|
|
# Raw base64
|
|
data = base64.b64decode(value)
|
|
|
|
# Detect format and parse
|
|
return self._parse_mesh_bytes(data, value)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to process mesh: {e}")
|
|
|
|
return default_mesh
|
|
|
|
def _parse_mesh_bytes(self, data: bytes, hint: str = "") -> dict:
|
|
"""Parse mesh from bytes based on format detection."""
|
|
default_mesh = {
|
|
"vertices": torch.zeros(1, 3, 3),
|
|
"faces": torch.zeros(1, 1, 3, dtype=torch.long)
|
|
}
|
|
|
|
# Detect format
|
|
if data.startswith(b'glTF'):
|
|
return self._parse_glb(data)
|
|
elif data.startswith(b'ply\n') or data.startswith(b'ply\r\n'):
|
|
return self._parse_ply(data)
|
|
elif data.startswith(b'solid ') or data[:80].count(b'\x00') < 10:
|
|
# Text STL or ASCII format
|
|
if b'v ' in data[:1000]:
|
|
return self._parse_obj(data)
|
|
elif len(data) > 84:
|
|
# Binary STL check
|
|
try:
|
|
import struct
|
|
num_triangles = struct.unpack('<I', data[80:84])[0]
|
|
expected_size = 84 + num_triangles * 50
|
|
if len(data) == expected_size:
|
|
return self._parse_stl_binary(data)
|
|
except:
|
|
pass
|
|
|
|
# Check hint for PLY extension
|
|
if hint.lower().endswith('.ply'):
|
|
return self._parse_ply(data)
|
|
|
|
# Try OBJ as fallback for text formats
|
|
if b'v ' in data[:1000]:
|
|
return self._parse_obj(data)
|
|
|
|
return default_mesh
|
|
|
|
def _parse_glb(self, data: bytes) -> dict:
|
|
"""Parse GLB (binary glTF) file."""
|
|
import struct
|
|
import json as json_mod
|
|
|
|
default_mesh = {
|
|
"vertices": torch.zeros(1, 3, 3),
|
|
"faces": torch.zeros(1, 1, 3, dtype=torch.long)
|
|
}
|
|
|
|
try:
|
|
# GLB header: magic (4), version (4), length (4)
|
|
if len(data) < 12:
|
|
return default_mesh
|
|
|
|
magic, version, length = struct.unpack('<4sII', data[:12])
|
|
if magic != b'glTF':
|
|
return default_mesh
|
|
|
|
# Read JSON chunk
|
|
chunk_length, chunk_type = struct.unpack('<II', data[12:20])
|
|
json_data = json_mod.loads(data[20:20+chunk_length])
|
|
|
|
# Read BIN chunk
|
|
bin_offset = 20 + chunk_length
|
|
if bin_offset + 8 > len(data):
|
|
return default_mesh
|
|
|
|
bin_length, bin_type = struct.unpack('<II', data[bin_offset:bin_offset+8])
|
|
bin_data = data[bin_offset+8:bin_offset+8+bin_length]
|
|
|
|
# Extract mesh from first primitive
|
|
meshes = json_data.get("meshes", [])
|
|
if not meshes:
|
|
return default_mesh
|
|
|
|
primitive = meshes[0].get("primitives", [{}])[0]
|
|
accessors = json_data.get("accessors", [])
|
|
buffer_views = json_data.get("bufferViews", [])
|
|
|
|
# Get vertices
|
|
pos_accessor_idx = primitive.get("attributes", {}).get("POSITION")
|
|
if pos_accessor_idx is None:
|
|
return default_mesh
|
|
|
|
pos_accessor = accessors[pos_accessor_idx]
|
|
pos_view = buffer_views[pos_accessor["bufferView"]]
|
|
pos_offset = pos_view.get("byteOffset", 0)
|
|
vertex_count = pos_accessor["count"]
|
|
|
|
vertices_flat = np.frombuffer(
|
|
bin_data[pos_offset:pos_offset + vertex_count * 12],
|
|
dtype=np.float32
|
|
)
|
|
vertices = vertices_flat.reshape(-1, 3)
|
|
|
|
# Get indices
|
|
indices_accessor_idx = primitive.get("indices")
|
|
if indices_accessor_idx is not None:
|
|
idx_accessor = accessors[indices_accessor_idx]
|
|
idx_view = buffer_views[idx_accessor["bufferView"]]
|
|
idx_offset = idx_view.get("byteOffset", 0)
|
|
idx_count = idx_accessor["count"]
|
|
|
|
# Handle different component types
|
|
comp_type = idx_accessor.get("componentType", 5125)
|
|
if comp_type == 5123: # UNSIGNED_SHORT
|
|
faces_flat = np.frombuffer(
|
|
bin_data[idx_offset:idx_offset + idx_count * 2],
|
|
dtype=np.uint16
|
|
)
|
|
else: # UNSIGNED_INT
|
|
faces_flat = np.frombuffer(
|
|
bin_data[idx_offset:idx_offset + idx_count * 4],
|
|
dtype=np.uint32
|
|
)
|
|
faces = faces_flat.reshape(-1, 3)
|
|
else:
|
|
# Generate faces from vertices
|
|
faces = np.arange(len(vertices)).reshape(-1, 3)
|
|
|
|
return {
|
|
"vertices": torch.from_numpy(vertices).unsqueeze(0).float(),
|
|
"faces": torch.from_numpy(faces.astype(np.int64)).unsqueeze(0)
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to parse GLB: {e}")
|
|
return default_mesh
|
|
|
|
def _parse_obj(self, data: bytes) -> dict:
|
|
"""Parse OBJ file."""
|
|
default_mesh = {
|
|
"vertices": torch.zeros(1, 3, 3),
|
|
"faces": torch.zeros(1, 1, 3, dtype=torch.long)
|
|
}
|
|
|
|
try:
|
|
text = data.decode('utf-8', errors='replace')
|
|
vertices = []
|
|
faces = []
|
|
|
|
for line in text.split('\n'):
|
|
line = line.strip()
|
|
if line.startswith('v '):
|
|
parts = line.split()[1:4]
|
|
vertices.append([float(p) for p in parts])
|
|
elif line.startswith('f '):
|
|
parts = line.split()[1:]
|
|
# Handle v/vt/vn format
|
|
face_indices = []
|
|
for p in parts[:3]: # Only first 3 vertices for triangles
|
|
idx = int(p.split('/')[0]) - 1 # OBJ is 1-indexed
|
|
face_indices.append(idx)
|
|
if len(face_indices) == 3:
|
|
faces.append(face_indices)
|
|
|
|
if vertices and faces:
|
|
return {
|
|
"vertices": torch.tensor(vertices).unsqueeze(0).float(),
|
|
"faces": torch.tensor(faces, dtype=torch.long).unsqueeze(0)
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to parse OBJ: {e}")
|
|
|
|
return default_mesh
|
|
|
|
def _parse_stl_binary(self, data: bytes) -> dict:
|
|
"""Parse binary STL file."""
|
|
import struct
|
|
|
|
default_mesh = {
|
|
"vertices": torch.zeros(1, 3, 3),
|
|
"faces": torch.zeros(1, 1, 3, dtype=torch.long)
|
|
}
|
|
|
|
try:
|
|
num_triangles = struct.unpack('<I', data[80:84])[0]
|
|
vertices = []
|
|
faces = []
|
|
|
|
offset = 84
|
|
for i in range(num_triangles):
|
|
# Skip normal (12 bytes)
|
|
offset += 12
|
|
# Read 3 vertices
|
|
for j in range(3):
|
|
v = struct.unpack('<fff', data[offset:offset+12])
|
|
vertices.append(list(v))
|
|
offset += 12
|
|
# Skip attribute byte count
|
|
offset += 2
|
|
|
|
# Create face indices
|
|
base_idx = i * 3
|
|
faces.append([base_idx, base_idx + 1, base_idx + 2])
|
|
|
|
if vertices and faces:
|
|
return {
|
|
"vertices": torch.tensor(vertices).unsqueeze(0).float(),
|
|
"faces": torch.tensor(faces, dtype=torch.long).unsqueeze(0)
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to parse STL: {e}")
|
|
|
|
return default_mesh
|
|
|
|
def _parse_ply(self, data: bytes) -> dict:
|
|
"""
|
|
Parse PLY file.
|
|
|
|
Detects whether it's a Gaussian Splatting file or a regular point cloud/mesh.
|
|
Returns appropriate dict structure.
|
|
"""
|
|
import struct
|
|
|
|
default_result = {
|
|
"vertices": torch.zeros(1, 3, 3),
|
|
"faces": torch.zeros(1, 1, 3, dtype=torch.long)
|
|
}
|
|
|
|
try:
|
|
# Find header end
|
|
header_end = data.find(b'end_header')
|
|
if header_end == -1:
|
|
return default_result
|
|
|
|
header = data[:header_end].decode('ascii', errors='ignore')
|
|
|
|
# Handle both Unix (\n) and Windows (\r\n) line endings
|
|
body_start = header_end + len(b'end_header')
|
|
if data[body_start:body_start + 2] == b'\r\n':
|
|
body_start += 2
|
|
elif data[body_start:body_start + 1] == b'\n':
|
|
body_start += 1
|
|
|
|
# Parse header
|
|
lines = header.split('\n')
|
|
vertex_count = 0
|
|
face_count = 0
|
|
properties = []
|
|
is_binary_le = False
|
|
is_binary_be = False
|
|
|
|
for line in lines:
|
|
line = line.strip()
|
|
if line.startswith('format binary_little_endian'):
|
|
is_binary_le = True
|
|
elif line.startswith('format binary_big_endian'):
|
|
is_binary_be = True
|
|
elif line.startswith('element vertex'):
|
|
vertex_count = int(line.split()[-1])
|
|
elif line.startswith('element face'):
|
|
face_count = int(line.split()[-1])
|
|
elif line.startswith('property'):
|
|
parts = line.split()
|
|
if len(parts) >= 3:
|
|
prop_type = parts[1]
|
|
prop_name = parts[2]
|
|
properties.append((prop_name, prop_type))
|
|
|
|
# Detect if this is a Gaussian Splatting file
|
|
prop_names = [p[0] for p in properties]
|
|
is_gaussian = 'opacity' in prop_names or any(p.startswith('f_dc_') or p.startswith('scale_') for p in prop_names)
|
|
|
|
if is_binary_le or is_binary_be:
|
|
return self._parse_ply_binary(
|
|
data[body_start:], vertex_count, face_count, properties,
|
|
is_gaussian, little_endian=is_binary_le
|
|
)
|
|
else:
|
|
return self._parse_ply_ascii(
|
|
data[body_start:].decode('ascii', errors='ignore'),
|
|
vertex_count, face_count, properties, is_gaussian
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to parse PLY: {e}")
|
|
return default_result
|
|
|
|
def _parse_ply_binary(
|
|
self,
|
|
body: bytes,
|
|
vertex_count: int,
|
|
face_count: int,
|
|
properties: list,
|
|
is_gaussian: bool,
|
|
little_endian: bool = True
|
|
) -> dict:
|
|
"""Parse binary PLY data."""
|
|
import struct
|
|
|
|
endian = '<' if little_endian else '>'
|
|
|
|
# Calculate vertex stride and build format string
|
|
type_map = {
|
|
'float': ('f', 4),
|
|
'float32': ('f', 4),
|
|
'double': ('d', 8),
|
|
'float64': ('d', 8),
|
|
'uchar': ('B', 1),
|
|
'uint8': ('B', 1),
|
|
'char': ('b', 1),
|
|
'int8': ('b', 1),
|
|
'ushort': ('H', 2),
|
|
'uint16': ('H', 2),
|
|
'short': ('h', 2),
|
|
'int16': ('h', 2),
|
|
'uint': ('I', 4),
|
|
'uint32': ('I', 4),
|
|
'int': ('i', 4),
|
|
'int32': ('i', 4),
|
|
}
|
|
|
|
vertex_format = endian
|
|
vertex_stride = 0
|
|
prop_offsets = {}
|
|
|
|
for prop_name, prop_type in properties:
|
|
if prop_type in type_map:
|
|
fmt_char, size = type_map[prop_type]
|
|
prop_offsets[prop_name] = (vertex_stride, fmt_char, size)
|
|
vertex_format += fmt_char
|
|
vertex_stride += size
|
|
|
|
# Validate vertex count against actual data size (security check)
|
|
if vertex_stride > 0:
|
|
max_possible_vertices = len(body) // vertex_stride
|
|
if vertex_count > max_possible_vertices:
|
|
logger.warning(
|
|
f"PLY claims {vertex_count} vertices but data only supports {max_possible_vertices}"
|
|
)
|
|
vertex_count = max_possible_vertices
|
|
|
|
# Read all vertex data
|
|
all_data = []
|
|
offset = 0
|
|
for i in range(vertex_count):
|
|
if offset + vertex_stride > len(body):
|
|
break
|
|
vertex_data = struct.unpack(vertex_format, body[offset:offset + vertex_stride])
|
|
all_data.append(vertex_data)
|
|
offset += vertex_stride
|
|
|
|
# Extract positions (x, y, z)
|
|
positions = []
|
|
x_idx = [i for i, (n, _) in enumerate(properties) if n == 'x']
|
|
y_idx = [i for i, (n, _) in enumerate(properties) if n == 'y']
|
|
z_idx = [i for i, (n, _) in enumerate(properties) if n == 'z']
|
|
|
|
if x_idx and y_idx and z_idx:
|
|
x_idx, y_idx, z_idx = x_idx[0], y_idx[0], z_idx[0]
|
|
for vertex in all_data:
|
|
positions.append([vertex[x_idx], vertex[y_idx], vertex[z_idx]])
|
|
|
|
positions = np.array(positions, dtype=np.float32)
|
|
|
|
if is_gaussian:
|
|
# Build Gaussian Splatting result
|
|
result = {
|
|
"positions": torch.from_numpy(positions).unsqueeze(0)
|
|
}
|
|
|
|
# Extract opacity
|
|
opacity_idx = [i for i, (n, _) in enumerate(properties) if n == 'opacity']
|
|
if opacity_idx:
|
|
opacity = np.array([v[opacity_idx[0]] for v in all_data], dtype=np.float32)
|
|
result["opacity"] = torch.from_numpy(opacity).unsqueeze(0)
|
|
|
|
# Extract scales
|
|
scale_indices = [(i, n) for i, (n, _) in enumerate(properties) if n.startswith('scale_')]
|
|
if scale_indices:
|
|
scale_indices.sort(key=lambda x: x[1])
|
|
scales = np.array([[v[idx] for idx, _ in scale_indices] for v in all_data], dtype=np.float32)
|
|
result["scales"] = torch.from_numpy(scales).unsqueeze(0)
|
|
|
|
# Extract rotations
|
|
rot_indices = [(i, n) for i, (n, _) in enumerate(properties) if n.startswith('rot_')]
|
|
if rot_indices:
|
|
rot_indices.sort(key=lambda x: x[1])
|
|
rotations = np.array([[v[idx] for idx, _ in rot_indices] for v in all_data], dtype=np.float32)
|
|
result["rotations"] = torch.from_numpy(rotations).unsqueeze(0)
|
|
|
|
# Extract SH coefficients
|
|
sh_indices = [(i, n) for i, (n, _) in enumerate(properties) if n.startswith('f_dc_') or n.startswith('f_rest_')]
|
|
if sh_indices:
|
|
sh_indices.sort(key=lambda x: (0 if 'f_dc_' in x[1] else 1, x[1]))
|
|
sh_data = np.array([[v[idx] for idx, _ in sh_indices] for v in all_data], dtype=np.float32)
|
|
result["sh_coefficients"] = torch.from_numpy(sh_data).unsqueeze(0)
|
|
|
|
return result
|
|
|
|
else:
|
|
# Regular point cloud or mesh
|
|
result = {
|
|
"points": torch.from_numpy(positions).unsqueeze(0)
|
|
}
|
|
|
|
# Extract colors
|
|
r_idx = [i for i, (n, _) in enumerate(properties) if n in ('red', 'r')]
|
|
g_idx = [i for i, (n, _) in enumerate(properties) if n in ('green', 'g')]
|
|
b_idx = [i for i, (n, _) in enumerate(properties) if n in ('blue', 'b')]
|
|
|
|
if r_idx and g_idx and b_idx:
|
|
colors = np.array([
|
|
[v[r_idx[0]], v[g_idx[0]], v[b_idx[0]]]
|
|
for v in all_data
|
|
], dtype=np.float32)
|
|
# Normalize if values are > 1
|
|
if colors.max() > 1.0:
|
|
colors = colors / 255.0
|
|
result["colors"] = torch.from_numpy(colors).unsqueeze(0)
|
|
|
|
# Extract normals
|
|
nx_idx = [i for i, (n, _) in enumerate(properties) if n == 'nx']
|
|
ny_idx = [i for i, (n, _) in enumerate(properties) if n == 'ny']
|
|
nz_idx = [i for i, (n, _) in enumerate(properties) if n == 'nz']
|
|
|
|
if nx_idx and ny_idx and nz_idx:
|
|
normals = np.array([
|
|
[v[nx_idx[0]], v[ny_idx[0]], v[nz_idx[0]]]
|
|
for v in all_data
|
|
], dtype=np.float32)
|
|
result["normals"] = torch.from_numpy(normals).unsqueeze(0)
|
|
|
|
# For mesh compatibility
|
|
if face_count == 0:
|
|
result["vertices"] = result["points"]
|
|
result["faces"] = torch.zeros(1, 0, 3, dtype=torch.long)
|
|
|
|
return result
|
|
|
|
def _parse_ply_ascii(
|
|
self,
|
|
body: str,
|
|
vertex_count: int,
|
|
face_count: int,
|
|
properties: list,
|
|
is_gaussian: bool
|
|
) -> dict:
|
|
"""Parse ASCII PLY data."""
|
|
lines = body.strip().split('\n')
|
|
prop_names = [p[0] for p in properties]
|
|
|
|
vertices = []
|
|
for i in range(min(vertex_count, len(lines))):
|
|
values = lines[i].split()
|
|
vertex = {}
|
|
for j, val in enumerate(values):
|
|
if j < len(prop_names):
|
|
try:
|
|
vertex[prop_names[j]] = float(val)
|
|
except ValueError:
|
|
vertex[prop_names[j]] = 0.0
|
|
vertices.append(vertex)
|
|
|
|
# Extract positions
|
|
positions = np.array([
|
|
[v.get('x', 0), v.get('y', 0), v.get('z', 0)]
|
|
for v in vertices
|
|
], dtype=np.float32)
|
|
|
|
if is_gaussian:
|
|
result = {"positions": torch.from_numpy(positions).unsqueeze(0)}
|
|
|
|
if 'opacity' in prop_names:
|
|
opacity = np.array([v.get('opacity', 1.0) for v in vertices], dtype=np.float32)
|
|
result["opacity"] = torch.from_numpy(opacity).unsqueeze(0)
|
|
|
|
return result
|
|
else:
|
|
return {
|
|
"points": torch.from_numpy(positions).unsqueeze(0),
|
|
"vertices": torch.from_numpy(positions).unsqueeze(0),
|
|
"faces": torch.zeros(1, 0, 3, dtype=torch.long)
|
|
}
|
|
|
|
def _parse_exr(self, data: bytes) -> torch.Tensor:
|
|
"""Parse EXR depth map file."""
|
|
try:
|
|
import OpenEXR
|
|
import Imath
|
|
import tempfile
|
|
import os
|
|
|
|
# OpenEXR requires a file path
|
|
with tempfile.NamedTemporaryFile(suffix=".exr", delete=False) as f:
|
|
f.write(data)
|
|
temp_path = f.name
|
|
|
|
try:
|
|
exr_file = OpenEXR.InputFile(temp_path)
|
|
header = exr_file.header()
|
|
dw = header['dataWindow']
|
|
width = dw.max.x - dw.min.x + 1
|
|
height = dw.max.y - dw.min.y + 1
|
|
|
|
# Read the R channel (depth is usually in R or Y)
|
|
pt = Imath.PixelType(Imath.PixelType.FLOAT)
|
|
channels = exr_file.header()['channels'].keys()
|
|
|
|
if 'R' in channels:
|
|
channel_data = exr_file.channel('R', pt)
|
|
elif 'Y' in channels:
|
|
channel_data = exr_file.channel('Y', pt)
|
|
else:
|
|
channel_data = exr_file.channel(list(channels)[0], pt)
|
|
|
|
depth = np.frombuffer(channel_data, dtype=np.float32).reshape(height, width)
|
|
return torch.from_numpy(depth).unsqueeze(0)
|
|
|
|
finally:
|
|
os.unlink(temp_path)
|
|
|
|
except ImportError:
|
|
logger.warning("OpenEXR not available for EXR parsing")
|
|
return torch.zeros(1, 64, 64)
|
|
except Exception as e:
|
|
logger.error(f"Failed to parse EXR: {e}")
|
|
return torch.zeros(1, 64, 64)
|