245 lines
8.9 KiB
Python
245 lines
8.9 KiB
Python
"""Reframe360 — ComfyUI node definitions."""
|
|
|
|
import io
|
|
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from .projection import (
|
|
ASPECT_PRESETS,
|
|
LENS_PROJECTIONS,
|
|
SENSOR_PRESETS,
|
|
fov_degrees,
|
|
reframe,
|
|
resolve_sensor,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Panorama cache.
|
|
#
|
|
# The preview needs the pano as a texture. Pushing it through the websocket as
|
|
# base64 on every run is slow and forgets itself on page reload. Caching it
|
|
# here and serving it over HTTP means the preview loads on its own and
|
|
# survives a refresh.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_PANO_CACHE = {}
|
|
_CACHE_LIMIT = 12
|
|
|
|
try:
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
|
|
@PromptServer.instance.routes.get("/reframe360/pano/{node_id}")
|
|
async def _serve_pano(request):
|
|
data = _PANO_CACHE.get(request.match_info["node_id"])
|
|
if data is None:
|
|
return web.Response(status=404)
|
|
return web.Response(
|
|
body=data,
|
|
content_type="image/jpeg",
|
|
headers={"Cache-Control": "no-store"},
|
|
)
|
|
|
|
except Exception as exc: # pragma: no cover
|
|
print(f"[Reframe360] preview route unavailable: {exc}")
|
|
|
|
|
|
def _cache_pano(node_id, equi, max_width):
|
|
"""Downscale the pano and stash it for the preview to fetch."""
|
|
img = equi[0].permute(1, 2, 0).cpu().numpy()
|
|
img = (img.clip(0, 1) * 255).astype("uint8")
|
|
pil = Image.fromarray(img)
|
|
|
|
if max_width > 0 and pil.width > max_width:
|
|
h = max(1, int(pil.height * max_width / pil.width))
|
|
pil = pil.resize((max_width, h), Image.LANCZOS)
|
|
|
|
buf = io.BytesIO()
|
|
pil.save(buf, format="JPEG", quality=88)
|
|
|
|
if len(_PANO_CACHE) > _CACHE_LIMIT:
|
|
_PANO_CACHE.pop(next(iter(_PANO_CACHE)))
|
|
_PANO_CACHE[str(node_id)] = buf.getvalue()
|
|
|
|
|
|
class Reframe360:
|
|
"""
|
|
Point a virtual camera around inside a 360 image.
|
|
|
|
Drag the preview to frame, then read exact numbers off the sliders. Both
|
|
write the same values, so rough framing and repeatable output are the
|
|
same operation.
|
|
"""
|
|
|
|
CATEGORY = "Reframe360"
|
|
FUNCTION = "execute"
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "camera_info")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# Order matters: the JS groups widgets by contiguous runs, so keep
|
|
# lens / camera / output params together and in this sequence.
|
|
return {
|
|
"required": {
|
|
"equirectangular": ("IMAGE",),
|
|
|
|
# --- LENS ---
|
|
"sensor": (list(SENSOR_PRESETS.keys()), {"default": "Full Frame (36 x 24)"}),
|
|
"sensor_fit": (["auto", "horizontal", "vertical"], {"default": "auto"}),
|
|
"sensor_width_custom": ("FLOAT", {"default": 36.0, "min": 0.1, "max": 200.0, "step": 0.01}),
|
|
"sensor_height_custom": ("FLOAT", {"default": 24.0, "min": 0.1, "max": 200.0, "step": 0.01}),
|
|
"focal_length": ("FLOAT", {
|
|
"default": 18.0, "min": 1.0, "max": 300.0, "step": 0.1,
|
|
"tooltip": "Millimetres. Meaningless without the sensor size above.",
|
|
}),
|
|
"lens_projection": (LENS_PROJECTIONS, {
|
|
"default": "rectilinear",
|
|
"tooltip": "Rectilinear keeps straight lines straight but falls apart below "
|
|
"~12mm on full frame. Use a fisheye projection for ultra-wide.",
|
|
}),
|
|
|
|
# --- CAMERA ---
|
|
"yaw": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1}),
|
|
"pitch": ("FLOAT", {"default": 0.0, "min": -90.0, "max": 90.0, "step": 0.1}),
|
|
"roll": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1}),
|
|
"perspective_warp": ("FLOAT", {
|
|
"default": 0.0, "min": -0.95, "max": 0.95, "step": 0.001,
|
|
"tooltip": "Slides the camera along its OWN Z axis, as a fraction of the "
|
|
"sphere radius. Bends straight lines, gives no parallax. "
|
|
"A warp, not a dolly.",
|
|
}),
|
|
|
|
# --- OUTPUT ---
|
|
"aspect_preset": (list(ASPECT_PRESETS.keys()), {"default": "16:9"}),
|
|
"output_width": ("INT", {"default": 1920, "min": 16, "max": 16384, "step": 2}),
|
|
"output_height": ("INT", {"default": 1080, "min": 16, "max": 16384, "step": 2}),
|
|
"interpolation": (["bicubic", "bilinear", "nearest"], {"default": "bicubic"}),
|
|
},
|
|
"optional": {
|
|
"preview_max_width": ("INT", {"default": 2048, "min": 256, "max": 8192, "step": 256}),
|
|
},
|
|
"hidden": {"unique_id": "UNIQUE_ID"},
|
|
}
|
|
|
|
def execute(
|
|
self, equirectangular, sensor, sensor_fit, sensor_width_custom,
|
|
sensor_height_custom, focal_length, lens_projection,
|
|
yaw, pitch, roll, perspective_warp,
|
|
aspect_preset, output_width, output_height, interpolation,
|
|
preview_max_width=2048, unique_id=None,
|
|
):
|
|
if sensor == "Custom":
|
|
sw, sh = sensor_width_custom, sensor_height_custom
|
|
else:
|
|
sw, sh = SENSOR_PRESETS[sensor]
|
|
|
|
# An aspect preset drives height from width. Custom leaves both alone.
|
|
target = ASPECT_PRESETS.get(aspect_preset)
|
|
if target is not None:
|
|
output_height = max(2, int(round(output_width / target / 2) * 2))
|
|
|
|
# ComfyUI hands us (B, H, W, C); the math wants (B, C, H, W).
|
|
equi = equirectangular.permute(0, 3, 1, 2).contiguous()
|
|
|
|
ratio = equi.shape[3] / equi.shape[2]
|
|
if abs(ratio - 2.0) > 0.02:
|
|
print(
|
|
f"[Reframe360] Input is {equi.shape[3]}x{equi.shape[2]} (ratio {ratio:.2f}:1). "
|
|
f"Equirectangular images are 2:1. Output will be distorted."
|
|
)
|
|
|
|
if unique_id is not None:
|
|
try:
|
|
_cache_pano(unique_id, equi, preview_max_width)
|
|
except Exception as exc:
|
|
print(f"[Reframe360] could not cache preview: {exc}")
|
|
|
|
out = reframe(
|
|
equi,
|
|
out_w=output_width, out_h=output_height,
|
|
sensor_w=sw, sensor_h=sh, sensor_fit=sensor_fit,
|
|
focal=focal_length, lens=lens_projection,
|
|
yaw=yaw, pitch=pitch, roll=roll,
|
|
warp=perspective_warp,
|
|
interpolation=interpolation,
|
|
)
|
|
|
|
eff_w, eff_h = resolve_sensor(sw, sh, sensor_fit, output_width, output_height)
|
|
hfov, vfov = fov_degrees(eff_w, eff_h, focal_length, lens_projection)
|
|
|
|
info = (
|
|
f"{focal_length:.1f}mm | {sensor} | {lens_projection}\n"
|
|
f"FOV {hfov:.1f} x {vfov:.1f} deg | effective sensor {eff_w:.2f} x {eff_h:.2f}mm\n"
|
|
f"yaw {yaw:.2f} pitch {pitch:.2f} roll {roll:.2f} | warp {perspective_warp:.3f}\n"
|
|
f"{output_width}x{output_height} ({output_width/output_height:.3f}:1) | "
|
|
f"{equi.shape[0]} frame(s)"
|
|
)
|
|
|
|
if hfov > 150 and lens_projection == "rectilinear":
|
|
info += "\nWARNING: rectilinear past 150 deg. Corners are badly stretched."
|
|
|
|
result = out.permute(0, 2, 3, 1).contiguous()
|
|
|
|
return {
|
|
"ui": {"pano_ready": [str(unique_id)], "camera_info": [info]},
|
|
"result": (result, info),
|
|
}
|
|
|
|
|
|
class Load360Image:
|
|
"""
|
|
Fetch an equirectangular image from a URL.
|
|
|
|
Works with any direct link to a 2:1 image — Poly Haven HDRIs, Mapillary,
|
|
Wikimedia, your own server.
|
|
|
|
Deliberately ships no site-specific scraper. Google's Geo Guidelines forbid
|
|
downloading Street View imagery for use outside their services, so that
|
|
source is yours to arrange, not this node's to ship.
|
|
"""
|
|
|
|
CATEGORY = "Reframe360"
|
|
FUNCTION = "execute"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"url": ("STRING", {"default": "", "multiline": False})},
|
|
"optional": {"timeout": ("INT", {"default": 30, "min": 5, "max": 300})},
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, url, **kw):
|
|
return url
|
|
|
|
def execute(self, url, timeout=30):
|
|
import numpy as np
|
|
import requests
|
|
|
|
url = url.strip()
|
|
if not url:
|
|
raise ValueError("No URL. Paste a direct link to an equirectangular image.")
|
|
if not url.startswith(("http://", "https://")):
|
|
raise ValueError(f"Not a URL: {url!r}")
|
|
|
|
r = requests.get(url, timeout=timeout, headers={"User-Agent": "ComfyUI-Reframe360"})
|
|
r.raise_for_status()
|
|
|
|
pil = Image.open(io.BytesIO(r.content)).convert("RGB")
|
|
arr = np.array(pil).astype("float32") / 255.0
|
|
return (torch.from_numpy(arr).unsqueeze(0),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Reframe360": Reframe360,
|
|
"Load360Image": Load360Image,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Reframe360": "Reframe 360 Camera",
|
|
"Load360Image": "Load 360 Image (URL)",
|
|
}
|