From 560ab3bfa43a2b3d94614451665254fdeed5df1f Mon Sep 17 00:00:00 2001 From: ethan Date: Sun, 23 Jun 2024 19:08:41 -0400 Subject: [PATCH] Modify pack_maps operator behavior - Add check to export_maps in case a pack map is disabled but is selected in one of the pack maps channels UI - Added is_pack_maps_enabled function for pack_maps operator - Refactored CGDJay PR to be closer to project lint standards - Optimized `poll` on pack_maps operator --- __init__.py | 8 +- constants.py | 2 +- operators/operators.py | 261 +++++++++++++++++++---------------------- preferences.py | 51 +++----- ui.py | 60 ++++------ utils/generic.py | 2 +- 6 files changed, 168 insertions(+), 216 deletions(-) diff --git a/__init__.py b/__init__.py index 9cc8924..04bc630 100644 --- a/__init__.py +++ b/__init__.py @@ -2,11 +2,9 @@ "name": "GrabDoc", "author": "Ethan Simon-Law", "location": "3D View > Sidebar > GrabDoc", - # NOTE: Also change: - # - Update check in preferences.py - # - Panel version in generic.py - "version": (1, 4, 3), - "blender": (4, 0, 2), + # NOTE: Also change update check in preferences.py + "version": (1, 4, 4), + "blender": (4, 1, 0), "tracker_url": "https://discord.com/invite/wHAyVZG", "category": "3D View" } diff --git a/constants.py b/constants.py index c329af4..f071778 100644 --- a/constants.py +++ b/constants.py @@ -29,7 +29,7 @@ class Global: COLOR_ID = "color" EMISSIVE_ID = "emissive" ROUGHNESS_ID = "roughness" - METALLIC_ID = "metallic" + METALLIC_ID = "metallic" NORMAL_NAME = NORMAL_ID.capitalize() CURVATURE_NAME = CURVATURE_ID.capitalize() diff --git a/operators/operators.py b/operators/operators.py index d304f2c..8565508 100644 --- a/operators/operators.py +++ b/operators/operators.py @@ -1,15 +1,11 @@ import os import time - import numpy - - - import bpy import blf from bpy.types import SpaceView3D, Event, Context, Operator, UILayout -from bpy.props import StringProperty, BoolProperty +from bpy.props import StringProperty from ..constants import Global, Error from ..utils.render import get_rendered_objects @@ -156,14 +152,19 @@ def export(context: Context, suffix: str, path: str = None) -> str: return path def execute(self, context: Context): - + gd = context.scene.gd report_value, report_string = \ bad_setup_check(context, active_export=True) if report_value: self.report({'ERROR'}, report_string) return {'CANCELLED'} + if gd.use_pack_maps is True and not is_pack_maps_enabled(): + self.report( + {'ERROR'}, + "Map packing enabled but incorrect export maps enabled" + ) + return {'CANCELLED'} - gd = context.scene.gd self.map_name = 'export' bake_maps = get_bake_maps() @@ -216,29 +217,20 @@ def execute(self, context: Context): if gd.export_plane: export_plane(context) - # Call for Original Context Mode, use bpy.ops - # so that Blenders viewport refreshes if active_selected: context.view_layer.objects.active = bpy.data.objects[activeCallback] - if bpy.ops.object.mode_set.poll(): bpy.ops.object.mode_set(mode=modeCallback) - # End the timer & UI progress bar end = time.time() exc_time = round(end - start, 2) - self.report( {'INFO'}, f"{Error.EXPORT_COMPLETE} (execution time: {exc_time}s)" ) - context.window_manager.progress_end() - #run the pack maps operator on export if the setting is enabled - if gd.use_pack_maps==True: - bpy.ops.grab_doc.pack_maps(remove_maps=False) - - + if gd.use_pack_maps is True: + bpy.ops.grab_doc.pack_maps() return {'FINISHED'} @@ -597,14 +589,15 @@ def draw(self, _context: Context): ################################################ -# TODO: CHANNEL PACKING +# CHANNEL PACKING ################################################ -# original code sourced from : -# https://blender.stackexchange.com/questions/274712/how-to-channel-pack-texture-in-python - -def pack_image_channels(pack_order,PackName): +def pack_image_channels(pack_order, PackName): + """ + NOTE: Original code sourced from: + https://blender.stackexchange.com/questions/274712/how-to-channel-pack-texture-in-python + """ dst_array = None has_alpha = False @@ -631,137 +624,130 @@ def pack_image_channels(pack_order,PackName): return dst_image -def Return_Channel_Path(context ,Channel): - gd = context.scene.gd - if Channel == 'none': - return ("") - if Channel == 'normals': - return ((os.path.join(gd.export_path,gd.export_name+'_'+gd.occlusion[0].suffix+get_format()))) - if Channel == 'curvature': - return((os.path.join(gd.export_path,gd.export_name+'_'+gd.curvature[0].suffix+get_format()))) - if Channel == 'occlusion': - return((os.path.join(gd.export_path,gd.export_name+'_'+gd.occlusion[0].suffix+get_format()))) - if Channel == 'height': - return((os.path.join(gd.export_path,gd.export_name+'_'+gd.height[0].suffix+get_format()))) - if Channel == 'id': - return((os.path.join(gd.export_path,gd.export_name+'_'+gd.id[0].suffix+get_format()))) - if Channel == 'alpha': - return((os.path.join(gd.export_path,gd.export_name+'_'+gd.alpha[0].suffix+get_format()))) - if Channel == 'color': - return((os.path.join(gd.export_path,gd.export_name+'_'+gd.color[0].suffix+get_format()))) - if Channel == 'emissive': - return((os.path.join(gd.export_path,gd.export_name+'_'+gd.emissive[0].suffix+get_format()))) - if Channel == 'roughness': - return((os.path.join(gd.export_path,gd.export_name+'_'+gd.roughness[0].suffix+get_format()))) - if Channel == 'metallic': - return((os.path.join(gd.export_path,gd.export_name+'_'+gd.metallic[0].suffix+get_format()))) - return False - - -class GRABDOC_OT_pack_maps(Operator): - """Pack bake maps into single texture""" - bl_idname = "grab_doc.pack_maps" - bl_label = "Pack Maps" - bl_options = {'REGISTER', 'UNDO'} - remove_maps : BoolProperty (name='Remove packed maps' ,default=False) +def get_channel_path(channel: str) -> str | None: + """Get the channel path of the given channel name. + + If the channel path is not found returns `None`.""" + gd = bpy.context.scene.gd + fmt = get_format() + filename = "" + if channel == 'normals': + filename = gd.export_name + '_' + gd.occlusion[0].suffix + fmt + elif channel == 'curvature': + filename = gd.export_name + '_' + gd.curvature[0].suffix + fmt + elif channel == 'occlusion': + filename = gd.export_name + '_' + gd.occlusion[0].suffix + fmt + elif channel == 'height': + filename = gd.export_name + '_' + gd.height[0].suffix + fmt + elif channel == 'id': + filename = gd.export_name + '_' + gd.id[0].suffix + fmt + elif channel == 'alpha': + filename = gd.export_name + '_' + gd.alpha[0].suffix + fmt + elif channel == 'color': + filename = gd.export_name + '_' + gd.color[0].suffix + fmt + elif channel == 'emissive': + filename = gd.export_name + '_' + gd.emissive[0].suffix + fmt + elif channel == 'roughness': + filename = gd.export_name + '_' + gd.roughness[0].suffix + fmt + elif channel == 'metallic': + filename = gd.export_name + '_' + gd.metallic[0].suffix + fmt + if filename == "": + return None + filepath = os.path.join(gd.export_path, filename) + if not os.path.exists(filepath): + return None + return filepath + + +def is_pack_maps_enabled() -> bool: + """Checks if the chosen pack channels + match the enabled maps to export. + + This function also returns True if a required + bake map is not enabled but the texture exists.""" + bake_maps = get_bake_maps() + bake_map_names = ['none'] + for bake_map in bake_maps: + bake_map_names.append(bake_map.ID) + + gd = bpy.context.scene.gd + if gd.channel_r not in bake_map_names \ + and get_channel_path(gd.channel_r) is None: + return False + if gd.channel_g not in bake_map_names \ + and get_channel_path(gd.channel_g) is None: + return False + if gd.channel_b not in bake_map_names \ + and get_channel_path(gd.channel_b) is None: + return False + if gd.channel_a not in bake_map_names \ + and get_channel_path(gd.channel_a) is None: + return False + return True + + +class GRABDOC_OT_pack_maps(OpInfo, Operator): + """Merge previously exported bake maps into single packed texture""" + bl_idname = "grab_doc.pack_maps" + bl_label = "Run Pack" - # Poll to check if all required images are correctly present in the export path @classmethod def poll(cls, context: Context) -> bool: gd = context.scene.gd - - RChannel=gd.channel_R - GChannel=gd.channel_G - BChannel=gd.channel_B - AChannel=gd.channel_A - - - r=(Return_Channel_Path(context,RChannel)) - g=(Return_Channel_Path(context,GChannel)) - b=(Return_Channel_Path(context,BChannel)) - a=(Return_Channel_Path(context,AChannel)) - - #Alpha requires a edge case as it should be the only option that uses the "none" setting - if gd.channel_A == 'none': - if os.path.exists(r) & os.path.exists(g) & os.path.exists(b)==True: - return True - else: - return False - else: - if os.path.exists(r) & os.path.exists(g) & os.path.exists(b) & os.path.exists(a)==True: - return True - else: - return False - + r = get_channel_path(gd.channel_r) + g = get_channel_path(gd.channel_g) + b = get_channel_path(gd.channel_b) + a = get_channel_path(gd.channel_a) + if not all((r, g, b)): + return False + if gd.channel_a != 'none' and a is None: + return False + return True def execute(self, context): - gd = context.scene.gd - Name = f"{gd.export_name}" - - PackName= (Name+"_"+gd.pack_name) - Path= gd.export_path - - #Loads all images into blender to avoid using a seperate python module to convert to np array - ImageR= bpy.data.images.load(Return_Channel_Path(context,gd.channel_R)) - ImageG=bpy.data.images.load(Return_Channel_Path(context,gd.channel_G)) - ImageB=bpy.data.images.load(Return_Channel_Path(context,gd.channel_B)) - if gd.channel_A != 'none': - ImageA=bpy.data.images.load(Return_Channel_Path(context,gd.channel_A)) - - - if gd.channel_A == 'none': - pack_order = [ - (ImageR, (0, 0)) - ,(ImageG, (0, 1)) - ,(ImageB, (0, 2)) - ] - else: - pack_order = [ - (ImageR, (0, 0)) - ,(ImageG, (0, 1)) - ,(ImageB, (0, 2)) - ,(ImageA, (0, 3)) - ] - - dst_image=pack_image_channels(pack_order,PackName) - - dst_image.filepath_raw = Path+"//"+PackName+get_format() + pack_name = gd.export_name + "_" + gd.pack_name + path = gd.export_path + + # Loads all images into blender to avoid using a + # separate python module to convert to np array + image_r = bpy.data.images.load(get_channel_path(gd.channel_r)) + image_g = bpy.data.images.load(get_channel_path(gd.channel_g)) + image_b = bpy.data.images.load(get_channel_path(gd.channel_b)) + pack_order = [ + (image_r, (0, 0)), + (image_g, (0, 1)), + (image_b, (0, 2)) + ] + if gd.channel_a != 'none': + image_a = bpy.data.images.load(get_channel_path(gd.channel_a)) + pack_order.append((image_a, (0, 3))) + + dst_image = pack_image_channels(pack_order, pack_name) + dst_image.filepath_raw = path+"//"+pack_name+get_format() dst_image.file_format = gd.format dst_image.save() - - #Remove images from blend file to keep it clean - - bpy.data.images.remove(ImageR) - bpy.data.images.remove(ImageG) - bpy.data.images.remove(ImageB) - if gd.channel_A != 'none': - bpy.data.images.remove(ImageA) + # Remove images from blend file to keep it clean + bpy.data.images.remove(image_r) + bpy.data.images.remove(image_g) + bpy.data.images.remove(image_b) + if gd.channel_a != 'none': + bpy.data.images.remove(image_a) bpy.data.images.remove(dst_image) - #option to delete the extra maps through the operator panel - if self.remove_maps==True: - if gd.channel_A != 'none': - os.remove(Return_Channel_Path(context,gd.channel_R)) - os.remove(Return_Channel_Path(context,gd.channel_G)) - os.remove(Return_Channel_Path(context,gd.channel_B)) - os.remove(Return_Channel_Path(context,gd.channel_A)) - else: - os.remove(Return_Channel_Path(context,gd.channel_R)) - os.remove(Return_Channel_Path(context,gd.channel_G)) - os.remove(Return_Channel_Path(context,gd.channel_B)) - - #reset value as this would be best a manual opt in on a final pack to prevent re exporting - self.remove_maps=False - + # Option to delete the extra maps through the operator panel + if gd.remove_original_maps is True: + os.remove(get_channel_path(gd.channel_r)) + os.remove(get_channel_path(gd.channel_g)) + os.remove(get_channel_path(gd.channel_b)) + if gd.channel_a != 'none': + os.remove(get_channel_path(gd.channel_a)) return {'FINISHED'} - - ################################################ # REGISTRATION ################################################ @@ -780,7 +766,6 @@ def execute(self, context): GRABDOC_OT_leave_map_preview, GRABDOC_OT_export_current_preview, GRABDOC_OT_config_maps, - #GRABDOC_OT_map_pack_info GRABDOC_OT_pack_maps ) diff --git a/preferences.py b/preferences.py index 7a581dc..dad6cdd 100644 --- a/preferences.py +++ b/preferences.py @@ -118,10 +118,10 @@ class GRABDOC_OT_add_preset(AddPresetBase, Operator): "gd.use_pack_maps", "gd.pack_name", - "gd.channel_R", - "gd.channel_G", - "gd.channel_B", - "gd.channel_A" + "gd.channel_r", + "gd.channel_g", + "gd.channel_b", + "gd.channel_a" ] # Where to store the preset @@ -395,39 +395,22 @@ def update_scale(self, context: Context): ) # Channel packing - # TODO: - # - Implement core functionality - # - Add all properties to presets use_pack_maps: BoolProperty( - name='Enable Packing on Export', - default=False - ) - - pack_name : StringProperty ( - name= 'Pack map name', - default= 'AORM') - - channel_R: EnumProperty( - items=MAP_TYPES[1:], - default="occlusion", - name='R' - ) - channel_G: EnumProperty( - items=MAP_TYPES[1:], - default="roughness", - name='G' - ) - channel_B: EnumProperty( - items=MAP_TYPES[1:], - default="metallic", - name='B' + name="Pack on Export", + description=\ + "After exporting, pack bake maps using the selected packing channels", + default=False ) - channel_A: EnumProperty( - items=MAP_TYPES, - default="none", - name='A' + remove_original_maps: BoolProperty( + name="Remove Original Maps", + description="Remove the original unpacked maps after exporting", + default=False ) - + pack_name: StringProperty(name="Packed Map Name", default="AORM") + channel_r: EnumProperty(items=MAP_TYPES[1:], default="occlusion", name='R') + channel_g: EnumProperty(items=MAP_TYPES[1:], default="roughness", name='G') + channel_b: EnumProperty(items=MAP_TYPES[1:], default="metallic", name='B') + channel_a: EnumProperty(items=MAP_TYPES, default="none", name='A') ################################## diff --git a/ui.py b/ui.py index 095782f..db81560 100644 --- a/ui.py +++ b/ui.py @@ -1,11 +1,7 @@ import os import bpy -from bpy.types import ( - Context, - Panel, - UILayout -) +from bpy.types import Context, Panel, UILayout from .constants import Global from .preferences import GRABDOC_PT_presets @@ -192,7 +188,6 @@ def draw(self, context: Context): image_format = "format" row.prop(gd, image_format) - # TODO: This is insane lol row2 = row.row() if gd.format == "OPEN_EXR": row2.prop(gd, "exr_depth", expand=True) @@ -224,6 +219,9 @@ def draw(self, context: Context): gd, "export_plane", text='Export Plane' ) + col.prop(gd, 'use_pack_maps') + if gd.use_pack_maps: + col.prop(gd, 'remove_original_maps') if gd.baker_type == "marmoset": col.prop( gd, 'marmo_auto_bake', @@ -278,41 +276,29 @@ def draw(self, context: Context): class GRABDOC_PT_pack_maps(PanelInfo, Panel): - bl_label = 'Pack Maps' - bl_parent_id = "GRABDOC_PT_grabdoc" - - @classmethod - def poll(cls, context: Context) -> bool: - return proper_scene_setup() and not context.scene.gd.preview_state - - def draw_header_preset(self, _context: Context): - self.layout.operator( - "grab_doc.pack_maps", - ) - - def draw_header(self, context: Context): - gd = context.scene.gd + bl_label = 'Pack Maps' + bl_parent_id = "GRABDOC_PT_grabdoc" - row = self.layout.row(align=True) - # row.prop(gd, '_use_pack_maps', text='') - row.separator(factor=.5) + @classmethod + def poll(cls, context: Context) -> bool: + return proper_scene_setup() and not context.scene.gd.preview_state - def draw(self, context: Context): - gd = context.scene.gd - + def draw_header_preset(self, _context: Context): + self.layout.operator("grab_doc.pack_maps", icon='IMAGE_DATA') + def draw(self, context: Context): + gd = context.scene.gd - layout = self.layout - layout.use_property_split = True - layout.use_property_decorate = False + layout = self.layout + layout.use_property_split = True + layout.use_property_decorate = False - col = layout.column(align=True) - col.prop(gd, 'use_pack_maps', text="Pack on export") - col.prop(gd, 'channel_R', text="channel R") - col.prop(gd, 'channel_G', text="channel G") - col.prop(gd, 'channel_B', text="channel B") - col.prop(gd, 'channel_A', text="channel A") - col.prop(gd, 'pack_name', text="Suffix") + col = layout.column(align=True) + col.prop(gd, 'channel_r') + col.prop(gd, 'channel_g') + col.prop(gd, 'channel_b') + col.prop(gd, 'channel_a') + col.prop(gd, 'pack_name', text="Suffix") ################################################ @@ -423,7 +409,7 @@ class GRABDOC_PT_metallic(BakerPanel, PanelInfo, Panel): GRABDOC_PT_color, GRABDOC_PT_emissive, GRABDOC_PT_roughness, - GRABDOC_PT_metallic + GRABDOC_PT_metallic ) diff --git a/utils/generic.py b/utils/generic.py index d6f971b..d35b0e4 100644 --- a/utils/generic.py +++ b/utils/generic.py @@ -84,7 +84,7 @@ def is_camera_in_3d_view() -> bool: def format_bl_label( name: str = "GrabDoc", # NOTE: MUST BE CHANGED ALONGSIDE BL_INFO - bl_version: str = (1, 4, 3) + bl_version: str = (1, 4, 4) ) -> str: tuples_version_pattern = r'\((\d+), (\d+), (\d+)\)' match = re.match(tuples_version_pattern, str(bl_version))