Spaces:
Paused
Paused
| # Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved. | |
| # | |
| # NVIDIA CORPORATION and its licensors retain all intellectual property | |
| # and proprietary rights in and to this software, related documentation | |
| # and any modifications thereto. Any use, reproduction, disclosure or | |
| # distribution of this software and related documentation without an express | |
| # license agreement from NVIDIA CORPORATION is strictly prohibited. | |
| import os | |
| import numpy as np | |
| import torch | |
| from . import util | |
| from . import texture | |
| from . import mesh | |
| ###################################################################################### | |
| # .mtl material format loading / storing | |
| ###################################################################################### | |
| def load_mtl(fn, clear_ks=True): | |
| import re | |
| mtl_path = os.path.dirname(fn) | |
| # Check if file exists | |
| if not os.path.exists(fn): | |
| print(f"Warning: Material file {fn} does not exist, returning empty material list") | |
| return [] | |
| # Read file | |
| try: | |
| with open(fn) as f: | |
| lines = f.readlines() | |
| except Exception as e: | |
| print(f"Warning: Could not read material file {fn}: {e}, returning empty material list") | |
| return [] | |
| # Parse materials | |
| materials = [] | |
| for line in lines: | |
| split_line = re.split(' +|\t+|\n+', line.strip()) | |
| if len(split_line) == 0: | |
| continue | |
| prefix = split_line[0].lower() | |
| data = split_line[1:] | |
| if 'newmtl' in prefix: | |
| material = {'name' : data[0]} | |
| materials += [material] | |
| elif materials: | |
| if 'bsdf' in prefix or 'map_kd' in prefix or 'map_ks' in prefix or 'bump' in prefix: | |
| material[prefix] = data[0] | |
| else: | |
| try: | |
| material[prefix] = torch.tensor(tuple(float(d) for d in data), dtype=torch.float32, device='cuda') | |
| except (ValueError, IndexError) as e: | |
| print(f"Warning: Could not parse material property {prefix} with data {data}: {e}") | |
| continue | |
| # Convert everything to textures. Our code expects 'kd' and 'ks' to be texture maps. So replace constants with 1x1 maps | |
| for mat in materials: | |
| if not 'bsdf' in mat: | |
| mat['bsdf'] = 'pbr' | |
| # Handle kd (diffuse color) | |
| if 'map_kd' in mat: | |
| try: | |
| mat['kd'] = texture.load_texture2D(os.path.join(mtl_path, mat['map_kd'])) | |
| except Exception as e: | |
| print(f"Warning: Could not load kd texture {mat['map_kd']}: {e}, using default") | |
| mat['kd'] = texture.Texture2D(torch.tensor([0.5, 0.5, 0.5], dtype=torch.float32, device='cuda')) | |
| elif 'kd' in mat: | |
| mat['kd'] = texture.Texture2D(mat['kd']) | |
| else: | |
| # Default diffuse color | |
| mat['kd'] = texture.Texture2D(torch.tensor([0.5, 0.5, 0.5], dtype=torch.float32, device='cuda')) | |
| # Handle ks (specular color) | |
| if 'map_ks' in mat: | |
| try: | |
| mat['ks'] = texture.load_texture2D(os.path.join(mtl_path, mat['map_ks']), channels=3) | |
| except Exception as e: | |
| print(f"Warning: Could not load ks texture {mat['map_ks']}: {e}, using default") | |
| mat['ks'] = texture.Texture2D(torch.tensor([0.0, 0.0, 0.0], dtype=torch.float32, device='cuda')) | |
| elif 'ks' in mat: | |
| mat['ks'] = texture.Texture2D(mat['ks']) | |
| else: | |
| # Default specular color | |
| mat['ks'] = texture.Texture2D(torch.tensor([0.0, 0.0, 0.0], dtype=torch.float32, device='cuda')) | |
| # Handle normal map | |
| if 'bump' in mat: | |
| try: | |
| mat['normal'] = texture.load_texture2D(os.path.join(mtl_path, mat['bump']), lambda_fn=lambda x: x * 2 - 1, channels=3) | |
| except Exception as e: | |
| print(f"Warning: Could not load normal texture {mat['bump']}: {e}, using default") | |
| mat['normal'] = texture.Texture2D(torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device='cuda')) | |
| # Convert Kd from sRGB to linear RGB | |
| try: | |
| mat['kd'] = texture.srgb_to_rgb(mat['kd']) | |
| except Exception as e: | |
| print(f"Warning: Could not convert kd to linear RGB: {e}") | |
| if clear_ks: | |
| # Override ORM occlusion (red) channel by zeros. We hijack this channel | |
| try: | |
| for mip in mat['ks'].getMips(): | |
| mip[..., 0] = 0.0 | |
| except Exception as e: | |
| print(f"Warning: Could not clear ks occlusion channel: {e}") | |
| return materials | |
| def save_mtl(fn, material): | |
| folder = os.path.dirname(fn) | |
| with open(fn, "w") as f: | |
| f.write('newmtl defaultMat\n') | |
| if material is not None: | |
| f.write('bsdf %s\n' % material['bsdf']) | |
| f.write('map_kd texture_kd.png\n') | |
| texture.save_texture2D(os.path.join(folder, 'texture_kd.png'), texture.rgb_to_srgb(material['kd'])) | |
| f.write('map_ks texture_ks.png\n') | |
| texture.save_texture2D(os.path.join(folder, 'texture_ks.png'), material['ks']) | |
| f.write('bump texture_n.png\n') | |
| texture.save_texture2D(os.path.join(folder, 'texture_n.png'), material['normal'], lambda_fn=lambda x:(x+1)*0.5) | |
| else: | |
| f.write('Kd 1 1 1\n') | |
| f.write('Ks 0 0 0\n') | |
| f.write('Ka 0 0 0\n') | |
| f.write('Tf 1 1 1\n') | |
| f.write('Ni 1\n') | |
| f.write('Ns 0\n') | |
| ###################################################################################### | |
| # Merge multiple materials into a single uber-material | |
| ###################################################################################### | |
| def _upscale_replicate(x, full_res): | |
| x = x.permute(0, 3, 1, 2) | |
| x = torch.nn.functional.pad(x, (0, full_res[1] - x.shape[3], 0, full_res[0] - x.shape[2]), 'replicate') | |
| return x.permute(0, 2, 3, 1).contiguous() | |
| def merge_materials(materials, texcoords, tfaces, mfaces): | |
| assert len(materials) > 0 | |
| for mat in materials: | |
| assert mat['bsdf'] == materials[0]['bsdf'], "All materials must have the same BSDF (uber shader)" | |
| assert ('normal' in mat) is ('normal' in materials[0]), "All materials must have either normal map enabled or disabled" | |
| uber_material = { | |
| 'name' : 'uber_material', | |
| 'bsdf' : materials[0]['bsdf'], | |
| } | |
| textures = ['kd', 'ks', 'normal'] | |
| # Find maximum texture resolution across all materials and textures | |
| max_res = None | |
| for mat in materials: | |
| for tex in textures: | |
| tex_res = np.array(mat[tex].getRes()) if tex in mat else np.array([1, 1]) | |
| max_res = np.maximum(max_res, tex_res) if max_res is not None else tex_res | |
| # Compute size of compund texture and round up to nearest PoT | |
| full_res = 2**np.ceil(np.log2(max_res * np.array([1, len(materials)]))).astype(np.int) | |
| # Normalize texture resolution across all materials & combine into a single large texture | |
| for tex in textures: | |
| if tex in materials[0]: | |
| tex_data = torch.cat(tuple(util.scale_img_nhwc(mat[tex].data, tuple(max_res)) for mat in materials), dim=2) # Lay out all textures horizontally, NHWC so dim2 is x | |
| tex_data = _upscale_replicate(tex_data, full_res) | |
| uber_material[tex] = texture.Texture2D(tex_data) | |
| # Compute scaling values for used / unused texture area | |
| s_coeff = [full_res[0] / max_res[0], full_res[1] / max_res[1]] | |
| # Recompute texture coordinates to cooincide with new composite texture | |
| new_tverts = {} | |
| new_tverts_data = [] | |
| for fi in range(len(tfaces)): | |
| matIdx = mfaces[fi] | |
| for vi in range(3): | |
| ti = tfaces[fi][vi] | |
| if not (ti in new_tverts): | |
| new_tverts[ti] = {} | |
| if not (matIdx in new_tverts[ti]): # create new vertex | |
| new_tverts_data.append([(matIdx + texcoords[ti][0]) / s_coeff[1], texcoords[ti][1] / s_coeff[0]]) # Offset texture coodrinate (x direction) by material id & scale to local space. Note, texcoords are (u,v) but texture is stored (w,h) so the indexes swap here | |
| new_tverts[ti][matIdx] = len(new_tverts_data) - 1 | |
| tfaces[fi][vi] = new_tverts[ti][matIdx] # reindex vertex | |
| return uber_material, new_tverts_data, tfaces | |