nodes/kijai-ComfyUI-MoGe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

reqs

Browse Source
This commit is contained in:
kijai
2024-11-02 17:20:15 +02:00
parent 02df2ca14c
commit f6bb73bdcd
4 changed files with 21 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
moge/utils/vis.py
+1 -1
View File
@@ -1,5 +1,5 @@
import numpy as np
import matplotlib
#import matplotlib
def colorize_depth(depth: np.ndarray, mask: np.ndarray = None, normalize: bool = True, cmap: str = 'Spectral') -> np.ndarray:
nodes.py
+13 -19
View File
@@ -1,23 +1,13 @@
import os
from pathlib import Path
import torch
import folder_paths
import comfy.model_management as mm
from comfy.utils import ProgressBar, load_torch_file
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
log = logging.getLogger(__name__)
import trimesh
import numpy as np
from PIL import Image
from .moge.model import MoGeModel
from .utils3d.numpy import image_mesh, image_uv, depth_edge
import trimesh
import numpy as np
from pathlib import Path
import uuid
import tempfile
from PIL import Image
from contextlib import nullcontext
try:
from accelerate import init_empty_weights
@@ -27,8 +17,16 @@ except:
is_accelerate_available = False
pass
import comfy.model_management as mm
from comfy.utils import load_torch_file
import folder_paths
script_directory = os.path.dirname(os.path.abspath(__file__))
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
log = logging.getLogger(__name__)
#region ModelLoading
class DownloadAndLoadMoGeModel:
@@ -159,10 +157,6 @@ class MoGeProcess:
tri=True
)
vertices, vertex_uvs = vertices * [1, -1, -1], vertex_uvs * [1, -1] + [0, 1]
run_id = str(uuid.uuid4())
tempdir = folder_paths.get_temp_directory()
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, folder_paths.get_output_directory())
@@ -183,7 +177,7 @@ class MoGeProcess:
process=False
).export(output_glb_path)
elif output_format == 'ply':
output_ply_path = Path(tempdir, f'{run_id}.ply')
output_ply_path = Path(full_output_folder, f'{filename}_{counter:05}_.ply')
output_ply_path.parent.mkdir(exist_ok=True)
trimesh.Trimesh(
vertices=vertices,
requirements.txt
+6
View File
@@ -0,0 +1,6 @@
trimesh
pillow
scipy
numpy
huggingface_hub
opencv-python
utils3d/numpy/quadmesh.py
+1 -59
View File
@@ -410,63 +410,5 @@ def tri_to_quad(
vertices (np.ndarray): [N_, 3] 3-dimensional vertices
faces (np.ndarray): [Q, 4] quad face indices
"""
raise NotImplementedError
if __name__ == '__main__':
import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')))
import utils3d
import numpy as np
import cv2
from vis import vis_edge_color
file = 'miku'
vertices, faces = utils3d.io.read_ply(f'test/assets/{file}.ply')
edges, edge2face, face2edge, face2face = calc_relations(faces)
quad_cands, quad2edge, quad2adj, quad_valid = calc_quad_candidates(edges, face2edge, edge2face)
distortion = calc_quad_distortion(vertices, quad_cands)
direction = calc_quad_direction(vertices, quad_cands)
smoothness = calc_quad_smoothness(quad2edge, quad2adj, direction)
boundary_edges = edges[edge2face[:, 1] == -1]
quads_weight, conn_min_weight, conn_max_weight = sovle_quad(face2edge, edge2face, quad2adj, distortion, smoothness, quad_valid)
quads = quad_cands[quads_weight > 0.5]
print('Mesh statistics')
print(f' #V = {vertices.shape[0]}')
print(f' #F = {faces.shape[0]}')
print(f' #E = {edges.shape[0]}')
print(f' #B = {boundary_edges.shape[0]}')
print(f' #Q_cand = {quad_cands.shape[0]}')
print(f' #Q = {quads.shape[0]}')
utils3d.io.write_ply(f'test/assets/{file}_boundary_edges.ply', vertices=vertices, edges=boundary_edges)
utils3d.io.write_ply(f'test/assets/{file}_quad_candidates.ply', vertices=vertices, faces=quads)
edge_colors = np.zeros([edges.shape[0], 3], dtype=np.uint8)
distortion = (distortion - distortion.min()) / (distortion.max() - distortion.min())
distortion = (distortion * 255).astype(np.uint8)
edge_colors[quad_valid] = cv2.cvtColor(cv2.applyColorMap(distortion, cv2.COLORMAP_JET), cv2.COLOR_BGR2RGB).reshape(-1, 3)
utils3d.io.write_ply(f'test/assets/{file}_quad_candidates_distortion.ply', **vis_edge_color(vertices, edges, edge_colors))
edge_colors = np.zeros([edges.shape[0], 3], dtype=np.uint8)
edge_colors[quad_valid] = cv2.cvtColor(cv2.applyColorMap((quads_weight * 255).astype(np.uint8), cv2.COLORMAP_JET), cv2.COLOR_BGR2RGB).reshape(-1, 3)
utils3d.io.write_ply(f'test/assets/{file}_quad_candidates_weights.ply', **vis_edge_color(vertices, edges, edge_colors))
utils3d.io.write_ply(f'test/assets/{file}_quad.ply', vertices=vertices, faces=quads)
quad_centers = vertices[quad_cands].mean(axis=1)
conns = np.stack([
np.arange(quad_cands.shape[0])[:, None].repeat(8, axis=1),
quad2adj,
], axis=-1)[quad2adj != -1] # [C, 2]
conns, conns_idx = np.unique(np.sort(conns, axis=-1), axis=0, return_index=True) # [C, 2], [C]
smoothness = smoothness[quad2adj != -1][conns_idx] # [C]
conns_color = cv2.cvtColor(cv2.applyColorMap((smoothness * 255).astype(np.uint8), cv2.COLORMAP_JET), cv2.COLOR_BGR2RGB).reshape(-1, 3)
utils3d.io.write_ply(f'test/assets/{file}_quad_conn_smoothness.ply', **vis_edge_color(quad_centers, conns, conns_color))
conns_color = cv2.cvtColor(cv2.applyColorMap((conn_min_weight * 255).astype(np.uint8), cv2.COLORMAP_JET), cv2.COLOR_BGR2RGB).reshape(-1, 3)
utils3d.io.write_ply(f'test/assets/{file}_quad_conn_min.ply', **vis_edge_color(quad_centers, conns, conns_color))
conns_color = cv2.cvtColor(cv2.applyColorMap((conn_max_weight * 255).astype(np.uint8), cv2.COLORMAP_JET), cv2.COLOR_BGR2RGB).reshape(-1, 3)
utils3d.io.write_ply(f'test/assets/{file}_quad_conn_max.ply', **vis_edge_color(quad_centers, conns, conns_color))
raise NotImplementedError