StdGEN / infer_refine.py
YulianSa's picture
update
914c133
raw
history blame
16 kB
import cv2
import numpy as np
import os
import trimesh
import argparse
import torch
import scipy
from PIL import Image
from refine.mesh_refine import geo_refine
from refine.func import make_star_cameras_orthographic
from refine.render import NormalsRenderer, calc_vertex_normals
from pytorch3d.structures import Meshes
from sklearn.neighbors import KDTree
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
# sam = sam_model_registry["vit_h"](checkpoint="./ckpt/sam_vit_h_4b8939.pth").cuda()
# generator = SamAutomaticMaskGenerator(
# model=sam,
# points_per_side=64,
# pred_iou_thresh=0.80,
# stability_score_thresh=0.92,
# crop_n_layers=1,
# crop_n_points_downscale_factor=2,
# min_mask_region_area=100,
# )
def fix_vert_color_glb(mesh_path):
from pygltflib import GLTF2, Material, PbrMetallicRoughness
obj1 = GLTF2().load(mesh_path)
obj1.meshes[0].primitives[0].material = 0
obj1.materials.append(Material(
pbrMetallicRoughness = PbrMetallicRoughness(
baseColorFactor = [1.0, 1.0, 1.0, 1.0],
metallicFactor = 0.,
roughnessFactor = 1.0,
),
emissiveFactor = [0.0, 0.0, 0.0],
doubleSided = True,
))
obj1.save(mesh_path)
def srgb_to_linear(c_srgb):
c_linear = np.where(c_srgb <= 0.04045, c_srgb / 12.92, ((c_srgb + 0.055) / 1.055) ** 2.4)
return c_linear.clip(0, 1.)
import trimesh
import numpy as np
from pytorch3d.structures import Meshes
from pytorch3d.renderer import TexturesUV
def save_py3dmesh_with_trimesh_fast(meshes: Meshes, save_glb_path, apply_sRGB_to_LinearRGB=True):
# Convert from pytorch3d meshes to trimesh mesh
vertices = meshes.verts_packed().cpu().float().numpy()
triangles = meshes.faces_packed().cpu().long().numpy()
# Check if the mesh uses TexturesUV
if isinstance(meshes.textures, TexturesUV):
# Extract UV coordinates and texture map
verts_uvs = meshes.textures.verts_uvs_padded()[0].cpu().numpy() # UV coordinates (N, 2)
faces_uvs = meshes.textures.faces_uvs_padded()[0].cpu().numpy() # UV face indices (M, 3)
texture_map = meshes.textures.maps_padded()[0].cpu().numpy() # Texture map (H, W, 3 or 4)
# Convert texture map to trimesh-compatible format
if apply_sRGB_to_LinearRGB:
texture_map = srgb_to_linear(texture_map)
texture_map = np.clip(texture_map, 0, 1) # Ensure values are in [0, 1]
material = trimesh.visual.texture.SimpleMaterial(image=texture_data, diffuse=(255, 255, 255))
# Create a trimesh.Trimesh object with UVs and texture
mesh = trimesh.Trimesh(
vertices=vertices,
faces=triangles,
visual=trimesh.visual.TextureVisuals(
uv=verts_uvs, # UV coordinates
image=texture_map, # Texture map
material=material # Material with texture
)
)
else:
# Fallback to vertex colors if TexturesUV is not used
np_color = meshes.textures.verts_features_packed().cpu().float().numpy()
if apply_sRGB_to_LinearRGB:
np_color = srgb_to_linear(np_color)
np_color = np.clip(np_color, 0, 1)
mesh = trimesh.Trimesh(vertices=vertices, faces=triangles, vertex_colors=np_color)
# Rotate 180 degrees along +Y if saving as GLB
if save_glb_path.endswith(".glb"):
vertices[:, [0, 2]] = -vertices[:, [0, 2]]
# Remove unreferenced vertices
mesh.remove_unreferenced_vertices()
# Save mesh
mesh.export(save_glb_path)
# if save_glb_path.endswith(".glb"):
# fix_vert_color_glb(save_glb_path)
print(f"Saving to {save_glb_path}")
def calc_horizontal_offset(target_img, source_img):
target_mask = target_img.astype(np.float32).sum(axis=-1) > 750
source_mask = source_img.astype(np.float32).sum(axis=-1) > 750
best_offset = -114514
for offset in range(-200, 200):
offset_mask = np.roll(source_mask, offset, axis=1)
overlap = (target_mask & offset_mask).sum()
if overlap > best_offset:
best_offset = overlap
best_offset_value = offset
return best_offset_value
def calc_horizontal_offset2(target_mask, source_img):
source_mask = source_img.astype(np.float32).sum(axis=-1) > 750
best_offset = -114514
for offset in range(-200, 200):
offset_mask = np.roll(source_mask, offset, axis=1)
overlap = (target_mask & offset_mask).sum()
if overlap > best_offset:
best_offset = overlap
best_offset_value = offset
return best_offset_value
def get_distract_mask(color_0, color_1, normal_0=None, normal_1=None, thres=0.25, ratio=0.50, outside_thres=0.10, outside_ratio=0.20):
distract_area = np.abs(color_0 - color_1).sum(axis=-1) > thres
if normal_0 is not None and normal_1 is not None:
distract_area |= np.abs(normal_0 - normal_1).sum(axis=-1) > thres
labeled_array, num_features = scipy.ndimage.label(distract_area)
results = []
random_sampled_points = []
for i in range(num_features + 1):
if np.sum(labeled_array == i) > 1000 and np.sum(labeled_array == i) < 100000:
results.append((i, np.sum(labeled_array == i)))
# random sample a point in the area
points = np.argwhere(labeled_array == i)
random_sampled_points.append(points[np.random.randint(0, points.shape[0])])
results = sorted(results, key=lambda x: x[1], reverse=True) # [1:]
distract_mask = np.zeros_like(distract_area)
distract_bbox = np.zeros_like(distract_area)
for i, _ in results:
distract_mask |= labeled_array == i
bbox = np.argwhere(labeled_array == i)
min_x, min_y = bbox.min(axis=0)
max_x, max_y = bbox.max(axis=0)
distract_bbox[min_x:max_x, min_y:max_y] = 1
return distract_mask, distract_bbox, _, _
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--input_mv_dir', type=str, default='result/multiview')
parser.add_argument('--input_obj_dir', type=str, default='result/slrm')
parser.add_argument('--output_dir', type=str, default='result/refined')
parser.add_argument('--outside_ratio', type=float, default=0.20)
parser.add_argument('--no_decompose', action='store_true')
args = parser.parse_args()
import time
start_time = time.time()
for test_idx in os.listdir(args.input_mv_dir):
mv_root_dir = os.path.join(args.input_mv_dir, test_idx)
obj_dir = os.path.join(args.input_obj_dir, test_idx)
fixed_v, fixed_f = None, None
flow_vert, flow_vector = None, None
last_colors, last_normals = None, None
last_front_color, last_front_normal = None, None
distract_mask = None
mv, proj = make_star_cameras_orthographic(8, 1, r=1.2)
mv = mv[[4, 3, 2, 0, 6, 5]]
renderer = NormalsRenderer(mv,proj,(1024,1024))
if not args.no_decompose:
for name_idx, level in zip([3, 1, 2], [2, 1, 0]):
mesh = trimesh.load(obj_dir + f'_{name_idx}.obj')
new_mesh = mesh.split(only_watertight=False)
new_mesh = [ j for j in new_mesh if len(j.vertices) >= 300 ]
mesh = trimesh.Scene(new_mesh).dump(concatenate=True)
mesh_v, mesh_f = mesh.vertices, mesh.faces
if last_colors is None:
images = renderer.render(
torch.tensor(mesh_v, device='cuda').float(),
torch.ones_like(torch.from_numpy(mesh_v), device='cuda').float(),
torch.tensor(mesh_f, device='cuda'),
)
mask = (images[..., 3] < 0.9).cpu().numpy()
colors, normals = [], []
for i in range(6):
color_path = os.path.join(mv_root_dir, f'level{level}', f'color_{i}.png')
normal_path = os.path.join(mv_root_dir, f'level{level}', f'normal_{i}.png')
color = cv2.imread(color_path)
normal = cv2.imread(normal_path)
color = color[..., ::-1]
normal = normal[..., ::-1]
if last_colors is not None:
offset = calc_horizontal_offset(np.array(last_colors[i]), color)
# print('offset', i, offset)
else:
offset = calc_horizontal_offset2(mask[i], color)
# print('init offset', i, offset)
if offset != 0:
color = np.roll(color, offset, axis=1)
normal = np.roll(normal, offset, axis=1)
color = Image.fromarray(color)
normal = Image.fromarray(normal)
colors.append(color)
normals.append(normal)
if last_front_color is not None and level == 0:
distract_mask, distract_bbox, _, _ = get_distract_mask(last_front_color, np.array(colors[0]).astype(np.float32) / 255.0, outside_ratio=args.outside_ratio)
cv2.imwrite(f'{args.output_dir}/{test_idx}/distract_mask.png', distract_mask.astype(np.uint8) * 255)
else:
distract_mask = None
distract_bbox = None
last_front_color = np.array(colors[0]).astype(np.float32) / 255.0
last_front_normal = np.array(normals[0]).astype(np.float32) / 255.0
if last_colors is None:
from copy import deepcopy
last_colors, last_normals = deepcopy(colors), deepcopy(normals)
# my mesh flow weight by nearest vertexs
if fixed_v is not None and fixed_f is not None and level == 1:
t = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f)
fixed_v_cpu = fixed_v.cpu().numpy()
kdtree_anchor = KDTree(fixed_v_cpu)
kdtree_mesh_v = KDTree(mesh_v)
_, idx_anchor = kdtree_anchor.query(mesh_v, k=1)
_, idx_mesh_v = kdtree_mesh_v.query(mesh_v, k=25)
idx_anchor = idx_anchor.squeeze()
neighbors = torch.tensor(mesh_v).cuda()[idx_mesh_v] # V, 25, 3
# calculate the distances neighbors [V, 25, 3]; mesh_v [V, 3] -> [V, 25]
neighbor_dists = torch.norm(neighbors - torch.tensor(mesh_v).cuda()[:, None], dim=-1)
neighbor_dists[neighbor_dists > 0.06] = 114514.
neighbor_weights = torch.exp(-neighbor_dists * 1.)
neighbor_weights = neighbor_weights / neighbor_weights.sum(dim=1, keepdim=True)
anchors = fixed_v[idx_anchor] # V, 3
anchor_normals = calc_vertex_normals(fixed_v, fixed_f)[idx_anchor] # V, 3
dis_anchor = torch.clamp(((anchors - torch.tensor(mesh_v).cuda()) * anchor_normals).sum(-1), min=0) + 0.01
vec_anchor = dis_anchor[:, None] * anchor_normals # V, 3
vec_anchor = vec_anchor[idx_mesh_v] # V, 25, 3
weighted_vec_anchor = (vec_anchor * neighbor_weights[:, :, None]).sum(1) # V, 3
mesh_v += weighted_vec_anchor.cpu().numpy()
t = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f)
mesh_v = torch.tensor(mesh_v, device='cuda', dtype=torch.float32)
mesh_f = torch.tensor(mesh_f, device='cuda')
new_mesh, simp_v, simp_f = geo_refine(mesh_v, mesh_f, colors, normals, fixed_v=fixed_v, fixed_f=fixed_f, distract_mask=distract_mask, distract_bbox=distract_bbox)
# my mesh flow weight by nearest vertexs
try:
if fixed_v is not None and fixed_f is not None and level != 0:
new_mesh_v = new_mesh.vertices.copy()
fixed_v_cpu = fixed_v.cpu().numpy()
kdtree_anchor = KDTree(fixed_v_cpu)
kdtree_mesh_v = KDTree(new_mesh_v)
_, idx_anchor = kdtree_anchor.query(new_mesh_v, k=1)
_, idx_mesh_v = kdtree_mesh_v.query(new_mesh_v, k=25)
idx_anchor = idx_anchor.squeeze()
neighbors = torch.tensor(new_mesh_v).cuda()[idx_mesh_v] # V, 25, 3
# calculate the distances neighbors [V, 25, 3]; new_mesh_v [V, 3] -> [V, 25]
neighbor_dists = torch.norm(neighbors - torch.tensor(new_mesh_v).cuda()[:, None], dim=-1)
neighbor_dists[neighbor_dists > 0.06] = 114514.
neighbor_weights = torch.exp(-neighbor_dists * 1.)
neighbor_weights = neighbor_weights / neighbor_weights.sum(dim=1, keepdim=True)
anchors = fixed_v[idx_anchor] # V, 3
anchor_normals = calc_vertex_normals(fixed_v, fixed_f)[idx_anchor] # V, 3
dis_anchor = torch.clamp(((anchors - torch.tensor(new_mesh_v).cuda()) * anchor_normals).sum(-1), min=0) + 0.01
vec_anchor = dis_anchor[:, None] * anchor_normals # V, 3
vec_anchor = vec_anchor[idx_mesh_v] # V, 25, 3
weighted_vec_anchor = (vec_anchor * neighbor_weights[:, :, None]).sum(1) # V, 3
new_mesh_v += weighted_vec_anchor.cpu().numpy()
new_mesh.vertices = new_mesh_v
except Exception as e:
pass
os.makedirs(f'{args.output_dir}/{test_idx}', exist_ok=True)
new_mesh.export(f'{args.output_dir}/{test_idx}/out_{level}.glb')
if fixed_v is None:
fixed_v, fixed_f = simp_v, simp_f
else:
fixed_f = torch.cat([fixed_f, simp_f + fixed_v.shape[0]], dim=0)
fixed_v = torch.cat([fixed_v, simp_v], dim=0)
# input("Press Enter to continue...")
print('finish', time.time() - start_time)
else:
mesh = trimesh.load(obj_dir + f'_0.obj')
mesh_v, mesh_f = mesh.vertices, mesh.faces
images = renderer.render(
torch.tensor(mesh_v, device='cuda').float(),
torch.ones_like(torch.from_numpy(mesh_v), device='cuda').float(),
torch.tensor(mesh_f, device='cuda'),
)
mask = (images[..., 3] < 0.9).cpu().numpy()
colors, normals = [], []
for i in range(6):
color_path = os.path.join(mv_root_dir, f'level0', f'color_{i}.png')
normal_path = os.path.join(mv_root_dir, f'level0', f'normal_{i}.png')
color = cv2.imread(color_path)
normal = cv2.imread(normal_path)
color = color[..., ::-1]
normal = normal[..., ::-1]
offset = calc_horizontal_offset2(mask[i], color)
if offset != 0:
color = np.roll(color, offset, axis=1)
normal = np.roll(normal, offset, axis=1)
color = Image.fromarray(color)
normal = Image.fromarray(normal)
colors.append(color)
normals.append(normal)
mesh_v = torch.tensor(mesh_v, device='cuda', dtype=torch.float32)
mesh_f = torch.tensor(mesh_f, device='cuda')
new_mesh, _, _ = geo_refine(mesh_v, mesh_f, colors, normals, no_decompose=True, expansion_weight=0.)
os.makedirs(f'{args.output_dir}/{test_idx}', exist_ok=True)
save_py3dmesh_with_trimesh_fast(new_mesh, f'{args.output_dir}/{test_idx}/out_nodecomp.glb', apply_sRGB_to_LinearRGB=False)