[CG Pipeline] Refactor

This commit is contained in:
brothermechanic 2023-11-13 13:07:33 +00:00 committed by Igor Brylyov
parent 6538f70d54
commit b3612d8655
23 changed files with 634 additions and 645 deletions

View file

@ -18,43 +18,94 @@ __version__ = '0.2'
import logging
import os
import shutil
import time
import bpy
import addon_utils
import BakeWrangler
from BakeWrangler.nodes.node_tree import BW_TREE_VERSION
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
# COLLECTIONS NAMIG CONVENTION
parts_col_name = 'Parts'
lcs_col_name = 'LCS'
hierarchy_col_name = 'Hierarchy'
lowpoly_col_name = 'Lowpoly'
# LCS POINT'S SUFFIXES CONVENTION
inlet = '_in'
outlet = '_out'
root = '_root'
# CG ASSETS SUFFIXES CONVENTION
hightpoly = '_hp'
midpoly = 'mp'
lowpoly = '_lp'
render = '_render'
bw_config_1 = {
'passes': {
'diffuse': {
'bake_cat': 'PBR',
'bake_subcat': 'ALBEDO',
'suffix': 'D',
},
'nornal': {
'bake_cat': 'CORE',
'bake_subcat': 'NORMAL',
'suffix': 'N',
},
'ao': {
'bake_cat': 'CORE',
'bake_subcat': 'AO',
'bake_samples': 16,
'suffix': 'AO',
},
'roughness': {
'bake_cat': 'CORE',
'bake_subcat': 'ROUGHNESS',
'suffix': 'R',
},
'metallic': {
'bake_cat': 'PBR',
'bake_subcat': 'METALLIC',
'suffix': 'M',
},
'uv': {
'bake_cat': 'WRANG',
'bake_subcat': 'ISLANDID',
'suffix': 'UV',
},
},
'margin': 8,
'ray_dist': 0.001,
'bake_samples': 4,
}
bw_config_10 = {
'passes': {
'nornal': {
'bake_cat': 'CORE',
'bake_subcat': 'NORMAL',
'suffix': 'N',
},
'ao': {
'bake_cat': 'CORE',
'bake_subcat': 'AO',
'bake_samples': 16,
'suffix': 'AO',
},
},
'margin': 8,
'ray_dist': 0.01,
'bake_samples': 4,
}
configs = {'bake_1': bw_config_1, 'bake_10': bw_config_10}
def bw_submit(lowpoly_obj_names, resolution=4096, tree_name='robossembler', area=None):
def bw_submit(lowpoly_obj_names,
bw_config,
resolution,
bake_path=None,
area=None,
**cg_config):
''' Submit session and bake textures with BakeWrangler addon. '''
BakeWrangler.register()
asm_name = os.path.basename(bpy.context.blend_data.filepath).split('.')[0]
asm_path = os.path.dirname(bpy.context.blend_data.filepath)
textures_path = os.path.join(asm_path,'textures').replace('\\', '/')
bake_path = os.path.join(textures_path,'bake')
bake_path_double = os.path.join(textures_path,'bake_double')
os.makedirs(bake_path, exist_ok=True)
if not bake_path:
asm_path = os.path.dirname(bpy.context.blend_data.filepath)
textures_path = os.path.join(asm_path, 'textures').replace('\\', '/')
bake_path = os.path.join(textures_path, 'bake').replace('\\', '/')
if os.path.exists(bake_path):
shutil.rmtree(bake_path)
os.makedirs(bake_path, exist_ok=True)
# create node tree
tree = bpy.data.node_groups.new(name=tree_name, type='BakeWrangler_Tree')
baking_tree_name = ''.join([item[0] for item in bw_config['passes'].keys()])
tree = bpy.data.node_groups.new(name=baking_tree_name, type='BakeWrangler_Tree')
# for default used Compositing Node Editor area
if not area:
@ -88,10 +139,10 @@ def bw_submit(lowpoly_obj_names, resolution=4096, tree_name='robossembler', area
tree.nodes.new('BakeWrangler_OutputSettings').location = (600, 0)
tree.nodes['Mesh Settings'].pinned = True
tree.nodes['Mesh Settings']['margin'] = 8
tree.nodes['Mesh Settings']['ray_dist'] = 0.001
tree.nodes['Mesh Settings']['margin'] = bw_config['margin']
tree.nodes['Mesh Settings']['ray_dist'] = bw_config['ray_dist']
tree.nodes['Sample Settings'].pinned = True
tree.nodes['Sample Settings']['bake_samples'] = 4
tree.nodes['Sample Settings']['bake_samples'] = bw_config['bake_samples']
tree.nodes['Pass Settings'].pinned = True
tree.nodes['Pass Settings']['res_bake_x'] = resolution
tree.nodes['Pass Settings']['res_bake_y'] = resolution
@ -105,187 +156,100 @@ def bw_submit(lowpoly_obj_names, resolution=4096, tree_name='robossembler', area
# batch bake node
node_batch = tree.nodes.new('BakeWrangler_Output_Batch_Bake')
node_batch.location = (1000, -500)
node_batch_double = tree.nodes.new('BakeWrangler_Output_Batch_Bake')
node_batch_double.location = (2000, -1500)
node_batch.name = 'Batch'
node_y_pos = 0
node_pos = 500
pass_socket = 0
pass_socket_double = 0
for lp_name in lowpoly_obj_names:
# run for eatch lowpoly object
lp = bpy.data.objects[lp_name]
mp = bpy.data.objects['_'.join(lp.name.split('_')[:-1] + [midpoly])]
mp = bpy.data.objects[
'_'.join(lp.name.split('_')[:-1] + [cg_config['midpoly']])
]
img_name = '_'.join(lp.name.split('_')[:-1]) + '_'
node_inputs = tree.nodes.new('BakeWrangler_Bake_Mesh')
node_inputs.location = (-700, node_y_pos - 500)
node_inputs.location = (-700, node_y_pos - node_pos)
node_inputs.inputs['Target'].value = lp
node_inputs.inputs['Source'].value = mp
def bw_pass(node_pos, bake_pass, pass_socket):
node_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
node_pass.location = (-200, node_y_pos - node_pos)
node_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
node_img.location = (200, node_y_pos - node_pos)
tree.links.new(node_pass.outputs['Color'], node_img.inputs['Color'])
node_pass.bake_cat = bake_pass['bake_cat']
if bake_pass.get('bake_cat') == 'CORE':
node_pass.bake_core = bake_pass['bake_subcat']
elif bake_pass.get('bake_cat') == 'PBR':
node_pass.bake_pbr = bake_pass['bake_subcat']
elif bake_pass.get('bake_cat') == 'WRANG':
node_pass.bake_wrang = bake_pass['bake_subcat']
if bake_pass.get('bake_samples'):
node_pass.bake_samples = bake_pass['bake_samples']
node_img.inputs['Color'].suffix = bake_pass['suffix']
node_img.inputs['Split Output'].disp_path = bake_path
node_img.inputs['Split Output'].img_name = img_name
tree.links.new(node_inputs.outputs['Mesh'], node_pass.inputs[1])
tree.links.new(node_img.outputs['Bake'], node_batch.inputs[pass_socket])
# bake passes
# Diffuse
node_d_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
node_d_pass.location = (-200, node_y_pos - 500)
node_d_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
node_d_img.location = (200, node_y_pos - 500)
tree.links.new(node_d_pass.outputs['Color'], node_d_img.inputs['Color'])
node_d_pass.bake_cat = 'PBR'
node_d_pass.bake_pbr = 'ALBEDO'
node_d_img.inputs['Color'].suffix = 'D'
node_d_img.inputs['Split Output'].disp_path = bake_path
node_d_img.inputs['Split Output'].img_name = img_name
'''
## Curvature
node_c_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
node_c_pass.location = (-200, node_y_pos -1000)
node_c_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
node_c_img.location = (200, node_y_pos -1000)
tree.links.new(node_c_pass.outputs['Color'], node_c_img.inputs['Color'])
node_c_pass.bake_cat = 'WRANG'
node_c_pass.bake_wrang = 'CURVATURE'
node_c_img.inputs['Color'].suffix = 'C'
node_c_img.inputs['Split Output'].disp_path = bake_path
node_c_img.inputs['Split Output'].img_name = img_name
'''
# Normal
node_n_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
node_n_pass.location = (-200, node_y_pos - 1500)
node_n_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
node_n_img.location = (200, node_y_pos - 1500)
node_n_img_double = tree.nodes.new('BakeWrangler_Output_Image_Path')
node_n_img_double.location = (200, node_y_pos - 1750)
tree.links.new(node_n_pass.outputs['Color'], node_n_img.inputs['Color'])
tree.links.new(node_n_pass.outputs['Color'], node_n_img_double.inputs['Color'])
node_n_pass.bake_cat = 'CORE'
node_n_pass.bake_core = 'NORMAL'
node_n_img.inputs['Color'].suffix = 'N'
node_n_img.inputs['Split Output'].disp_path = bake_path
node_n_img.inputs['Split Output'].img_name = img_name
node_n_img_double.inputs['Color'].suffix = 'N'
node_n_img_double.inputs['Split Output'].disp_path = bake_path_double
node_n_img_double.inputs['Split Output'].img_name = img_name
# AO
node_ao_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
node_ao_pass.location = (-200, node_y_pos - 2000)
node_ao_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
node_ao_img.location = (200, node_y_pos - 2000)
node_ao_img_double = tree.nodes.new('BakeWrangler_Output_Image_Path')
node_ao_img_double.location = (200, node_y_pos - 2250)
tree.links.new(node_ao_pass.outputs['Color'], node_ao_img.inputs['Color'])
tree.links.new(node_ao_pass.outputs['Color'], node_ao_img_double.inputs['Color'])
node_ao_pass.bake_cat = 'CORE'
node_ao_pass.bake_core = 'AO'
node_ao_pass.bake_samples = 32
node_ao_img.inputs['Color'].suffix = 'AO'
node_ao_img.inputs['Split Output'].disp_path = bake_path
node_ao_img.inputs['Split Output'].img_name = img_name
node_ao_img_double.inputs['Color'].suffix = 'AO'
node_ao_img_double.inputs['Split Output'].disp_path = bake_path_double
node_ao_img_double.inputs['Split Output'].img_name = img_name
# Roughness
node_r_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
node_r_pass.location = (-200, node_y_pos - 2500)
node_r_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
node_r_img.location = (200, node_y_pos - 2500)
tree.links.new(node_r_pass.outputs['Color'], node_r_img.inputs['Color'])
node_r_pass.bake_cat = 'CORE'
node_r_pass.bake_core = 'ROUGHNESS'
node_r_img.inputs['Color'].suffix = 'R'
node_r_img.inputs['Split Output'].disp_path = bake_path
node_r_img.inputs['Split Output'].img_name = img_name
# Metallic
node_m_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
node_m_pass.location = (-200, node_y_pos - 3000)
node_m_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
node_m_img.location = (200, node_y_pos - 3000)
tree.links.new(node_m_pass.outputs['Color'], node_m_img.inputs['Color'])
node_m_pass.bake_cat = 'PBR'
node_m_pass.bake_pbr = 'METALLIC'
node_m_img.inputs['Color'].suffix = 'M'
node_m_img.inputs['Split Output'].disp_path = bake_path
node_m_img.inputs['Split Output'].img_name = img_name
# UV
node_uv_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
node_uv_pass.location = (-200, node_y_pos - 3500)
node_uv_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
node_uv_img.location = (200, node_y_pos - 3500)
tree.links.new(node_uv_pass.outputs['Color'], node_uv_img.inputs['Color'])
node_uv_pass.bake_cat = 'WRANG'
node_uv_pass.bake_wrang = 'ISLANDID'
node_uv_img.inputs['Color'].suffix = 'UV'
node_uv_img.inputs['Split Output'].disp_path = bake_path
node_uv_img.inputs['Split Output'].img_name = img_name
# connect meshes to passes
tree.links.new(node_inputs.outputs['Mesh'], node_d_pass.inputs[1])
'''
tree.links.new(node_inputs.outputs['Mesh'], node_c_pass.inputs[1])
'''
tree.links.new(node_inputs.outputs['Mesh'], node_n_pass.inputs[1])
tree.links.new(node_inputs.outputs['Mesh'], node_ao_pass.inputs[1])
tree.links.new(node_inputs.outputs['Mesh'], node_r_pass.inputs[1])
tree.links.new(node_inputs.outputs['Mesh'], node_m_pass.inputs[1])
tree.links.new(node_inputs.outputs['Mesh'], node_uv_pass.inputs[1])
# batch bake node
tree.links.new(node_d_img.outputs['Bake'], node_batch.inputs[pass_socket])
'''
pass_socket += 1
tree.links.new(node_c_img.outputs['Bake'], node_batch.inputs[pass_socket])
'''
pass_socket += 1
tree.links.new(node_n_img.outputs['Bake'], node_batch.inputs[pass_socket])
pass_socket += 1
tree.links.new(node_ao_img.outputs['Bake'], node_batch.inputs[pass_socket])
pass_socket += 1
tree.links.new(node_r_img.outputs['Bake'], node_batch.inputs[pass_socket])
pass_socket += 1
tree.links.new(node_m_img.outputs['Bake'], node_batch.inputs[pass_socket])
pass_socket += 1
tree.links.new(node_uv_img.outputs['Bake'], node_batch.inputs[pass_socket])
pass_socket += 1
# batch bake node double
tree.links.new(node_n_img_double.outputs['Bake'],
node_batch_double.inputs[pass_socket_double])
pass_socket_double += 1
tree.links.new(node_ao_img_double.outputs['Bake'],
node_batch_double.inputs[pass_socket_double])
pass_socket_double += 1
for bake_pass in bw_config['passes'].keys():
bw_pass(node_pos, bw_config['passes'][bake_pass], pass_socket)
pass_socket += 1
node_pos += 500
reroute1 = tree.nodes.new('NodeReroute')
reroute1.location = (-1000, 0)
reroute2 = tree.nodes.new('NodeReroute')
reroute2.location = (2000, 0)
tree.links.new(reroute1.outputs[0], reroute2.inputs[0])
node_pos = 500
node_y_pos -= 4000
bpy.ops.bake_wrangler.bake_pass(tree=tree_name, node=node_batch.name, sock=-1)
logger.info('Baking first pass is finished!')
return tree, pass_socket
# double pass
tree.nodes['Mesh Settings']['ray_dist'] = 0.01
bpy.ops.bake_wrangler.bake_pass(tree=tree_name, node=node_batch_double.name, sock=-1)
logger.info('Baking double pass is finished!')
def bw_bake(lowpoly_obj_names, textures_path, resolution, **cg_config):
BakeWrangler.register()
output_paths = []
for config in configs:
# 1 pass
bake_path = os.path.join(textures_path, config).replace('\\', '/')
if os.path.exists(bake_path):
shutil.rmtree(bake_path)
os.makedirs(bake_path, exist_ok=True)
tree_obj, count = bw_submit(
lowpoly_obj_names, configs[config], resolution, bake_path, **cg_config)
return textures_path
node_batch = tree_obj.nodes['Batch']
bpy.ops.bake_wrangler.bake_pass(tree=tree_obj.name, node=node_batch.name, sock=-1)
logger.info('Baking pipeline is started!')
# delay until render will be complete
# ______________________________
exit_delay = 180
start_time = time.time()
while bpy.context.window_manager.bw_status == 1:
time.sleep(2)
if not os.listdir(bake_path):
if time.time() > start_time + exit_delay:
break
if not len(os.listdir(bake_path)) < count:
break
# ______________________________
assert len(os.listdir(bake_path)) == count, (
'Warning! Baked only {} of {} textures!'.format(
len(os.listdir(bake_path)), count))
logger.info('Baking %s textures is finished!', len(os.listdir(bake_path)))
output_paths.append(bake_path)
return output_paths