diff --git a/nodes/debug.py b/nodes/debug.py index 05aa335..364ad75 100644 --- a/nodes/debug.py +++ b/nodes/debug.py @@ -1,10 +1,8 @@ -import base64 -import io +import base64, io from pathlib import Path from typing import Optional -import folder_paths -import torch +import folder_paths, torch from ..log import log from ..utils import tensor2pil @@ -43,18 +41,29 @@ def process_list(anything): ) elif isinstance(first_element, torch.Tensor): - text.append(f"List of Tensors: {first_element.shape} (x{len(anything)})") + text.append( + f"List of Tensors: {first_element.shape} (x{len(anything)})" + ) return {"text": text} def process_dict(anything): - text = [] + if "mesh" in anything: + m = {"geometry": {}} + m["geometry"]["mesh"] = mesh_to_json(anything["mesh"]) + if "material" in anything: + m["geometry"]["material"] = anything["material"] + return m + + res = [] if "samples" in anything: - is_empty = "(empty)" if torch.count_nonzero(anything["samples"]) == 0 else "" - text.append(f"Latent Samples: {anything['samples'].shape} {is_empty}") + is_empty = ( + "(empty)" if torch.count_nonzero(anything["samples"]) == 0 else "" + ) + res.append(f"Latent Samples: {anything['samples'].shape} {is_empty}") - return {"text": text} + return {"text": res} def process_bool(anything): @@ -65,6 +74,7 @@ def process_text(anything): return {"text": [str(anything)]} +# NOT USED ANYMORE def process_geometry(anything): return {"geometry": [mesh_to_json(anything)]} @@ -102,8 +112,6 @@ def do_debug(self, output_to_console, **kwargs): bool: process_bool, o3d.geometry.Geometry: process_geometry, } - if output_to_console: - print("bouh!") for anything in kwargs.values(): processor = processors.get(type(anything)) @@ -118,11 +126,21 @@ def do_debug(self, output_to_console, **kwargs): processed_data = processor(anything) for ui_key, ui_value in processed_data.items(): - output["ui"][ui_key].extend(ui_value) + if isinstance(ui_value, list): + output["ui"][ui_key].extend(ui_value) + else: + output["ui"][ui_key].append(ui_value) # log.debug( # f"Processed input {k}, found {len(processed_data.get('b64_images', []))} images and {len(processed_data.get('text', []))} text items." # ) + if output_to_console: + from rich.console import Console + + cons = Console() + cons.print("OUTPUT:") + cons.print(output) + return output diff --git a/nodes/geo_tools.py b/nodes/geo_tools.py index 0cad1fa..39054e4 100644 --- a/nodes/geo_tools.py +++ b/nodes/geo_tools.py @@ -1,10 +1,18 @@ +import copy import itertools +import json +import os + import numpy as np import open3d as o3d -import json + from ..utils import log -import os + +def spread_geo(geo, *, cp=False): + mesh = geo["mesh"] if not cp else copy.copy(geo["mesh"]) + material = geo.get("material", {}) + return (mesh, material) def euler_to_rotation_matrix(x_deg, y_deg, z_deg): @@ -14,13 +22,19 @@ def euler_to_rotation_matrix(x_deg, y_deg, z_deg): z = np.radians(z_deg) # Rotation matrix around x-axis - Rx = np.array([[1, 0, 0], [0, np.cos(x), -np.sin(x)], [0, np.sin(x), np.cos(x)]]) + Rx = np.array( + [[1, 0, 0], [0, np.cos(x), -np.sin(x)], [0, np.sin(x), np.cos(x)]] + ) # Rotation matrix around y-axis - Ry = np.array([[np.cos(y), 0, np.sin(y)], [0, 1, 0], [-np.sin(y), 0, np.cos(y)]]) + Ry = np.array( + [[np.cos(y), 0, np.sin(y)], [0, 1, 0], [-np.sin(y), 0, np.cos(y)]] + ) # Rotation matrix around z-axis - Rz = np.array([[np.cos(z), -np.sin(z), 0], [np.sin(z), np.cos(z), 0], [0, 0, 1]]) + Rz = np.array( + [[np.cos(z), -np.sin(z), 0], [np.sin(z), np.cos(z), 0], [0, 0, 1]] + ) return Rz @ Ry @ Rx @@ -75,7 +89,9 @@ def create_grid(scale=(1, 1, 1), rows=10, columns=10): # Create vertices vertices = [] for i in np.linspace(-dy / 2, dy / 2, rows + 1): - vertices.extend([j, 0, i] for j in np.linspace(-dx / 2, dx / 2, columns + 1)) + vertices.extend( + [j, 0, i] for j in np.linspace(-dx / 2, dx / 2, columns + 1) + ) # Generate triangles triangles = [] for i, j in itertools.product(range(rows), range(columns)): @@ -101,7 +117,9 @@ def create_box(scale=(1, 1, 1), divisions=(1, 1, 1)): vertices = [] for i in np.linspace(-dx / 2, dx / 2, div_x + 1): for j in np.linspace(-dy / 2, dy / 2, div_y + 1): - vertices.extend([i, j, k] for k in np.linspace(-dz / 2, dz / 2, div_z + 1)) + vertices.extend( + [i, j, k] for k in np.linspace(-dz / 2, dz / 2, div_z + 1) + ) # Generate triangles for the box faces triangles = [] for x, y in itertools.product(range(div_x), range(div_y)): @@ -240,20 +258,125 @@ def create_torus(torus_radius=1, ring_radius=0.5, rows=10, columns=10): # CATEGORY = "mtb/uv" +def default_material(color=None): + return { + "color": color or "0x00ff00", + "roughness": 1.0, + "metalness": 0.0, + "emissive": "0x000000", + "displacementScale": 1.0, + "displacementMap": None, + } + + +class MTB_Material: + """Make a std material.""" + + @classmethod + def INPUT_TYPES(cls): + base = default_material() + return { + "required": { + "color": ("COLOR", {"default": base["color"]}), + "roughness": ( + "FLOAT", + { + "default": base["roughness"], + "min": 0.005, + "max": 4.0, + "step": 0.01, + }, + ), + "metalness": ( + "FLOAT", + { + "default": base["metalness"], + "min": 0.0, + "max": 1.0, + "step": 0.01, + }, + ), + "emissive": ("COLOR", {"default": base["emissive"]}), + "displacementScale": ( + "FLOAT", + {"default": 1.0, "min": -10.0, "max": 10.0}, + ), + }, + "optional": {"displacementMap": ("IMAGE",)}, + } + + RETURN_TYPES = ("GEO_MATERIAL",) + RETURN_NAMES = ("material",) + FUNCTION = "make_material" + CATEGORY = "mtb/3D" + + def make_material( + self, **kwargs + ): # color, roughness, metalness, emissive, displacementScalen displacementMap=None): + return (kwargs,) + + +class MTB_ApplyMaterial: + """Apply a Material to a geometry.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": {"geometry": ("GEOMETRY",), "color": ("COLOR",)}, + "optional": {"material": ("GEO_MATERIAL",)}, + } + + RETURN_TYPES = ("GEOMETRY",) + RETURN_NAMES = ("geometry",) + FUNCTION = "apply" + CATEGORY = "mtb/3D" + + def apply( + self, + geometry, + color, + material=None, + ): + if material is None: + material = default_material(color) + # + geometry["material"] = material + + return (geometry,) + + class TransformGeometry: - """Transforms the input geometry""" + """Transforms the input geometry.""" @classmethod def INPUT_TYPES(cls): return { "required": { "mesh": ("GEOMETRY",), - "position_x": ("FLOAT", {"default": 0.0, "step": 0.1}), - "position_y": ("FLOAT", {"default": 0.0, "step": 0.1}), - "position_z": ("FLOAT", {"default": 0.0, "step": 0.1}), - "rotation_x": ("FLOAT", {"default": 0.0, "step": 1}), - "rotation_y": ("FLOAT", {"default": 0.0, "step": 1}), - "rotation_z": ("FLOAT", {"default": 0.0, "step": 1}), + "position_x": ( + "FLOAT", + {"default": 0.0, "step": 0.1, "min": -10000, "max": 10000}, + ), + "position_y": ( + "FLOAT", + {"default": 0.0, "step": 0.1, "min": -10000, "max": 10000}, + ), + "position_z": ( + "FLOAT", + {"default": 0.0, "step": 0.1, "min": -10000, "max": 10000}, + ), + "rotation_x": ( + "FLOAT", + {"default": 0.0, "step": 1, "min": -10000, "max": 10000}, + ), + "rotation_y": ( + "FLOAT", + {"default": 0.0, "step": 1, "min": -10000, "max": 10000}, + ), + "rotation_z": ( + "FLOAT", + {"default": 0.0, "step": 1, "min": -10000, "max": 10000}, + ), "scale_x": ("FLOAT", {"default": 1.0, "step": 0.1}), "scale_y": ("FLOAT", {"default": 1.0, "step": 0.1}), "scale_z": ("FLOAT", {"default": 1.0, "step": 0.1}), @@ -267,7 +390,7 @@ def INPUT_TYPES(cls): def transform_geometry( self, - mesh, + mesh: o3d.geometry.TriangleMesh, position_x=0.0, position_y=0.0, position_z=0.0, @@ -290,12 +413,21 @@ def transform_geometry( rotation = (rotation_x, rotation_y, rotation_z) scale = np.array([scale_x, scale_y, scale_z]) - transformation_matrix = get_transformation_matrix(position, rotation, scale) - return (mesh.transform(transformation_matrix),) + transformation_matrix = get_transformation_matrix( + position, rotation, scale + ) + mesh, material = spread_geo(mesh, cp=True) + + return ( + { + "mesh": mesh.transform(transformation_matrix), + "material": material, + }, + ) class GeometrySphere: - """Makes a Sphere 3D geometry""" + """Makes a Sphere 3D geometry..""" @classmethod def INPUT_TYPES(cls): @@ -320,11 +452,11 @@ def make_sphere(self, create_uv_map, radius, resolution): ) mesh.compute_vertex_normals() - return (mesh,) + return ({"mesh": mesh},) class GeometryTest: - """Fetches an Open3D data geometry""" + """Fetches an Open3D data geometry..""" @classmethod def INPUT_TYPES(cls): @@ -358,11 +490,11 @@ def fetch_data(self, name): model = getattr(o3d.data, name)() mesh = o3d.io.read_triangle_mesh(model.path) mesh.compute_vertex_normals() - return (mesh,) + return ({"mesh": mesh},) class GeometryBox: - """Makes a Box 3D geometry""" + """Makes a Box 3D geometry.""" @classmethod def INPUT_TYPES(cls): @@ -408,11 +540,11 @@ def make_box( (width, height, depth), (divisions_x, divisions_y, divisions_z) ) - return (mesh,) + return ({"mesh": mesh},) class LoadGeometry: - """Load a 3D geometry""" + """Load a 3D geometry.""" @classmethod def INPUT_TYPES(cls): @@ -428,35 +560,47 @@ def load_geo(self, path): raise ValueError(f"Path {path} does not exist") mesh = o3d.io.read_triangle_mesh(path) + + if len(mesh.vertices) == 0: + mesh = o3d.io.read_triangle_model(path) + mesh_count = len(mesh.meshes) + if mesh_count == 0: + raise ValueError("Couldn't parse input file") + + if mesh_count > 1: + log.warn( + f"Found {mesh_count} meshes, only the first will be used..." + ) + + mesh = mesh.meshes[0].mesh + mesh.compute_vertex_normals() return { - "result": (mesh,), + "result": ({"mesh": mesh},), } class GeometryInfo: - """Retrieve information about a 3D geometry""" + """Retrieve information about a 3D geometry.""" @classmethod def INPUT_TYPES(cls): return {"required": {"geometry": ("GEOMETRY", {})}} - RETURN_TYPES = ("INT", "INT") - RETURN_NAMES = ("num_vertices", "num_triangles") + RETURN_TYPES = ("INT", "INT", "MATERIAL") + RETURN_NAMES = ("num_vertices", "num_triangles", "material") FUNCTION = "get_info" CATEGORY = "mtb/3D" def get_info(self, geometry): - log.debug(geometry) - return ( - len(geometry.vertices), - len(geometry.triangles), - ) + mesh, material = spread_geo(geometry) + log.debug(mesh) + return (len(mesh.vertices), len(mesh.triangles), material) class GeometryDecimater: - """Optimized the geometry to match the target number of triangles""" + """Optimized the geometry to match the target number of triangles.""" @classmethod def INPUT_TYPES(cls): @@ -473,14 +617,16 @@ def INPUT_TYPES(cls): CATEGORY = "mtb/3D" def decimate(self, mesh, target): - mesh = mesh.simplify_quadric_decimation(target_number_of_triangles=target) + mesh = mesh.simplify_quadric_decimation( + target_number_of_triangles=target + ) mesh.compute_vertex_normals() - return (mesh,) + return ({"mesh": mesh},) class GeometrySceneSetup: - """Scene setup for the renderer""" + """Scene setup for the renderer.""" @classmethod def INPUT_TYPES(cls): @@ -496,11 +642,11 @@ def INPUT_TYPES(cls): CATEGORY = "mtb/3D" def setup(self, mesh, target): - return ({"geometry": mesh, "camera": cam},) + return ({"geometry": {"mesh": mesh}, "camera": cam},) class GeometryRender: - """Renders a Geometry to an image""" + """Renders a Geometry to an image.""" @classmethod def INPUT_TYPES(cls): @@ -536,6 +682,9 @@ def render(self, geometry, width, height, background, camera): GeometryTest, GeometryDecimater, GeometrySphere, + GeometryRender, TransformGeometry, GeometryBox, + MTB_ApplyMaterial, + MTB_Material, ] diff --git a/requirements.txt b/requirements.txt index 96d14b6..cc618a7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,5 @@ rembg imageio_ffmpeg rich rich_argparse -matplotlib \ No newline at end of file +matplotlib +open3d==0.17.0 \ No newline at end of file diff --git a/web/debug.js b/web/debug.js index eca8202..11fe901 100644 --- a/web/debug.js +++ b/web/debug.js @@ -10,9 +10,7 @@ import { app } from '../../scripts/app.js' import * as shared from './comfy_shared.js' -import { log } from './comfy_shared.js' import { MtbWidgets } from './mtb_widgets.js' -import { o3d_to_three } from './geometry_nodes.js' // TODO: respect inputs order... diff --git a/web/geometry_nodes.js b/web/geometry_nodes.js index 90d0491..64ceb42 100644 --- a/web/geometry_nodes.js +++ b/web/geometry_nodes.js @@ -27,7 +27,9 @@ export const make_wireframe = (mesh) => { return wireframe } -export const o3d_to_three = (data) => { +export const o3d_to_three = (data, material_opts) => { + + material_opts = material_opts || { color: "0x00ff00" } // Parse the JSON data const meshData = JSON.parse(data) @@ -58,13 +60,16 @@ export const o3d_to_three = (data) => { const uvs = new Float32Array(meshData.triangle_uvs.flat()) geometry.setAttribute('uv', new THREE.BufferAttribute(uvs, 2)) } - const material_opts = { - // wireframe: true, + + if (material_opts.displacementB64 != undefined) { + material_opts.displacementMap = THREE.ImageUtils.loadTexture(material_opts.displacementB64) + delete (material_opts.displacementB64) } + // For visualization, you might choose to use the MeshPhongMaterial to get the benefit of lighting with normals const material = meshData.vertex_colors - ? new THREE.MeshPhongMaterial({ ...material_opts, vertexColors: true }) - : new THREE.MeshPhongMaterial({ ...material_opts, color: 0x00ff00 }) + ? new THREE.MeshStandardMaterial({ ...material_opts, vertexColors: true }) + : new THREE.MeshStandardMaterial({ ...material_opts }) const threeMesh = new THREE.Mesh(geometry, material) diff --git a/web/mtb_widgets.js b/web/mtb_widgets.js index bccf9a1..c9b58a9 100644 --- a/web/mtb_widgets.js +++ b/web/mtb_widgets.js @@ -406,7 +406,7 @@ export const MtbWidgets = { const animate = () => { requestAnimationFrame(animate) if (this.mesh && this.animate) { - this.group.rotation.x += 0.005 + //this.group.rotation.x += 0.005 this.group.rotation.y += 0.005 } this.renderer.render(this.scene, this.camera) @@ -432,9 +432,16 @@ export const MtbWidgets = { } return [width, width] }, + onRemoved: function () { + if (this.inputEl) { + this.inputEl.remove() + } + } } log('Creating canvas') w.inputEl = document.createElement('canvas') + w.inputEl.width = 768 + w.inputEl.height = 768 // add context menu with "animate" and "show wireframe" w.inputEl.addEventListener('contextmenu', (e) => { @@ -483,9 +490,9 @@ export const MtbWidgets = { document.body.appendChild(w.menu) }) - w.initThreeJS(w.inputEl, val) + w.initThreeJS(w.inputEl) - w.mesh = o3d_to_three(val) + w.mesh = o3d_to_three(val?.mesh ? val.mesh : val, val?.material) w.mesh_wireframe = make_wireframe(w.mesh) w.group = new THREE.Group()