From 85e94a3d0de223e0330525fee5bf37d84403ca4f Mon Sep 17 00:00:00 2001 From: Rusfort Date: Wed, 14 May 2025 18:07:28 +0200 Subject: [PATCH] Debugsession N2 - New fallback for LOWRES images --- Documentation/01_User_Guide/02_Features.md | 1 + .../04_Configuration_and_Presets.md | 12 + .../01_User_Guide/09_Output_Structure.md | 1 + .../04_Configuration_System_and_Presets.md | 3 + .../05_Processing_Pipeline.md | 41 ++- autotest.py | 8 +- config/app_settings.json | 5 +- config/file_type_definitions.json | 2 +- configuration.py | 21 +- processing/pipeline/asset_context.py | 6 +- processing/pipeline/orchestrator.py | 331 +++++++++++------- processing/pipeline/stages/initial_scaling.py | 90 +++-- .../stages/metadata_initialization.py | 13 +- .../stages/prepare_processing_items.py | 197 ++++++++--- .../pipeline/stages/regular_map_processor.py | 13 +- processing/pipeline/stages/save_variants.py | 15 +- processing/utils/image_processing_utils.py | 10 + rule_structure.py | 46 ++- 18 files changed, 575 insertions(+), 240 deletions(-) diff --git a/Documentation/01_User_Guide/02_Features.md b/Documentation/01_User_Guide/02_Features.md index 37124e9..34a54cc 100644 --- a/Documentation/01_User_Guide/02_Features.md +++ b/Documentation/01_User_Guide/02_Features.md @@ -16,6 +16,7 @@ This document outlines the key features of the Asset Processor Tool. * Saves maps in appropriate formats (JPG, PNG, EXR) based on complex rules involving map type (`FORCE_LOSSLESS_MAP_TYPES`), resolution (`RESOLUTION_THRESHOLD_FOR_JPG`), bit depth, and source format. * Calculates basic image statistics (Min/Max/Mean) for a reference resolution. * Calculates and stores the relative aspect ratio change string in metadata (e.g., `EVEN`, `X150`, `Y125`). + * **Low-Resolution Fallback:** If enabled (`ENABLE_LOW_RESOLUTION_FALLBACK`), automatically saves an additional "LOWRES" variant of source images if their largest dimension is below a configurable threshold (`LOW_RESOLUTION_THRESHOLD`). This "LOWRES" variant uses the original image dimensions and is saved in addition to any standard resolution outputs. * **Channel Merging:** Combines channels from different maps into packed textures (e.g., NRMRGH) based on preset rules (`MAP_MERGE_RULES` in `config.py`). * **Metadata Generation:** Creates a `metadata.json` file for each asset containing details about maps, category, archetype, aspect ratio change, processing settings, etc. * **Output Organization:** Creates a clean, structured output directory (`///`). diff --git a/Documentation/01_User_Guide/04_Configuration_and_Presets.md b/Documentation/01_User_Guide/04_Configuration_and_Presets.md index 0748067..20a4e2e 100644 --- a/Documentation/01_User_Guide/04_Configuration_and_Presets.md +++ b/Documentation/01_User_Guide/04_Configuration_and_Presets.md @@ -13,6 +13,18 @@ The `app_settings.json` file is structured into several key sections, including: * `ASSET_TYPE_DEFINITIONS`: Defines known asset types (like Surface, Model, Decal) and their properties. * `MAP_MERGE_RULES`: Defines how multiple input maps can be merged into a single output map (e.g., combining Normal and Roughness into one). +### Low-Resolution Fallback Settings + +These settings control the generation of low-resolution "fallback" variants for source images: + +* `ENABLE_LOW_RESOLUTION_FALLBACK` (boolean, default: `true`): + * If `true`, the tool will generate an additional "LOWRES" variant for source images whose largest dimension is smaller than the `LOW_RESOLUTION_THRESHOLD`. + * This "LOWRES" variant uses the original dimensions of the source image and is saved in addition to any other standard resolution outputs (e.g., 1K, PREVIEW). + * If `false`, this feature is disabled. +* `LOW_RESOLUTION_THRESHOLD` (integer, default: `512`): + * Defines the pixel dimension (for the largest side of an image) below which the "LOWRES" fallback variant will be generated (if enabled). + * For example, if set to `512`, any source image smaller than 512x512 (e.g., 256x512, 128x128) will have a "LOWRES" variant created. + ### LLM Predictor Settings For users who wish to utilize the experimental LLM Predictor feature, the following settings are available in `config/llm_settings.json`: diff --git a/Documentation/01_User_Guide/09_Output_Structure.md b/Documentation/01_User_Guide/09_Output_Structure.md index a45a4b1..c64a141 100644 --- a/Documentation/01_User_Guide/09_Output_Structure.md +++ b/Documentation/01_User_Guide/09_Output_Structure.md @@ -58,6 +58,7 @@ The `` (the root folder where processing output starts) i Each asset directory contains the following: * Processed texture maps (e.g., `WoodFloor_Albedo_4k.png`, `MetalPanel_Normal_2k.exr`). The exact filenames depend on the `OUTPUT_FILENAME_PATTERN`. These are the resized, format-converted, and bit-depth adjusted texture files. + * **LOWRES Variants:** If the "Low-Resolution Fallback" feature is enabled and a source image's dimensions are below the configured threshold, an additional variant with "LOWRES" as its resolution token (e.g., `MyTexture_COL_LOWRES.png`) will be saved. This variant uses the original dimensions of the source image. * Merged texture maps (e.g., `WoodFloor_Combined_4k.png`). The exact filenames depend on the `OUTPUT_FILENAME_PATTERN`. These are maps created by combining channels from different source maps based on the configured merge rules. * Model files (if present in the source asset). * `metadata.json`: A JSON file containing detailed information about the asset and the processing that was performed. This includes details about the maps (resolutions, formats, bit depths, and for roughness maps, a `derived_from_gloss_filename: true` flag if it was inverted from an original gloss map), merged map details, calculated image statistics, aspect ratio change information, asset category and archetype, the source preset used, and a list of ignored source files. This file is intended for use by downstream tools or scripts (like the Blender integration scripts). diff --git a/Documentation/02_Developer_Guide/04_Configuration_System_and_Presets.md b/Documentation/02_Developer_Guide/04_Configuration_System_and_Presets.md index 1c1b4f8..7636d7e 100644 --- a/Documentation/02_Developer_Guide/04_Configuration_System_and_Presets.md +++ b/Documentation/02_Developer_Guide/04_Configuration_System_and_Presets.md @@ -12,6 +12,9 @@ The tool's configuration is loaded from several JSON files, providing a layered 1. **Application Settings (`config/app_settings.json`):** This JSON file defines the core global default settings, constants, and rules that apply generally across different asset sources (e.g., the global `OUTPUT_DIRECTORY_PATTERN` and `OUTPUT_FILENAME_PATTERN`, standard image resolutions, map merge rules, output format rules, Blender paths, temporary directory prefix, initial scaling mode, merge dimension mismatch strategy). See the [User Guide: Output Structure](../01_User_Guide/09_Output_Structure.md#available-tokens) for a list of available tokens for these patterns. * *Note:* `ASSET_TYPE_DEFINITIONS` and `FILE_TYPE_DEFINITIONS` are no longer stored here; they have been moved to dedicated files. + * It also includes settings for new features like the "Low-Resolution Fallback": + * `ENABLE_LOW_RESOLUTION_FALLBACK` (boolean): Enables or disables the generation of "LOWRES" variants for small source images. Defaults to `true`. + * `LOW_RESOLUTION_THRESHOLD` (integer): The pixel dimension threshold (largest side) below which a "LOWRES" variant is created if the feature is enabled. Defaults to `512`. 2. **User Settings (`config/user_settings.json`):** This optional JSON file allows users to override specific settings defined in `config/app_settings.json`. If this file exists, its values for corresponding keys will take precedence over the base application settings. This file is primarily managed through the GUI's Application Preferences Editor. diff --git a/Documentation/02_Developer_Guide/05_Processing_Pipeline.md b/Documentation/02_Developer_Guide/05_Processing_Pipeline.md index 4dead92..b5a77b1 100644 --- a/Documentation/02_Developer_Guide/05_Processing_Pipeline.md +++ b/Documentation/02_Developer_Guide/05_Processing_Pipeline.md @@ -50,27 +50,44 @@ These stages are executed sequentially once for each asset before the core item ### Core Item Processing Loop -The [`PipelineOrchestrator`](processing/pipeline/orchestrator.py:36) iterates through the `context.processing_items` list (populated by the [`PrepareProcessingItemsStage`](processing/pipeline/stages/prepare_processing_items.py:10)). For each item (either a [`FileRule`](rule_structure.py:5) for a regular map or a [`MergeTaskDefinition`](processing/pipeline/asset_context.py:16) for a merged map), the following stages are executed sequentially: +The [`PipelineOrchestrator`](processing/pipeline/orchestrator.py:36) iterates through the `context.processing_items` list (populated by the [`PrepareProcessingItemsStage`](processing/pipeline/stages/prepare_processing_items.py:10)). Each `item` in this list is now either a [`ProcessingItem`](rule_structure.py:0) (representing a specific variant of a source map, e.g., Color at 1K, or Color at LOWRES) or a [`MergeTaskDefinition`](processing/pipeline/asset_context.py:16). 1. **[`PrepareProcessingItemsStage`](processing/pipeline/stages/prepare_processing_items.py:10)** (`processing/pipeline/stages/prepare_processing_items.py`): - * **Responsibility**: (Executed once before the loop) Creates the `context.processing_items` list by combining [`FileRule`](rule_structure.py:5)s from `context.files_to_process` and [`MergeTaskDefinition`](processing/pipeline/asset_context.py:16)s derived from the global `map_merge_rules` configuration. It correctly accesses `map_merge_rules` from `context.config_obj` and validates each merge rule for the presence of `output_map_type` and a dictionary for `inputs`. Initializes `context.intermediate_results`. - * **Context Interaction**: Reads from `context.files_to_process` and `context.config_obj` (accessing `map_merge_rules`). Populates `context.processing_items` and initializes `context.intermediate_results`. + * **Responsibility**: (Executed once before the loop) This stage is now responsible for "exploding" each relevant [`FileRule`](rule_structure.py:5) into one or more [`ProcessingItem`](rule_structure.py:0) objects. + * For each [`FileRule`](rule_structure.py:5) that represents an image map: + * It loads the source image data and determines its original dimensions and bit depth. + * It creates standard [`ProcessingItem`](rule_structure.py:0)s for each required output resolution (e.g., "1K", "PREVIEW"), populating them with a copy of the source image data and the respective `resolution_key`. + * If the "Low-Resolution Fallback" feature is enabled (`ENABLE_LOW_RESOLUTION_FALLBACK` in config) and the source image's largest dimension is below `LOW_RESOLUTION_THRESHOLD`, it creates an additional [`ProcessingItem`](rule_structure.py:0) with `resolution_key="LOWRES"`, using the original image data and dimensions. + * It also adds [`MergeTaskDefinition`](processing/pipeline/asset_context.py:16)s derived from global `map_merge_rules`. + * **Context Interaction**: Reads `context.files_to_process` and `context.config_obj`. Populates `context.processing_items` with a list of [`ProcessingItem`](rule_structure.py:0) and [`MergeTaskDefinition`](processing/pipeline/asset_context.py:16) objects. Initializes `context.intermediate_results`. -2. **[`RegularMapProcessorStage`](processing/pipeline/stages/regular_map_processor.py:18)** (`processing/pipeline/stages/regular_map_processor.py`): - * **Responsibility**: (Executed per [`FileRule`](rule_structure.py:5) item) Checks if the `FileRule.item_type` starts with "MAP_". If not, the item is skipped. Otherwise, it loads the image data for the file, determines its potentially suffixed internal map type (e.g., "MAP_COL-1"), applies in-memory transformations (Gloss-to-Rough, Normal Green Invert) using the shared utility function [`apply_common_map_transformations`](processing/utils/image_processing_utils.py), and returns the processed image data and details in a [`ProcessedRegularMapData`](processing/pipeline/asset_context.py:23) object. The `internal_map_type` in the output reflects any transformations (e.g., "MAP_GLOSS" becomes "MAP_ROUGH"). - * **Context Interaction**: Reads from the input [`FileRule`](rule_structure.py:5) (checking `item_type`) and [`Configuration`](configuration.py:68). Returns a [`ProcessedRegularMapData`](processing/pipeline/asset_context.py:23) object which is stored in `context.intermediate_results`. +For each `item` in `context.processing_items`: + +2. **Transformations (Implicit or via a dedicated stage - formerly `RegularMapProcessorStage` logic):** + * **Responsibility**: If the `item` is a [`ProcessingItem`](rule_structure.py:0), its `image_data` (loaded by `PrepareProcessingItemsStage`) may need transformations (Gloss-to-Rough, Normal Green Invert). This logic, previously in `RegularMapProcessorStage`, might be integrated into `PrepareProcessingItemsStage` before `ProcessingItem` creation, or handled by a new dedicated transformation stage that operates on `ProcessingItem.image_data`. The `item.map_type_identifier` would be updated if a transformation like Gloss-to-Rough occurs. + * **Context Interaction**: Modifies `item.image_data` and `item.map_type_identifier` within the [`ProcessingItem`](rule_structure.py:0) object. 3. **[`MergedTaskProcessorStage`](processing/pipeline/stages/merged_task_processor.py:68)** (`processing/pipeline/stages/merged_task_processor.py`): - * **Responsibility**: (Executed per [`MergeTaskDefinition`](processing/pipeline/asset_context.py:16) item) Validates that all input map types specified in the merge rule start with "MAP_". If not, the task is failed. It dynamically loads input images by looking up the required input map types (e.g., "MAP_NRM") in `context.processed_maps_details` and using the temporary file paths from their `saved_files_info`. It applies in-memory transformations to inputs using [`apply_common_map_transformations`](processing/utils/image_processing_utils.py), handles dimension mismatches (with fallback creation if configured and `source_dimensions` are available), performs the channel merging operation, and returns the merged image data and details in a [`ProcessedMergedMapData`](processing/pipeline/asset_context.py:35) object. The `output_map_type` of the merged map must also be "MAP_" prefixed in the configuration. - * **Context Interaction**: Reads from the input [`MergeTaskDefinition`](processing/pipeline/asset_context.py:16) (checking input map types), `context.workspace_path`, `context.processed_maps_details` (for input image data), and [`Configuration`](configuration.py:68). Returns a [`ProcessedMergedMapData`](processing/pipeline/asset_context.py:35) object which is stored in `context.intermediate_results`. + * **Responsibility**: (Executed if `item` is a [`MergeTaskDefinition`](processing/pipeline/asset_context.py:16)) Same as before: validates inputs, loads source map data (likely from `ProcessingItem`s in `context.processing_items` or a cache populated from them), applies transformations, merges channels, and returns [`ProcessedMergedMapData`](processing/pipeline/asset_context.py:35). + * **Context Interaction**: Reads [`MergeTaskDefinition`](processing/pipeline/asset_context.py:16), potentially `context.processing_items` (or a cache derived from it) for input image data. Returns [`ProcessedMergedMapData`](processing/pipeline/asset_context.py:35). 4. **[`InitialScalingStage`](processing/pipeline/stages/initial_scaling.py:14)** (`processing/pipeline/stages/initial_scaling.py`): - * **Responsibility**: (Executed per item) Applies initial scaling (e.g., Power-of-Two downscaling) to the image data from the previous processing stage based on the `initial_scaling_mode` configuration. - * **Context Interaction**: Takes a [`InitialScalingInput`](processing/pipeline/asset_context.py:46) (containing image data and config) and returns an [`InitialScalingOutput`](processing/pipeline/asset_context.py:54) object, which updates the item's entry in `context.intermediate_results`. + * **Responsibility**: (Executed per item) + * If `item` is a [`ProcessingItem`](rule_structure.py:0): Takes `item.image_data`, `item.current_dimensions`, and `item.resolution_key` as input. If `item.resolution_key` is "LOWRES", POT scaling is skipped. Otherwise, applies POT scaling if configured. + * If `item` is from a `MergeTaskDefinition` (i.e., `processed_data` from `MergedTaskProcessorStage`): Applies POT scaling as before. + * **Context Interaction**: Takes [`InitialScalingInput`](processing/pipeline/asset_context.py:46) (now including `resolution_key`). Returns [`InitialScalingOutput`](processing/pipeline/asset_context.py:54) (also including `resolution_key`), which updates `context.intermediate_results`. The `current_image_data` and `current_dimensions` for saving are taken from this output. 5. **[`SaveVariantsStage`](processing/pipeline/stages/save_variants.py:15)** (`processing/pipeline/stages/save_variants.py`): - * **Responsibility**: (Executed per item) Takes the final processed image data (potentially scaled) and configuration, and calls a utility to save the image to temporary files in various resolutions and formats as defined by the configuration. - * **Context Interaction**: Takes a [`SaveVariantsInput`](processing/pipeline/asset_context.py:61) object (which includes the "MAP_" prefixed `internal_map_type`). It uses the `get_filename_friendly_map_type` utility to convert this to a "standard type" (e.g., "COL") for output naming. Returns a [`SaveVariantsOutput`](processing/pipeline/asset_context.py:79) object containing details about the saved temporary files. The orchestrator stores these details, including the original "MAP_" prefixed `internal_map_type`, in `context.processed_maps_details` for the item. + * **Responsibility**: (Executed per item) Saves the (potentially scaled) `current_image_data`. + * **Context Interaction**: + * Takes [`SaveVariantsInput`](processing/pipeline/asset_context.py:61). + * `internal_map_type` is set from `item.map_type_identifier` (for `ProcessingItem`) or `processed_data.output_map_type` (for merged). + * `output_filename_pattern_tokens['resolution']` is set to the `resolution_key` obtained from `scaled_data_output.resolution_key` (which originates from `item.resolution_key` for `ProcessingItem`s, or is `None` for merged items that get all standard resolutions). + * `image_resolutions` argument for `SaveVariantsInput`: + * If `resolution_key == "LOWRES"`: Set to `{"LOWRES": width_of_lowres_data}`. + * If `resolution_key` is a standard key (e.g., "1K"): Set to `{resolution_key: configured_dimension}`. + * For merged items (where `resolution_key` from scaling is likely `None`): Set to the full `config.image_resolutions` map to generate all applicable standard sizes. + * Returns [`SaveVariantsOutput`](processing/pipeline/asset_context.py:79). Orchestrator stores details in `context.processed_maps_details`. ### Post-Item Stages diff --git a/autotest.py b/autotest.py index 1782ef2..96f147e 100644 --- a/autotest.py +++ b/autotest.py @@ -96,6 +96,8 @@ class InfoSummaryFilter(logging.Filter): "verify: processingengine.process called", ": effective supplier set to", ": metadata initialized.", + "path", + "\\asset_processor", ": file rules queued for processing", "successfully loaded base application settings", "successfully loaded and merged asset_type_definitions", @@ -108,12 +110,6 @@ class InfoSummaryFilter(logging.Filter): "worker thread: finished processing for rule:", "task finished signal received for", # Autotest step markers (not global summaries) - "step 1: loading zip file:", - "step 2: selecting preset:", - "step 4: retrieving and comparing rules...", - "step 5: starting processing...", - "step 7: checking output path:", - "output path check completed.", ] def filter(self, record): diff --git a/config/app_settings.json b/config/app_settings.json index 1343766..dbba665 100644 --- a/config/app_settings.json +++ b/config/app_settings.json @@ -46,7 +46,10 @@ "TEMP_DIR_PREFIX": "_PROCESS_ASSET_", "INITIAL_SCALING_MODE": "POT_DOWNSCALE", "MERGE_DIMENSION_MISMATCH_STRATEGY": "USE_LARGEST", + "ENABLE_LOW_RESOLUTION_FALLBACK": true, + "LOW_RESOLUTION_THRESHOLD": 512, "general_settings": { - "invert_normal_map_green_channel_globally": false + "invert_normal_map_green_channel_globally": false, + "app_version": "Pre-Alpha" } } \ No newline at end of file diff --git a/config/file_type_definitions.json b/config/file_type_definitions.json index 018a0d0..b8aac76 100644 --- a/config/file_type_definitions.json +++ b/config/file_type_definitions.json @@ -190,7 +190,7 @@ ], "is_grayscale": false, "keybind": "E", - "standard_type": "" + "standard_type": "EXTRA" }, "FILE_IGNORE": { "bit_depth_rule": "", diff --git a/configuration.py b/configuration.py index 6a4bee1..b494d25 100644 --- a/configuration.py +++ b/configuration.py @@ -4,6 +4,7 @@ from pathlib import Path import logging import re import collections.abc +from typing import Optional log = logging.getLogger(__name__) @@ -12,7 +13,7 @@ APP_SETTINGS_PATH = BASE_DIR / "config" / "app_settings.json" LLM_SETTINGS_PATH = BASE_DIR / "config" / "llm_settings.json" ASSET_TYPE_DEFINITIONS_PATH = BASE_DIR / "config" / "asset_type_definitions.json" FILE_TYPE_DEFINITIONS_PATH = BASE_DIR / "config" / "file_type_definitions.json" -USER_SETTINGS_PATH = BASE_DIR / "config" / "user_settings.json" # New path for user settings +USER_SETTINGS_PATH = BASE_DIR / "config" / "user_settings.json" SUPPLIERS_CONFIG_PATH = BASE_DIR / "config" / "suppliers.json" PRESETS_DIR = BASE_DIR / "Presets" @@ -649,6 +650,24 @@ class Configuration: """Returns the LLM request timeout in seconds from LLM settings.""" return self._llm_settings.get('llm_request_timeout', 120) + @property + def app_version(self) -> Optional[str]: + """Returns the application version from general_settings.""" + gs = self._core_settings.get('general_settings') + if isinstance(gs, dict): + return gs.get('app_version') + return None + + @property + def enable_low_resolution_fallback(self) -> bool: + """Gets the setting for enabling low-resolution fallback.""" + return self._core_settings.get('ENABLE_LOW_RESOLUTION_FALLBACK', True) + + @property + def low_resolution_threshold(self) -> int: + """Gets the pixel dimension threshold for low-resolution fallback.""" + return self._core_settings.get('LOW_RESOLUTION_THRESHOLD', 512) + @property def FILE_TYPE_DEFINITIONS(self) -> dict: return self._file_type_definitions diff --git a/processing/pipeline/asset_context.py b/processing/pipeline/asset_context.py index f6363e5..42c3d28 100644 --- a/processing/pipeline/asset_context.py +++ b/processing/pipeline/asset_context.py @@ -1,3 +1,4 @@ +import dataclasses # Added import from dataclasses import dataclass from pathlib import Path from typing import Dict, List, Optional @@ -27,6 +28,7 @@ class ProcessedRegularMapData: original_bit_depth: Optional[int] original_dimensions: Optional[Tuple[int, int]] # (width, height) transformations_applied: List[str] + resolution_key: Optional[str] = None # Added field status: str = "Processed" error_message: Optional[str] = None @@ -45,9 +47,10 @@ class ProcessedMergedMapData: @dataclass class InitialScalingInput: image_data: np.ndarray + initial_scaling_mode: str # Moved before fields with defaults original_dimensions: Optional[Tuple[int, int]] # (width, height) + resolution_key: Optional[str] = None # Added field # Configuration needed - initial_scaling_mode: str # Output for InitialScalingStage @dataclass @@ -55,6 +58,7 @@ class InitialScalingOutput: scaled_image_data: np.ndarray scaling_applied: bool final_dimensions: Tuple[int, int] # (width, height) + resolution_key: Optional[str] = None # Added field # Input for SaveVariantsStage @dataclass diff --git a/processing/pipeline/orchestrator.py b/processing/pipeline/orchestrator.py index 6c8fe7a..2d89408 100644 --- a/processing/pipeline/orchestrator.py +++ b/processing/pipeline/orchestrator.py @@ -8,7 +8,7 @@ from typing import List, Dict, Optional, Any, Union # Added Any, Union import numpy as np # Added numpy from configuration import Configuration -from rule_structure import SourceRule, AssetRule, FileRule # Added FileRule +from rule_structure import SourceRule, AssetRule, FileRule, ProcessingItem # Added ProcessingItem # Import new context classes and stages from .asset_context import ( @@ -200,145 +200,224 @@ class PipelineOrchestrator: current_image_data: Optional[np.ndarray] = None # Track current image data ref try: - # 1. Process (Load/Merge + Transform) - if isinstance(item, FileRule): - if item.item_type == 'EXTRA': - log.debug(f"{item_log_prefix}: Skipping image processing for EXTRA FileRule '{item.file_path}'.") - # Add a basic entry to processed_maps_details to acknowledge it was seen - context.processed_maps_details[item.file_path] = { - "status": "Skipped (EXTRA file)", - "internal_map_type": "EXTRA", - "source_file": str(item.file_path) - } - continue # Skip to the next item - item_key = item.file_path # Use file_path string as key - log.debug(f"{item_log_prefix}: Processing FileRule '{item.file_path}'...") - processed_data = self._regular_processor_stage.execute(context, item) - elif isinstance(item, MergeTaskDefinition): - item_key = item.task_key # Use task_key string as key - log.info(f"{item_log_prefix}: Executing MergedTaskProcessorStage for MergeTask '{item_key}'...") # Log call - processed_data = self._merged_processor_stage.execute(context, item) - # Log status/error from merge processor - if processed_data: - log.info(f"{item_log_prefix}: MergedTaskProcessorStage result - Status: {processed_data.status}, Error: {processed_data.error_message}") - else: - log.warning(f"{item_log_prefix}: MergedTaskProcessorStage returned None for MergeTask '{item_key}'.") - else: - log.warning(f"{item_log_prefix}: Unknown item type '{type(item)}'. Skipping.") - item_key = f"unknown_item_{item_index}" - context.processed_maps_details[item_key] = {"status": "Skipped", "notes": f"Unknown item type {type(item)}"} - asset_had_item_errors = True - continue # Next item + # The 'item' is now expected to be a ProcessingItem or MergeTaskDefinition + + if isinstance(item, ProcessingItem): + item_key = f"{item.source_file_info_ref}_{item.map_type_identifier}_{item.resolution_key}" + item_log_prefix = f"Asset '{asset_name}', ProcItem '{item_key}'" + log.info(f"{item_log_prefix}: Starting processing.") - # Check for processing failure - if not processed_data or processed_data.status != "Processed": - error_msg = processed_data.error_message if processed_data else "Processor returned None" - log.error(f"{item_log_prefix}: Failed during processing stage. Error: {error_msg}") - context.processed_maps_details[item_key] = {"status": "Failed", "notes": f"Processing Error: {error_msg}", "stage": processed_data.__class__.__name__ if processed_data else "UnknownProcessor"} - asset_had_item_errors = True - continue # Next item + # Data for ProcessingItem is already loaded by PrepareProcessingItemsStage + current_image_data = item.image_data + current_dimensions = item.current_dimensions + item_resolution_key = item.resolution_key + + # Transformations (like gloss to rough, normal invert) are assumed to be applied + # by RegularMapProcessorStage if it's still used, or directly in PrepareProcessingItemsStage + # before creating the ProcessingItem, or a new dedicated transformation stage. + # For now, assume item.image_data is ready for scaling/saving. + + # Store initial ProcessingItem data as "processed_data" for consistency if RegularMapProcessor is bypassed + # This is a simplification; a dedicated transformation stage would be cleaner. + # For now, we assume transformations happened before or within PrepareProcessingItemsStage. + # The 'processed_data' variable here is more of a placeholder for what would feed into scaling. + + # Create a simple ProcessedRegularMapData-like structure for logging/details if needed, + # or adapt the final_details population later. + # For now, we'll directly use 'item' fields. - # Store intermediate result & get current image data - context.intermediate_results[item_key] = processed_data - current_image_data = processed_data.processed_image_data if isinstance(processed_data, ProcessedRegularMapData) else processed_data.merged_image_data - current_dimensions = processed_data.original_dimensions if isinstance(processed_data, ProcessedRegularMapData) else processed_data.final_dimensions - - # 2. Scale (Optional) - scaling_mode = getattr(context.config_obj, "INITIAL_SCALING_MODE", "NONE") - if scaling_mode != "NONE" and current_image_data is not None and current_image_data.size > 0: - if isinstance(item, MergeTaskDefinition): # Log scaling call for merge tasks - log.info(f"{item_log_prefix}: Calling InitialScalingStage for MergeTask '{item_key}' (Mode: {scaling_mode})...") - log.debug(f"{item_log_prefix}: Applying initial scaling (Mode: {scaling_mode})...") + # 2. Scale (Optional) + scaling_mode = getattr(context.config_obj, "INITIAL_SCALING_MODE", "NONE") + # Pass the item's resolution_key to InitialScalingInput scale_input = InitialScalingInput( image_data=current_image_data, - original_dimensions=current_dimensions, # Pass original/merged dims - initial_scaling_mode=scaling_mode + original_dimensions=current_dimensions, + initial_scaling_mode=scaling_mode, + resolution_key=item_resolution_key # Pass the key ) + # Add _source_file_path for logging within InitialScalingStage if available + setattr(scale_input, '_source_file_path', item.source_file_info_ref) + + log.debug(f"{item_log_prefix}: Calling InitialScalingStage. Input res_key: {scale_input.resolution_key}") scaled_data_output = self._scaling_stage.execute(scale_input) - # Update intermediate result and current image data reference - context.intermediate_results[item_key] = scaled_data_output # Overwrite previous intermediate - current_image_data = scaled_data_output.scaled_image_data # Use scaled data for saving - log.debug(f"{item_log_prefix}: Scaling applied: {scaled_data_output.scaling_applied}. New Dims: {scaled_data_output.final_dimensions}") - else: - log.debug(f"{item_log_prefix}: Initial scaling skipped (Mode: NONE or empty image).") - # Create dummy output if scaling skipped, using current dims - final_dims = current_dimensions if current_dimensions else (current_image_data.shape[1], current_image_data.shape[0]) if current_image_data is not None else (0,0) - scaled_data_output = InitialScalingOutput(scaled_image_data=current_image_data, scaling_applied=False, final_dimensions=final_dims) + current_image_data = scaled_data_output.scaled_image_data + current_dimensions = scaled_data_output.final_dimensions # Dimensions after scaling + # The resolution_key from item is passed through by InitialScalingOutput + output_resolution_key = scaled_data_output.resolution_key + log.debug(f"{item_log_prefix}: InitialScalingStage output. Scaled: {scaled_data_output.scaling_applied}, New Dims: {current_dimensions}, Output ResKey: {output_resolution_key}") + context.intermediate_results[item_key] = scaled_data_output - # 3. Save Variants - if current_image_data is None or current_image_data.size == 0: - log.warning(f"{item_log_prefix}: Skipping save stage because image data is empty.") - context.processed_maps_details[item_key] = {"status": "Skipped", "notes": "No image data to save", "stage": "SaveVariantsStage"} - # Don't mark as asset error, just skip this item's saving - continue # Next item + # 3. Save Variants + if current_image_data is None or current_image_data.size == 0: + log.warning(f"{item_log_prefix}: Skipping save stage because image data is empty.") + context.processed_maps_details[item_key] = {"status": "Skipped", "notes": "No image data to save", "stage": "SaveVariantsStage"} + continue - if isinstance(item, MergeTaskDefinition): # Log save call for merge tasks - log.info(f"{item_log_prefix}: Calling SaveVariantsStage for MergeTask '{item_key}'...") - log.debug(f"{item_log_prefix}: Saving variants...") - # Prepare input for save stage - internal_map_type = processed_data.final_internal_map_type if isinstance(processed_data, ProcessedRegularMapData) else processed_data.output_map_type - source_bit_depth = [processed_data.original_bit_depth] if isinstance(processed_data, ProcessedRegularMapData) and processed_data.original_bit_depth is not None else processed_data.source_bit_depths if isinstance(processed_data, ProcessedMergedMapData) else [8] # Default bit depth if unknown - - # Construct filename tokens (ensure temp dir is used) - output_filename_tokens = { - 'asset_name': asset_name, - 'output_base_directory': context.engine_temp_dir, # Save variants to temp dir - # Add other tokens from context/config as needed by the pattern - 'supplier': context.effective_supplier or 'UnknownSupplier', - } - - # Log the value being read for the threshold before creating the input object - log.info(f"ORCHESTRATOR_DEBUG: Reading RESOLUTION_THRESHOLD_FOR_JPG from config for SaveVariantsInput: {getattr(context.config_obj, 'RESOLUTION_THRESHOLD_FOR_JPG', None)}") - save_input = SaveVariantsInput( - image_data=current_image_data, # Use potentially scaled data - internal_map_type=internal_map_type, - source_bit_depth_info=source_bit_depth, - output_filename_pattern_tokens=output_filename_tokens, - # Pass config values needed by save stage - image_resolutions=context.config_obj.image_resolutions, - file_type_defs=getattr(context.config_obj, "FILE_TYPE_DEFINITIONS", {}), - output_format_8bit=context.config_obj.get_8bit_output_format(), - output_format_16bit_primary=context.config_obj.get_16bit_output_formats()[0], - output_format_16bit_fallback=context.config_obj.get_16bit_output_formats()[1], - png_compression_level=context.config_obj.png_compression_level, - jpg_quality=context.config_obj.jpg_quality, - output_filename_pattern=context.config_obj.output_filename_pattern, - resolution_threshold_for_jpg=getattr(context.config_obj, "resolution_threshold_for_jpg", None) # Corrected case - ) - saved_data = self._save_stage.execute(save_input) - # Log saved_data for merge tasks - if isinstance(item, MergeTaskDefinition): - log.info(f"{item_log_prefix}: SaveVariantsStage result for MergeTask '{item_key}' - Status: {saved_data.status if saved_data else 'N/A'}, Saved Files: {len(saved_data.saved_files_details) if saved_data else 0}") - - # Check save status and finalize item result - if saved_data and saved_data.status.startswith("Processed"): - item_status = saved_data.status # e.g., "Processed" or "Processed (No Output)" - log.info(f"{item_log_prefix}: Item successfully processed and saved. Status: {item_status}") - # Populate final details for this item - final_details = { - "status": item_status, - "saved_files_info": saved_data.saved_files_details, # List of dicts from save util - "internal_map_type": internal_map_type, - "original_dimensions": processed_data.original_dimensions if isinstance(processed_data, ProcessedRegularMapData) else None, - "final_dimensions": scaled_data_output.final_dimensions if scaled_data_output else current_dimensions, - "transformations": processed_data.transformations_applied if isinstance(processed_data, ProcessedRegularMapData) else processed_data.transformations_applied_to_inputs, - # Add source file if regular map - "source_file": str(processed_data.source_file_path) if isinstance(processed_data, ProcessedRegularMapData) else None, + log.debug(f"{item_log_prefix}: Preparing to save variant with resolution key '{output_resolution_key}'...") + + output_filename_tokens = { + 'asset_name': asset_name, + 'output_base_directory': context.engine_temp_dir, + 'supplier': context.effective_supplier or 'UnknownSupplier', + 'resolution': output_resolution_key # Use the key from the item/scaling stage } - # Log final details addition for merge tasks - if isinstance(item, MergeTaskDefinition): - log.info(f"{item_log_prefix}: Adding final details to context.processed_maps_details for MergeTask '{item_key}'. Details: {final_details}") - context.processed_maps_details[item_key] = final_details + + # Determine image_resolutions argument for save_image_variants + save_specific_resolutions = {} + if output_resolution_key == "LOWRES": + # For LOWRES, the "resolution value" is its actual dimension. + # image_saving_utils needs a dict like {"LOWRES": 64} if current_dim is 64x64 + # Assuming current_dimensions[0] is width. + save_specific_resolutions = {"LOWRES": current_dimensions[0] if current_dimensions else 0} + log.debug(f"{item_log_prefix}: Preparing to save LOWRES variant. Dimensions: {current_dimensions}. Save resolutions arg: {save_specific_resolutions}") + elif output_resolution_key in context.config_obj.image_resolutions: + save_specific_resolutions = {output_resolution_key: context.config_obj.image_resolutions[output_resolution_key]} + else: + log.warning(f"{item_log_prefix}: Resolution key '{output_resolution_key}' not found in config.image_resolutions and not LOWRES. Saving might fail or use full res.") + # Fallback: pass all configured resolutions, image_saving_utils will try to match by size. + # This might not be ideal if the key is truly unknown. + # Or, more strictly, fail here if key is unknown and not LOWRES. + # For now, let image_saving_utils handle it by passing all. + save_specific_resolutions = context.config_obj.image_resolutions + + + save_input = SaveVariantsInput( + image_data=current_image_data, + internal_map_type=item.map_type_identifier, + source_bit_depth_info=[item.bit_depth] if item.bit_depth is not None else [8], # Default to 8 if not set + output_filename_pattern_tokens=output_filename_tokens, + image_resolutions=save_specific_resolutions, # Pass the specific resolution(s) + file_type_defs=getattr(context.config_obj, "FILE_TYPE_DEFINITIONS", {}), + output_format_8bit=context.config_obj.get_8bit_output_format(), + output_format_16bit_primary=context.config_obj.get_16bit_output_formats()[0], + output_format_16bit_fallback=context.config_obj.get_16bit_output_formats()[1], + png_compression_level=context.config_obj.png_compression_level, + jpg_quality=context.config_obj.jpg_quality, + output_filename_pattern=context.config_obj.output_filename_pattern, + resolution_threshold_for_jpg=getattr(context.config_obj, "resolution_threshold_for_jpg", None) + ) + saved_data = self._save_stage.execute(save_input) + + if saved_data and saved_data.status.startswith("Processed"): + item_status = saved_data.status + log.info(f"{item_log_prefix}: Item successfully processed and saved. Status: {item_status}") + context.processed_maps_details[item_key] = { + "status": item_status, + "saved_files_info": saved_data.saved_files_details, + "internal_map_type": item.map_type_identifier, + "resolution_key": output_resolution_key, + "original_dimensions": item.original_dimensions, + "final_dimensions": current_dimensions, # Dimensions after scaling + "source_file": item.source_file_info_ref, + } + else: + error_msg = saved_data.error_message if saved_data else "Save stage returned None" + log.error(f"{item_log_prefix}: Failed during save stage. Error: {error_msg}") + context.processed_maps_details[item_key] = {"status": "Failed", "notes": f"Save Error: {error_msg}", "stage": "SaveVariantsStage"} + asset_had_item_errors = True + item_status = "Failed" + + elif isinstance(item, MergeTaskDefinition): + # --- This part needs similar refactoring for resolution_key if merged outputs can be LOWRES --- + # --- For now, assume merged tasks always produce standard resolutions --- + item_key = item.task_key + item_log_prefix = f"Asset '{asset_name}', MergeTask '{item_key}'" + log.info(f"{item_log_prefix}: Processing MergeTask.") + + # 1. Process Merge Task + processed_data = self._merged_processor_stage.execute(context, item) + if not processed_data or processed_data.status != "Processed": + error_msg = processed_data.error_message if processed_data else "Merge processor returned None" + log.error(f"{item_log_prefix}: Failed during merge processing. Error: {error_msg}") + context.processed_maps_details[item_key] = {"status": "Failed", "notes": f"Merge Error: {error_msg}", "stage": "MergedTaskProcessorStage"} + asset_had_item_errors = True + continue + + context.intermediate_results[item_key] = processed_data + current_image_data = processed_data.merged_image_data + current_dimensions = processed_data.final_dimensions + + # 2. Scale Merged Output (Optional) + # Merged tasks typically don't have a single "resolution_key" like LOWRES from source. + # They produce an image that then gets downscaled to 1K, PREVIEW etc. + # So, resolution_key for InitialScalingInput here would be None or a default. + scaling_mode = getattr(context.config_obj, "INITIAL_SCALING_MODE", "NONE") + scale_input = InitialScalingInput( + image_data=current_image_data, + original_dimensions=current_dimensions, + initial_scaling_mode=scaling_mode, + resolution_key=None # Merged outputs are not "LOWRES" themselves before this scaling + ) + setattr(scale_input, '_source_file_path', f"MergeTask_{item_key}") # For logging + + log.debug(f"{item_log_prefix}: Calling InitialScalingStage for merged data.") + scaled_data_output = self._scaling_stage.execute(scale_input) + current_image_data = scaled_data_output.scaled_image_data + current_dimensions = scaled_data_output.final_dimensions + # Merged items don't have a specific output_resolution_key from source, + # they will be saved to all applicable resolutions from config. + # So scaled_data_output.resolution_key will be None here. + context.intermediate_results[item_key] = scaled_data_output + + # 3. Save Merged Variants + if current_image_data is None or current_image_data.size == 0: + log.warning(f"{item_log_prefix}: Skipping save for merged task, image data is empty.") + context.processed_maps_details[item_key] = {"status": "Skipped", "notes": "No merged image data to save", "stage": "SaveVariantsStage"} + continue + + output_filename_tokens = { + 'asset_name': asset_name, + 'output_base_directory': context.engine_temp_dir, + 'supplier': context.effective_supplier or 'UnknownSupplier', + # 'resolution' token will be filled by image_saving_utils for each variant + } + + # For merged tasks, we usually want to generate all standard resolutions. + # The `resolution_key` from the item itself is not applicable here for the `resolution` token. + # The `image_saving_utils.save_image_variants` will iterate through `context.config_obj.image_resolutions`. + save_input = SaveVariantsInput( + image_data=current_image_data, + internal_map_type=processed_data.output_map_type, + source_bit_depth_info=processed_data.source_bit_depths, + output_filename_pattern_tokens=output_filename_tokens, + image_resolutions=context.config_obj.image_resolutions, # Pass all configured resolutions + file_type_defs=getattr(context.config_obj, "FILE_TYPE_DEFINITIONS", {}), + output_format_8bit=context.config_obj.get_8bit_output_format(), + output_format_16bit_primary=context.config_obj.get_16bit_output_formats()[0], + output_format_16bit_fallback=context.config_obj.get_16bit_output_formats()[1], + png_compression_level=context.config_obj.png_compression_level, + jpg_quality=context.config_obj.jpg_quality, + output_filename_pattern=context.config_obj.output_filename_pattern, + resolution_threshold_for_jpg=getattr(context.config_obj, "resolution_threshold_for_jpg", None) + ) + saved_data = self._save_stage.execute(save_input) + + if saved_data and saved_data.status.startswith("Processed"): + item_status = saved_data.status + log.info(f"{item_log_prefix}: Merged task successfully processed and saved. Status: {item_status}") + context.processed_maps_details[item_key] = { + "status": item_status, + "saved_files_info": saved_data.saved_files_details, + "internal_map_type": processed_data.output_map_type, + "final_dimensions": current_dimensions, + } + else: + error_msg = saved_data.error_message if saved_data else "Save stage for merged task returned None" + log.error(f"{item_log_prefix}: Failed during save stage for merged task. Error: {error_msg}") + context.processed_maps_details[item_key] = {"status": "Failed", "notes": f"Save Error (Merged): {error_msg}", "stage": "SaveVariantsStage"} + asset_had_item_errors = True + item_status = "Failed" else: - error_msg = saved_data.error_message if saved_data else "Save stage returned None" - log.error(f"{item_log_prefix}: Failed during save stage. Error: {error_msg}") - context.processed_maps_details[item_key] = {"status": "Failed", "notes": f"Save Error: {error_msg}", "stage": "SaveVariantsStage"} + log.warning(f"{item_log_prefix}: Unknown item type in loop: {type(item)}. Skipping.") + # Ensure some key exists to prevent KeyError if item_key was not set + unknown_item_key = f"unknown_item_at_index_{item_index}" + context.processed_maps_details[unknown_item_key] = {"status": "Skipped", "notes": f"Unknown item type {type(item)}"} asset_had_item_errors = True - item_status = "Failed" # Ensure item status reflects failure + continue except Exception as e: - log.exception(f"{item_log_prefix}: Unhandled exception during item processing loop: {e}") + log.exception(f"Asset '{asset_name}', Item Loop Index {item_index}: Unhandled exception: {e}") # Ensure details are recorded even on unhandled exception if item_key is not None: context.processed_maps_details[item_key] = {"status": "Failed", "notes": f"Unhandled Loop Error: {e}", "stage": "OrchestratorLoop"} diff --git a/processing/pipeline/stages/initial_scaling.py b/processing/pipeline/stages/initial_scaling.py index 6fc27ab..9d88c5a 100644 --- a/processing/pipeline/stages/initial_scaling.py +++ b/processing/pipeline/stages/initial_scaling.py @@ -1,5 +1,5 @@ import logging -from typing import Tuple +from typing import Tuple, Optional # Added Optional import cv2 # Assuming cv2 is available for interpolation flags import numpy as np @@ -7,77 +7,93 @@ import numpy as np from .base_stage import ProcessingStage # Import necessary context classes and utils from ..asset_context import InitialScalingInput, InitialScalingOutput +# ProcessingItem is no longer created here, so its import can be removed if not used otherwise. +# For now, keep rule_structure import if other elements from it might be needed, +# but ProcessingItem itself is not directly instantiated by this stage anymore. +# from rule_structure import ProcessingItem from ...utils import image_processing_utils as ipu +import numpy as np +import cv2 # Added cv2 for interpolation flags (already used implicitly by ipu.resize_image) log = logging.getLogger(__name__) class InitialScalingStage(ProcessingStage): """ - Applies initial scaling (e.g., Power-of-Two downscaling) to image data - if configured via the InitialScalingInput. + Applies initial Power-of-Two (POT) downscaling to image data if configured + and if the item is not already a 'LOWRES' variant. """ def execute(self, input_data: InitialScalingInput) -> InitialScalingOutput: """ - Applies scaling based on input_data.initial_scaling_mode. + Applies POT scaling based on input_data.initial_scaling_mode, + unless input_data.resolution_key is 'LOWRES'. + Passes through the resolution_key. """ - log.debug(f"Initial Scaling Stage: Mode '{input_data.initial_scaling_mode}'.") + # Safely access source_file_path for logging, if provided by orchestrator via underscore attribute + source_file_path = getattr(input_data, '_source_file_path', "UnknownSourcePath") + log_prefix = f"InitialScalingStage (Source: {source_file_path}, ResKey: {input_data.resolution_key})" + + log.debug(f"{log_prefix}: Mode '{input_data.initial_scaling_mode}'. Received resolution_key: '{input_data.resolution_key}'") image_to_scale = input_data.image_data - original_dims_wh = input_data.original_dimensions + current_dimensions_wh = input_data.original_dimensions # Dimensions of the image_to_scale scaling_mode = input_data.initial_scaling_mode - scaling_applied = False - final_image_data = image_to_scale # Default to original if no scaling happens + + output_resolution_key = input_data.resolution_key # Pass through the resolution key if image_to_scale is None or image_to_scale.size == 0: - log.warning("Initial Scaling Stage: Input image data is None or empty. Skipping.") - # Return original (empty) data and indicate no scaling + log.warning(f"{log_prefix}: Input image data is None or empty. Skipping POT scaling.") return InitialScalingOutput( scaled_image_data=np.array([]), scaling_applied=False, - final_dimensions=(0, 0) + final_dimensions=(0, 0), + resolution_key=output_resolution_key ) - if original_dims_wh is None: - log.warning("Initial Scaling Stage: Original dimensions not provided. Using current image shape.") - h_pre_scale, w_pre_scale = image_to_scale.shape[:2] - original_dims_wh = (w_pre_scale, h_pre_scale) + if not current_dimensions_wh: + log.warning(f"{log_prefix}: Original dimensions not provided for POT scaling. Using current image shape.") + h_pre_pot_scale, w_pre_pot_scale = image_to_scale.shape[:2] else: - w_pre_scale, h_pre_scale = original_dims_wh + w_pre_pot_scale, h_pre_pot_scale = current_dimensions_wh + + final_image_data = image_to_scale # Default to original if no scaling happens + scaling_applied = False + # Skip POT scaling if the item is already a LOWRES variant or scaling mode is NONE + if output_resolution_key == "LOWRES": + log.info(f"{log_prefix}: Item is a 'LOWRES' variant. Skipping POT downscaling.") + elif scaling_mode == "NONE": + log.info(f"{log_prefix}: Mode is NONE. No POT scaling applied.") + elif scaling_mode == "POT_DOWNSCALE": + pot_w = ipu.get_nearest_power_of_two_downscale(w_pre_pot_scale) + pot_h = ipu.get_nearest_power_of_two_downscale(h_pre_pot_scale) - if scaling_mode == "POT_DOWNSCALE": - pot_w = ipu.get_nearest_power_of_two_downscale(w_pre_scale) - pot_h = ipu.get_nearest_power_of_two_downscale(h_pre_scale) - - if (pot_w, pot_h) != (w_pre_scale, h_pre_scale): - log.info(f"Initial Scaling: Applying POT Downscale from ({w_pre_scale},{h_pre_scale}) to ({pot_w},{pot_h}).") - # Use INTER_AREA for downscaling generally + if (pot_w, pot_h) != (w_pre_pot_scale, h_pre_pot_scale): + log.info(f"{log_prefix}: Applying POT Downscale from ({w_pre_pot_scale},{h_pre_pot_scale}) to ({pot_w},{pot_h}).") resized_img = ipu.resize_image(image_to_scale, pot_w, pot_h, interpolation=cv2.INTER_AREA) if resized_img is not None: final_image_data = resized_img scaling_applied = True - log.debug("Initial Scaling: POT Downscale applied successfully.") + log.debug(f"{log_prefix}: POT Downscale applied successfully.") else: - log.warning("Initial Scaling: POT Downscale resize failed. Using original data.") - # final_image_data remains image_to_scale + log.warning(f"{log_prefix}: POT Downscale resize failed. Using pre-POT-scaled data.") else: - log.info("Initial Scaling: POT Downscale - Image already POT or smaller. No scaling needed.") - # final_image_data remains image_to_scale - - elif scaling_mode == "NONE": - log.info("Initial Scaling: Mode is NONE. No scaling applied.") - # final_image_data remains image_to_scale + log.info(f"{log_prefix}: Image already POT or smaller. No POT scaling needed.") else: - log.warning(f"Initial Scaling: Unknown INITIAL_SCALING_MODE '{scaling_mode}'. Defaulting to NONE.") - # final_image_data remains image_to_scale + log.warning(f"{log_prefix}: Unknown INITIAL_SCALING_MODE '{scaling_mode}'. Defaulting to NONE (no scaling).") # Determine final dimensions - final_h, final_w = final_image_data.shape[:2] - final_dims_wh = (final_w, final_h) + if final_image_data is not None and final_image_data.size > 0: + final_h, final_w = final_image_data.shape[:2] + final_dims_wh = (final_w, final_h) + else: + final_dims_wh = (0,0) + if final_image_data is None: # Ensure it's an empty array for consistency if None + final_image_data = np.array([]) return InitialScalingOutput( scaled_image_data=final_image_data, scaling_applied=scaling_applied, - final_dimensions=final_dims_wh + final_dimensions=final_dims_wh, + resolution_key=output_resolution_key # Pass through the resolution key ) \ No newline at end of file diff --git a/processing/pipeline/stages/metadata_initialization.py b/processing/pipeline/stages/metadata_initialization.py index e77ef96..8870253 100644 --- a/processing/pipeline/stages/metadata_initialization.py +++ b/processing/pipeline/stages/metadata_initialization.py @@ -148,12 +148,15 @@ class MetadataInitializationStage(ProcessingStage): context.asset_metadata['processing_start_time'] = datetime.datetime.now().isoformat() context.asset_metadata['status'] = "Pending" - if context.config_obj and hasattr(context.config_obj, 'general_settings') and \ - hasattr(context.config_obj.general_settings, 'app_version'): - context.asset_metadata['version'] = context.config_obj.general_settings.app_version + app_version_value = None + if context.config_obj and hasattr(context.config_obj, 'app_version'): + app_version_value = context.config_obj.app_version + + if app_version_value: + context.asset_metadata['version'] = app_version_value else: - logger.warning("App version not found in config_obj.general_settings. Setting version to 'N/A'.") - context.asset_metadata['version'] = "N/A" # Default or placeholder + logger.warning("App version not found using config_obj.app_version. Setting version to 'N/A'.") + context.asset_metadata['version'] = "N/A" if context.incrementing_value is not None: context.asset_metadata['incrementing_value'] = context.incrementing_value diff --git a/processing/pipeline/stages/prepare_processing_items.py b/processing/pipeline/stages/prepare_processing_items.py index cdfc2ac..9c71cfd 100644 --- a/processing/pipeline/stages/prepare_processing_items.py +++ b/processing/pipeline/stages/prepare_processing_items.py @@ -1,21 +1,69 @@ import logging -from typing import List, Union, Optional +from typing import List, Union, Optional, Tuple, Dict # Added Dict +from pathlib import Path # Added Path from .base_stage import ProcessingStage from ..asset_context import AssetProcessingContext, MergeTaskDefinition -from rule_structure import FileRule # Assuming FileRule is imported correctly +from rule_structure import FileRule, ProcessingItem # Added ProcessingItem +from processing.utils import image_processing_utils as ipu # Added ipu log = logging.getLogger(__name__) class PrepareProcessingItemsStage(ProcessingStage): """ - Identifies and prepares a unified list of items (FileRule, MergeTaskDefinition) - to be processed in subsequent stages. Performs initial validation. + Identifies and prepares a unified list of ProcessingItem and MergeTaskDefinition objects + to be processed in subsequent stages. Performs initial validation and explodes + FileRules into specific ProcessingItems for each required output variant. """ + def _get_target_resolutions(self, source_w: int, source_h: int, config_resolutions: dict, file_rule: FileRule) -> Dict[str, int]: + """ + Determines the target output resolutions for a given source image. + Placeholder logic: Uses all config resolutions smaller than or equal to source, plus PREVIEW if smaller. + Needs to be refined to consider FileRule.resolution_override and actual project requirements. + """ + # For now, very basic logic: + # If FileRule has a resolution_override (e.g., (1024,1024)), that might be the *only* target. + # This needs to be clarified. Assuming override means *only* that size. + if file_rule.resolution_override and isinstance(file_rule.resolution_override, tuple) and len(file_rule.resolution_override) == 2: + # How to get a "key" for an arbitrary override? For now, skip if overridden. + # This part of the design (how overrides interact with standard resolutions) is unclear. + # Let's assume for now that if resolution_override is set, we don't generate standard named resolutions. + # This is likely incorrect for a full implementation. + log.warning(f"FileRule '{file_rule.file_path}' has resolution_override. Standard resolution key generation skipped (needs design refinement).") + return {} + + + target_res = {} + max_source_dim = max(source_w, source_h) + + for key, res_val in config_resolutions.items(): + if key == "PREVIEW": # Always consider PREVIEW if its value is smaller + if res_val < max_source_dim : # Or just always include PREVIEW? For now, if smaller. + target_res[key] = res_val + elif res_val <= max_source_dim: + target_res[key] = res_val + + # Ensure PREVIEW is included if it's defined and smaller than the smallest other target, or if no other targets. + # This logic is still a bit naive. + if "PREVIEW" in config_resolutions and config_resolutions["PREVIEW"] < max_source_dim: + if not target_res or config_resolutions["PREVIEW"] < min(v for k,v in target_res.items() if k != "PREVIEW" and isinstance(v,int)): + target_res["PREVIEW"] = config_resolutions["PREVIEW"] + elif "PREVIEW" in config_resolutions and not target_res : # if only preview is applicable + if config_resolutions["PREVIEW"] <= max_source_dim: + target_res["PREVIEW"] = config_resolutions["PREVIEW"] + + + if not target_res and max_source_dim > 0 : # If no standard res is smaller, but image exists + log.debug(f"No standard resolutions from config are <= source dimension {max_source_dim}. Only LOWRES (if applicable) or PREVIEW (if smaller) might be generated.") + + log.debug(f"Determined target resolutions for source {source_w}x{source_h}: {target_res}") + return target_res + + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: """ - Populates context.processing_items with FileRule and MergeTaskDefinition objects. + Populates context.processing_items with ProcessingItem and MergeTaskDefinition objects. """ asset_name_for_log = context.asset_rule.asset_name if context.asset_rule else "Unknown Asset" log.info(f"Asset '{asset_name_for_log}': Preparing processing items...") @@ -25,72 +73,135 @@ class PrepareProcessingItemsStage(ProcessingStage): context.processing_items = [] return context - items_to_process: List[Union[FileRule, MergeTaskDefinition]] = [] + # Output list will now be List[Union[ProcessingItem, MergeTaskDefinition]] + items_to_process: List[Union[ProcessingItem, MergeTaskDefinition]] = [] preparation_failed = False + config = context.config_obj - # --- Add regular files --- + # --- Process FileRules into ProcessingItems --- if context.files_to_process: - # Validate source path early for regular files source_path_valid = True if not context.source_rule or not context.source_rule.input_path: - log.error(f"Asset '{asset_name_for_log}': SourceRule or SourceRule.input_path is not set. Cannot process regular files.") + log.error(f"Asset '{asset_name_for_log}': SourceRule or SourceRule.input_path is not set.") source_path_valid = False - preparation_failed = True # Mark as failed if source path is missing + preparation_failed = True context.status_flags['prepare_items_failed_reason'] = "SourceRule.input_path missing" elif not context.workspace_path or not context.workspace_path.is_dir(): - log.error(f"Asset '{asset_name_for_log}': Workspace path '{context.workspace_path}' is not a valid directory. Cannot process regular files.") + log.error(f"Asset '{asset_name_for_log}': Workspace path '{context.workspace_path}' is invalid.") source_path_valid = False - preparation_failed = True # Mark as failed if workspace path is bad + preparation_failed = True context.status_flags['prepare_items_failed_reason'] = "Workspace path invalid" if source_path_valid: for file_rule in context.files_to_process: - # Basic validation for FileRule itself + log_prefix_fr = f"Asset '{asset_name_for_log}', FileRule '{file_rule.file_path}'" if not file_rule.file_path: - log.warning(f"Asset '{asset_name_for_log}': Skipping FileRule with empty file_path.") - continue # Skip this specific rule, but don't fail the whole stage - items_to_process.append(file_rule) - log.debug(f"Asset '{asset_name_for_log}': Added {len(context.files_to_process)} potential FileRule items.") - else: - log.warning(f"Asset '{asset_name_for_log}': Skipping addition of all FileRule items due to invalid source/workspace path.") + log.warning(f"{log_prefix_fr}: Skipping FileRule with empty file_path.") + continue + + item_type = file_rule.item_type_override or file_rule.item_type + if not item_type or item_type == "EXTRA" or not item_type.startswith("MAP_"): + log.debug(f"{log_prefix_fr}: Item type is '{item_type}'. Not creating map ProcessingItems.") + # Optionally, create a different kind of ProcessingItem for EXTRAs if they need pipeline processing + continue + + source_image_path = context.workspace_path / file_rule.file_path + if not source_image_path.is_file(): + log.error(f"{log_prefix_fr}: Source image file not found at '{source_image_path}'. Skipping this FileRule.") + preparation_failed = True # Individual file error can contribute to overall stage failure + context.status_flags.setdefault('prepare_items_file_errors', []).append(str(source_image_path)) + continue + + # Load image data to get dimensions and for LOWRES variant + # This data will be passed to subsequent stages via ProcessingItem. + # Consider caching this load if RegularMapProcessorStage also loads. + # For now, load here as dimensions are needed for LOWRES decision. + log.debug(f"{log_prefix_fr}: Loading image from '{source_image_path}' to determine dimensions and prepare items.") + source_image_data = ipu.load_image(str(source_image_path)) + if source_image_data is None: + log.error(f"{log_prefix_fr}: Failed to load image from '{source_image_path}'. Skipping this FileRule.") + preparation_failed = True + context.status_flags.setdefault('prepare_items_file_errors', []).append(f"Failed to load {source_image_path}") + continue + + orig_h, orig_w = source_image_data.shape[:2] + original_dimensions_wh = (orig_w, orig_h) + source_bit_depth = ipu.get_image_bit_depth(str(source_image_path)) # Get bit depth from file + source_channels = ipu.get_image_channels(source_image_data) - # --- Add merged tasks --- - # --- Add merged tasks from global configuration --- - # merged_image_tasks are expected to be loaded into context.config_obj - # by the Configuration class from app_settings.json. - - merged_tasks_list = getattr(context.config_obj, 'map_merge_rules', None) + # Determine standard resolutions to generate + # This logic needs to be robust and consider file_rule.resolution_override, etc. + # Using a placeholder _get_target_resolutions for now. + target_resolutions = self._get_target_resolutions(orig_w, orig_h, config.image_resolutions, file_rule) + for res_key, _res_val in target_resolutions.items(): + pi = ProcessingItem( + source_file_info_ref=str(source_image_path), # Using full path as ref + map_type_identifier=item_type, + resolution_key=res_key, + image_data=source_image_data.copy(), # Give each PI its own copy + original_dimensions=original_dimensions_wh, + current_dimensions=original_dimensions_wh, + bit_depth=source_bit_depth, + channels=source_channels, + status="Pending" + ) + items_to_process.append(pi) + log.debug(f"{log_prefix_fr}: Created standard ProcessingItem: {pi.map_type_identifier}_{pi.resolution_key}") + + # Create LOWRES variant if applicable + if config.enable_low_resolution_fallback and max(orig_w, orig_h) < config.low_resolution_threshold: + # Check if a LOWRES item for this source_file_info_ref already exists (e.g. if target_resolutions was empty) + # This check is important if _get_target_resolutions might return empty for small images. + # A more robust way is to ensure LOWRES is distinct from standard resolutions. + + # Avoid duplicate LOWRES if _get_target_resolutions somehow already made one (unlikely with current placeholder) + is_lowres_already_added = any(p.resolution_key == "LOWRES" and p.source_file_info_ref == str(source_image_path) for p in items_to_process if isinstance(p, ProcessingItem)) + + if not is_lowres_already_added: + pi_lowres = ProcessingItem( + source_file_info_ref=str(source_image_path), + map_type_identifier=item_type, + resolution_key="LOWRES", + image_data=source_image_data.copy(), # Fresh copy for LOWRES + original_dimensions=original_dimensions_wh, + current_dimensions=original_dimensions_wh, + bit_depth=source_bit_depth, + channels=source_channels, + status="Pending" + ) + items_to_process.append(pi_lowres) + log.info(f"{log_prefix_fr}: Created LOWRES ProcessingItem because {orig_w}x{orig_h} < {config.low_resolution_threshold}px threshold.") + else: + log.debug(f"{log_prefix_fr}: LOWRES item for this source already added by target resolution logic. Skipping duplicate LOWRES creation.") + elif config.enable_low_resolution_fallback: + log.debug(f"{log_prefix_fr}: Image {orig_w}x{orig_h} not below LOWRES threshold {config.low_resolution_threshold}px.") + + + else: # Source path not valid + log.warning(f"Asset '{asset_name_for_log}': Skipping creation of ProcessingItems from FileRules due to invalid source/workspace path.") + + # --- Add MergeTaskDefinitions --- (This part remains largely the same) + merged_tasks_list = getattr(config, 'map_merge_rules', None) if merged_tasks_list and isinstance(merged_tasks_list, list): log.debug(f"Asset '{asset_name_for_log}': Found {len(merged_tasks_list)} merge tasks in global config.") for task_idx, task_data in enumerate(merged_tasks_list): if isinstance(task_data, dict): task_key = f"merged_task_{task_idx}" - # Basic validation for merge task data: requires output_map_type and an inputs dictionary if not task_data.get('output_map_type') or not isinstance(task_data.get('inputs'), dict): - log.warning(f"Asset '{asset_name_for_log}', Task Index {task_idx}: Skipping merge task due to missing 'output_map_type' or valid 'inputs' dictionary. Task data: {task_data}") - continue # Skip this specific task - log.debug(f"Asset '{asset_name_for_log}', Preparing Merge Task Index {task_idx}: Raw task_data: {task_data}") + log.warning(f"Asset '{asset_name_for_log}', Task Index {task_idx}: Skipping merge task due to missing 'output_map_type' or valid 'inputs'. Task data: {task_data}") + continue merge_def = MergeTaskDefinition(task_data=task_data, task_key=task_key) - log.debug(f"Asset '{asset_name_for_log}': Created MergeTaskDefinition object: {merge_def}") - log.info(f"Asset '{asset_name_for_log}': Successfully CREATED MergeTaskDefinition: Key='{merge_def.task_key}', OutputType='{merge_def.task_data.get('output_map_type', 'N/A')}'") items_to_process.append(merge_def) + log.info(f"Asset '{asset_name_for_log}': Added MergeTaskDefinition: Key='{merge_def.task_key}', OutputType='{merge_def.task_data.get('output_map_type', 'N/A')}'") else: - log.warning(f"Asset '{asset_name_for_log}': Item at index {task_idx} in config_obj.merged_image_tasks is not a dictionary. Skipping. Item: {task_data}") - # The log for "Added X potential MergeTaskDefinition items" will be covered by the final log. - elif merged_tasks_list is None: - log.debug(f"Asset '{asset_name_for_log}': 'merged_image_tasks' not found in config_obj. No global merge tasks to add.") - elif not isinstance(merged_tasks_list, list): - log.warning(f"Asset '{asset_name_for_log}': 'merged_image_tasks' in config_obj is not a list. Skipping global merge tasks. Type: {type(merged_tasks_list)}") - else: # Empty list - log.debug(f"Asset '{asset_name_for_log}': 'merged_image_tasks' in config_obj is empty. No global merge tasks to add.") + log.warning(f"Asset '{asset_name_for_log}': Item at index {task_idx} in config.map_merge_rules is not a dict. Skipping. Item: {task_data}") + # ... (rest of merge task handling) ... + if not items_to_process and not preparation_failed: # Check preparation_failed too + log.info(f"Asset '{asset_name_for_log}': No valid items (ProcessingItem or MergeTaskDefinition) found to process.") - if not items_to_process: - log.info(f"Asset '{asset_name_for_log}': No valid items found to process after preparation.") - - log.debug(f"Asset '{asset_name_for_log}': Final items_to_process before assigning to context: {items_to_process}") context.processing_items = items_to_process context.intermediate_results = {} # Initialize intermediate results storage diff --git a/processing/pipeline/stages/regular_map_processor.py b/processing/pipeline/stages/regular_map_processor.py index 964aaf8..e07ba64 100644 --- a/processing/pipeline/stages/regular_map_processor.py +++ b/processing/pipeline/stages/regular_map_processor.py @@ -37,7 +37,7 @@ class RegularMapProcessorStage(ProcessingStage): """ final_internal_map_type = initial_internal_map_type # Default - base_map_type_match = re.match(r"(MAP_[A-Z]{3})", initial_internal_map_type) + base_map_type_match = re.match(r"(MAP_[A-Z]+)", initial_internal_map_type) if not base_map_type_match or not asset_rule or not asset_rule.files: return final_internal_map_type # Cannot determine suffix without base type or asset rule files @@ -47,7 +47,7 @@ class RegularMapProcessorStage(ProcessingStage): peers_of_same_base_type = [] for fr_asset in asset_rule.files: fr_asset_item_type = fr_asset.item_type_override or fr_asset.item_type or "UnknownMapType" - fr_asset_base_match = re.match(r"(MAP_[A-Z]{3})", fr_asset_item_type) + fr_asset_base_match = re.match(r"(MAP_[A-Z]+)", fr_asset_item_type) if fr_asset_base_match and fr_asset_base_match.group(1) == true_base_map_type: peers_of_same_base_type.append(fr_asset) @@ -197,10 +197,17 @@ class RegularMapProcessorStage(ProcessingStage): result.final_internal_map_type = final_map_type # Update if Gloss->Rough changed it result.transformations_applied = transform_notes + # --- Determine Resolution Key for LOWRES --- + if config.enable_low_resolution_fallback and result.original_dimensions: + w, h = result.original_dimensions + if max(w, h) < config.low_resolution_threshold: + result.resolution_key = "LOWRES" + log.info(f"{log_prefix}: Image dimensions ({w}x{h}) are below threshold ({config.low_resolution_threshold}px). Flagging as LOWRES.") + # --- Success --- result.status = "Processed" result.error_message = None - log.info(f"{log_prefix}: Successfully processed regular map. Final type: '{result.final_internal_map_type}'.") + log.info(f"{log_prefix}: Successfully processed regular map. Final type: '{result.final_internal_map_type}', ResolutionKey: {result.resolution_key}.") except Exception as e: log.exception(f"{log_prefix}: Unhandled exception during processing: {e}") diff --git a/processing/pipeline/stages/save_variants.py b/processing/pipeline/stages/save_variants.py index 482b1cc..7c76482 100644 --- a/processing/pipeline/stages/save_variants.py +++ b/processing/pipeline/stages/save_variants.py @@ -23,8 +23,17 @@ class SaveVariantsStage(ProcessingStage): Calls isu.save_image_variants with data from input_data. """ internal_map_type = input_data.internal_map_type - log_prefix = f"Save Variants Stage (Type: {internal_map_type})" + # The input_data for SaveVariantsStage doesn't directly contain the ProcessingItem. + # It receives data *derived* from a ProcessingItem by previous stages. + # For debugging, we'd need to pass more context or rely on what's in output_filename_pattern_tokens. + resolution_key_from_tokens = input_data.output_filename_pattern_tokens.get('resolution', 'UnknownResKey') + log_prefix = f"Save Variants Stage (Type: {internal_map_type}, ResKey: {resolution_key_from_tokens})" + log.info(f"{log_prefix}: Starting.") + log.debug(f"{log_prefix}: Input image_data shape: {input_data.image_data.shape if input_data.image_data is not None else 'None'}") + log.debug(f"{log_prefix}: Input source_bit_depth_info: {input_data.source_bit_depth_info}") + log.debug(f"{log_prefix}: Configured image_resolutions for saving: {input_data.image_resolutions}") + log.debug(f"{log_prefix}: Output filename pattern tokens: {input_data.output_filename_pattern_tokens}") # Initialize output object with default failure state result = SaveVariantsOutput( @@ -64,11 +73,11 @@ class SaveVariantsStage(ProcessingStage): "resolution_threshold_for_jpg": input_data.resolution_threshold_for_jpg, # Added } - log.debug(f"{log_prefix}: Calling save_image_variants utility.") + log.debug(f"{log_prefix}: Calling save_image_variants utility with args: {save_args}") saved_files_details: List[Dict] = isu.save_image_variants(**save_args) if saved_files_details: - log.info(f"{log_prefix}: Save utility completed successfully. Saved {len(saved_files_details)} variants.") + log.info(f"{log_prefix}: Save utility completed successfully. Saved {len(saved_files_details)} variants: {[details.get('filepath') for details in saved_files_details]}") result.saved_files_details = saved_files_details result.status = "Processed" result.error_message = None diff --git a/processing/utils/image_processing_utils.py b/processing/utils/image_processing_utils.py index 70da34a..143b6bb 100644 --- a/processing/utils/image_processing_utils.py +++ b/processing/utils/image_processing_utils.py @@ -194,6 +194,16 @@ def get_image_bit_depth(image_path_str: str) -> Optional[int]: print(f"Error getting bit depth for {image_path_str}: {e}") return None +def get_image_channels(image_data: np.ndarray) -> Optional[int]: + """Determines the number of channels in an image.""" + if image_data is None: + return None + if len(image_data.shape) == 2: # Grayscale + return 1 + elif len(image_data.shape) == 3: # Color + return image_data.shape[2] + return None # Unknown shape + def calculate_image_stats(image_data: np.ndarray) -> Optional[Dict]: """ Calculates min, max, mean for a given numpy image array. diff --git a/rule_structure.py b/rule_structure.py index 14b2313..5b04fa9 100644 --- a/rule_structure.py +++ b/rule_structure.py @@ -1,6 +1,7 @@ import dataclasses import json from typing import List, Dict, Any, Tuple, Optional +import numpy as np # Added for ProcessingItem @dataclasses.dataclass class FileRule: file_path: str = None @@ -10,8 +11,12 @@ class FileRule: resolution_override: Tuple[int, int] = None channel_merge_instructions: Dict[str, Any] = dataclasses.field(default_factory=dict) output_format_override: str = None + processing_items: List['ProcessingItem'] = dataclasses.field(default_factory=list) # Added field def to_json(self) -> str: + # Need to handle ProcessingItem serialization if it contains non-serializable types like np.ndarray + # For now, assume asdict handles it or it's handled before calling to_json for persistence. + # A custom asdict_factory might be needed for robust serialization. return json.dumps(dataclasses.asdict(self), indent=4) @classmethod @@ -54,4 +59,43 @@ class SourceRule: data = json.loads(json_string) # Manually deserialize nested AssetRule objects data['assets'] = [AssetRule.from_json(json.dumps(asset_data)) for asset_data in data.get('assets', [])] - return cls(**data) \ No newline at end of file + # Need to handle ProcessingItem deserialization if it was serialized + # For now, from_json for FileRule doesn't explicitly handle processing_items from JSON. + return cls(**data) + +@dataclasses.dataclass +class ProcessingItem: + """ + Represents a specific version of an image map to be processed and saved. + This could be a standard resolution (1K, 2K), a preview, or a special + variant like 'LOWRES'. + """ + source_file_info_ref: str # Reference to the original SourceFileInfo or unique ID of the source image + map_type_identifier: str # The internal map type (e.g., "MAP_COL", "MAP_ROUGH") + resolution_key: str # The resolution identifier (e.g., "1K", "PREVIEW", "LOWRES") + image_data: np.ndarray # The actual image data for this item + original_dimensions: Tuple[int, int] # (width, height) of the source image for this item + current_dimensions: Tuple[int, int] # (width, height) of the image_data in this item + target_filename: str = "" # Will be populated by SaveVariantsStage + is_extra: bool = False # If this item should be treated as an 'extra' file + bit_depth: Optional[int] = None + channels: Optional[int] = None + file_extension: Optional[str] = None # Determined during saving based on format + processing_applied_log: List[str] = dataclasses.field(default_factory=list) + status: str = "Pending" # e.g., Pending, Processed, Failed + error_message: Optional[str] = None + + # __getstate__ and __setstate__ might be needed if we pickle these objects + # and np.ndarray causes issues. For JSON, image_data would typically not be serialized. + def __getstate__(self): + state = self.__dict__.copy() + # Don't pickle image_data if it's large or not needed for state + if 'image_data' in state: # Or a more sophisticated check + del state['image_data'] # Example: remove it + return state + + def __setstate__(self, state): + self.__dict__.update(state) + # Potentially re-initialize or handle missing 'image_data' + if 'image_data' not in self.__dict__: + self.image_data = None # Or load it if a path was stored instead \ No newline at end of file