From 12cf557dd7e3636edc2c65823e217279bb3272d2 Mon Sep 17 00:00:00 2001 From: Rusfort Date: Fri, 9 May 2025 11:32:16 +0200 Subject: [PATCH] Uncompleted Processing Refactor --- ProjectNotes/ProcessingEngineRefactorPlan.md | 181 ++ main.py | 7 + processing/pipeline/asset_context.py | 24 + processing/pipeline/orchestrator.py | 133 ++ .../stages/alpha_extraction_to_mask.py | 175 ++ .../pipeline/stages/asset_skip_logic.py | 48 + processing/pipeline/stages/base_stage.py | 22 + .../pipeline/stages/file_rule_filter.py | 80 + .../stages/gloss_to_rough_conversion.py | 156 ++ .../stages/individual_map_processing.py | 245 +++ processing/pipeline/stages/map_merging.py | 310 ++++ .../stages/metadata_finalization_save.py | 119 ++ .../stages/metadata_initialization.py | 163 ++ .../stages/normal_map_green_channel.py | 154 ++ .../pipeline/stages/output_organization.py | 155 ++ .../pipeline/stages/supplier_determination.py | 61 + processing/utils/__init__.py | 1 + processing/utils/image_processing_utils.py | 357 ++++ processing_engine.py | 1589 +---------------- tests/__init__.py | 1 + tests/processing/pipeline/__init__.py | 1 + tests/processing/pipeline/stages/__init__.py | 1 + .../stages/test_alpha_extraction_to_mask.py | 273 +++ .../pipeline/stages/test_asset_skip_logic.py | 213 +++ .../pipeline/stages/test_file_rule_filter.py | 330 ++++ .../stages/test_gloss_to_rough_conversion.py | 486 +++++ .../stages/test_individual_map_processing.py | 555 ++++++ .../pipeline/stages/test_map_merging.py | 538 ++++++ .../stages/test_metadata_finalization_save.py | 359 ++++ .../stages/test_metadata_initialization.py | 169 ++ .../stages/test_normal_map_green_channel.py | 323 ++++ .../stages/test_output_organization.py | 417 +++++ .../stages/test_supplier_determination.py | 213 +++ .../processing/pipeline/test_orchestrator.py | 383 ++++ .../utils/test_image_processing_utils.py | 504 ++++++ tests/utils/__init__.py | 1 + tests/utils/test_path_utils.py | 252 +++ utils/path_utils.py | 9 + 38 files changed, 7472 insertions(+), 1536 deletions(-) create mode 100644 ProjectNotes/ProcessingEngineRefactorPlan.md create mode 100644 processing/pipeline/asset_context.py create mode 100644 processing/pipeline/orchestrator.py create mode 100644 processing/pipeline/stages/alpha_extraction_to_mask.py create mode 100644 processing/pipeline/stages/asset_skip_logic.py create mode 100644 processing/pipeline/stages/base_stage.py create mode 100644 processing/pipeline/stages/file_rule_filter.py create mode 100644 processing/pipeline/stages/gloss_to_rough_conversion.py create mode 100644 processing/pipeline/stages/individual_map_processing.py create mode 100644 processing/pipeline/stages/map_merging.py create mode 100644 processing/pipeline/stages/metadata_finalization_save.py create mode 100644 processing/pipeline/stages/metadata_initialization.py create mode 100644 processing/pipeline/stages/normal_map_green_channel.py create mode 100644 processing/pipeline/stages/output_organization.py create mode 100644 processing/pipeline/stages/supplier_determination.py create mode 100644 processing/utils/__init__.py create mode 100644 processing/utils/image_processing_utils.py create mode 100644 tests/__init__.py create mode 100644 tests/processing/pipeline/__init__.py create mode 100644 tests/processing/pipeline/stages/__init__.py create mode 100644 tests/processing/pipeline/stages/test_alpha_extraction_to_mask.py create mode 100644 tests/processing/pipeline/stages/test_asset_skip_logic.py create mode 100644 tests/processing/pipeline/stages/test_file_rule_filter.py create mode 100644 tests/processing/pipeline/stages/test_gloss_to_rough_conversion.py create mode 100644 tests/processing/pipeline/stages/test_individual_map_processing.py create mode 100644 tests/processing/pipeline/stages/test_map_merging.py create mode 100644 tests/processing/pipeline/stages/test_metadata_finalization_save.py create mode 100644 tests/processing/pipeline/stages/test_metadata_initialization.py create mode 100644 tests/processing/pipeline/stages/test_normal_map_green_channel.py create mode 100644 tests/processing/pipeline/stages/test_output_organization.py create mode 100644 tests/processing/pipeline/stages/test_supplier_determination.py create mode 100644 tests/processing/pipeline/test_orchestrator.py create mode 100644 tests/processing/utils/test_image_processing_utils.py create mode 100644 tests/utils/__init__.py create mode 100644 tests/utils/test_path_utils.py diff --git a/ProjectNotes/ProcessingEngineRefactorPlan.md b/ProjectNotes/ProcessingEngineRefactorPlan.md new file mode 100644 index 0000000..1364ab6 --- /dev/null +++ b/ProjectNotes/ProcessingEngineRefactorPlan.md @@ -0,0 +1,181 @@ +# Project Plan: Modularizing the Asset Processing Engine + +**Last Updated:** May 9, 2025 + +**1. Project Vision & Goals** + +* **Vision:** Transform the asset processing pipeline into a highly modular, extensible, and testable system. +* **Primary Goals:** + 1. Decouple processing steps into independent, reusable stages. + 2. Simplify the addition of new processing capabilities (e.g., GLOSS > ROUGH conversion, Alpha to MASK, Normal Map Green Channel inversion). + 3. Improve code maintainability and readability. + 4. Enhance unit and integration testing capabilities for each processing component. + 5. Centralize common utility functions (image manipulation, path generation). + +**2. Proposed Architecture Overview** + +* **Core Concept:** A `PipelineOrchestrator` will manage a sequence of `ProcessingStage`s. Each stage will operate on an `AssetProcessingContext` object, which carries all necessary data and state for a single asset through the pipeline. +* **Key Components:** + * `AssetProcessingContext`: Data class holding asset-specific data, configuration, temporary paths, and status. + * `PipelineOrchestrator`: Class to manage the overall processing flow for a `SourceRule`, iterating through assets and executing the pipeline of stages for each. + * `ProcessingStage` (Base Class/Interface): Defines the contract for all individual processing stages (e.g., `execute(context)` method). + * Specific Stage Classes: (e.g., `SupplierDeterminationStage`, `IndividualMapProcessingStage`, etc.) + * Utility Modules: `image_processing_utils.py`, enhancements to `utils/path_utils.py`. + +**3. Proposed File Structure** + +* `processing/` + * `pipeline/` + * `__init__.py` + * `asset_context.py` (Defines `AssetProcessingContext`) + * `orchestrator.py` (Defines `PipelineOrchestrator`) + * `stages/` + * `__init__.py` + * `base_stage.py` (Defines `ProcessingStage` interface) + * `supplier_determination.py` + * `asset_skip_logic.py` + * `metadata_initialization.py` + * `file_rule_filter.py` + * `gloss_to_rough_conversion.py` + * `alpha_extraction_to_mask.py` + * `normal_map_green_channel.py` + * `individual_map_processing.py` + * `map_merging.py` + * `metadata_finalization.py` + * `output_organization.py` + * `utils/` + * `__init__.py` + * `image_processing_utils.py` (New module for image functions) +* `utils/` (Top-level existing directory) + * `path_utils.py` (To be enhanced with `sanitize_filename` from `processing_engine.py`) + +**4. Detailed Phases and Tasks** + +**Phase 0: Setup & Core Structures Definition** +*Goal: Establish the foundational classes for the new pipeline.* +* **Task 0.1: Define `AssetProcessingContext`** + * Create `processing/pipeline/asset_context.py`. + * Define the `AssetProcessingContext` data class with fields: `source_rule: SourceRule`, `asset_rule: AssetRule`, `workspace_path: Path`, `engine_temp_dir: Path`, `output_base_path: Path`, `effective_supplier: Optional[str]`, `asset_metadata: Dict`, `processed_maps_details: Dict[str, Dict[str, Dict]]`, `merged_maps_details: Dict[str, Dict[str, Dict]]`, `files_to_process: List[FileRule]`, `loaded_data_cache: Dict`, `config_obj: Configuration`, `status_flags: Dict`, `incrementing_value: Optional[str]`, `sha5_value: Optional[str]`. + * Ensure proper type hinting. +* **Task 0.2: Define `ProcessingStage` Base Class/Interface** + * Create `processing/pipeline/stages/base_stage.py`. + * Define an abstract base class `ProcessingStage` with an abstract method `execute(self, context: AssetProcessingContext) -> AssetProcessingContext`. +* **Task 0.3: Implement Initial `PipelineOrchestrator`** + * Create `processing/pipeline/orchestrator.py`. + * Define the `PipelineOrchestrator` class. + * Implement `__init__(self, config_obj: Configuration, stages: List[ProcessingStage])`. + * Implement `process_source_rule(self, source_rule: SourceRule, workspace_path: Path, output_base_path: Path, overwrite: bool, incrementing_value: Optional[str], sha5_value: Optional[str]) -> Dict[str, List[str]]`. + * Handles creation/cleanup of the main engine temporary directory. + * Loops through `source_rule.assets`, initializes `AssetProcessingContext` for each. + * Iterates `self.stages`, calling `stage.execute(context)`. + * Collects overall status. + +**Phase 1: Utility Module Refactoring** +*Goal: Consolidate and centralize common utility functions.* +* **Task 1.1: Refactor Path Utilities** + * Move `_sanitize_filename` from `processing_engine.py` to `utils/path_utils.py`. + * Update uses to call the new utility function. +* **Task 1.2: Create `image_processing_utils.py`** + * Create `processing/utils/image_processing_utils.py`. + * Move general-purpose image functions from `processing_engine.py`: + * `is_power_of_two` + * `get_nearest_pot` + * `calculate_target_dimensions` + * `calculate_image_stats` + * `normalize_aspect_ratio_change` + * Core image loading, BGR<>RGB conversion, generic resizing (from `_load_and_transform_source`). + * Core data type conversion for saving, color conversion for saving, `cv2.imwrite` call (from `_save_image`). + * Ensure functions are pure and testable. + +**Phase 2: Implementing Core Processing Stages (Migrating Existing Logic)** +*Goal: Migrate existing functionalities from `processing_engine.py` into the new stage-based architecture.* +(For each task: create stage file, implement class, move logic, adapt to `AssetProcessingContext`) +* **Task 2.1: Implement `SupplierDeterminationStage`** +* **Task 2.2: Implement `AssetSkipLogicStage`** +* **Task 2.3: Implement `MetadataInitializationStage`** +* **Task 2.4: Implement `FileRuleFilterStage`** (New logic for `item_type == "FILE_IGNORE"`) +* **Task 2.5: Implement `IndividualMapProcessingStage`** (Adapts `_process_individual_maps`, uses `image_processing_utils.py`) +* **Task 2.6: Implement `MapMergingStage`** (Adapts `_merge_maps`, uses `image_processing_utils.py`) +* **Task 2.7: Implement `MetadataFinalizationAndSaveStage`** (Adapts `_generate_metadata_file`, uses `utils.path_utils.generate_path_from_pattern`) +* **Task 2.8: Implement `OutputOrganizationStage`** (Adapts `_organize_output_files`) + +**Phase 3: Implementing New Feature Stages** +*Goal: Add the new desired processing capabilities as distinct stages.* +* **Task 3.1: Implement `GlossToRoughConversionStage`** (Identify gloss, convert, invert, save temp, update `FileRule`) +* **Task 3.2: Implement `AlphaExtractionToMaskStage`** (Check existing mask, find MAP_COL with alpha, extract, save temp, add new `FileRule`) +* **Task 3.3: Implement `NormalMapGreenChannelStage`** (Identify normal maps, invert green based on config, save temp, update `FileRule`) + +**Phase 4: Integration, Testing & Finalization** +*Goal: Assemble the pipeline, test thoroughly, and deprecate old code.* +* **Task 4.1: Configure `PipelineOrchestrator`** + * Instantiate `PipelineOrchestrator` in main application logic with the ordered list of stage instances. +* **Task 4.2: Unit Testing** + * Unit tests for each `ProcessingStage` (mocking `AssetProcessingContext`). + * Unit tests for `image_processing_utils.py` and `utils/path_utils.py` functions. +* **Task 4.3: Integration Testing** + * Test `PipelineOrchestrator` end-to-end with sample data. + * Compare outputs with the existing engine for consistency. +* **Task 4.4: Documentation Update** + * Update developer documentation (e.g., `Documentation/02_Developer_Guide/05_Processing_Pipeline.md`). + * Document `AssetProcessingContext` and stage responsibilities. +* **Task 4.5: Deprecate/Remove Old `ProcessingEngine` Code** + * Gradually remove refactored logic from `processing_engine.py`. + +**5. Workflow Diagram** + +```mermaid +graph TD + AA[Load SourceRule & Config] --> BA(PipelineOrchestrator: process_source_rule); + BA --> CA{For Each Asset in SourceRule}; + CA -- Yes --> DA(Orchestrator: Create AssetProcessingContext); + DA --> EA(SupplierDeterminationStage); + EA -- context --> FA(AssetSkipLogicStage); + FA -- context --> GA{context.skip_asset?}; + GA -- Yes --> HA(Orchestrator: Record Skipped); + HA --> CA; + GA -- No --> IA(MetadataInitializationStage); + IA -- context --> JA(FileRuleFilterStage); + JA -- context --> KA(GlossToRoughConversionStage); + KA -- context --> LA(AlphaExtractionToMaskStage); + LA -- context --> MA(NormalMapGreenChannelStage); + MA -- context --> NA(IndividualMapProcessingStage); + NA -- context --> OA(MapMergingStage); + OA -- context --> PA(MetadataFinalizationAndSaveStage); + PA -- context --> QA(OutputOrganizationStage); + QA -- context --> RA(Orchestrator: Record Processed/Failed); + RA --> CA; + CA -- No --> SA(Orchestrator: Cleanup Engine Temp Dir); + SA --> TA[Processing Complete]; + + subgraph Stages + direction LR + EA + FA + IA + JA + KA + LA + MA + NA + OA + PA + QA + end + + subgraph Utils + direction LR + U1[image_processing_utils.py] + U2[utils/path_utils.py] + end + + NA -.-> U1; + OA -.-> U1; + KA -.-> U1; + LA -.-> U1; + MA -.-> U1; + + PA -.-> U2; + QA -.-> U2; + + classDef context fill:#f9f,stroke:#333,stroke-width:2px; + class DA,EA,FA,IA,JA,KA,LA,MA,NA,OA,PA,QA context; \ No newline at end of file diff --git a/main.py b/main.py index acc0713..25f8049 100644 --- a/main.py +++ b/main.py @@ -21,6 +21,11 @@ from PySide6.QtCore import Qt from PySide6.QtWidgets import QApplication # --- Backend Imports --- +# Add current directory to sys.path for direct execution +import sys +import os +sys.path.append(os.path.dirname(__file__)) + try: from configuration import Configuration, ConfigurationError from processing_engine import ProcessingEngine @@ -29,6 +34,8 @@ try: from utils.workspace_utils import prepare_processing_workspace except ImportError as e: script_dir = Path(__file__).parent.resolve() + print(f"ERROR: Cannot import Configuration or rule_structure classes.") + print(f"Ensure configuration.py and rule_structure.py are in the same directory or Python path.") print(f"ERROR: Failed to import necessary classes: {e}") print(f"Ensure 'configuration.py' and 'asset_processor.py' exist in the directory:") print(f" {script_dir}") diff --git a/processing/pipeline/asset_context.py b/processing/pipeline/asset_context.py new file mode 100644 index 0000000..5b411d7 --- /dev/null +++ b/processing/pipeline/asset_context.py @@ -0,0 +1,24 @@ +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional + +from rule_structure import AssetRule, FileRule, SourceRule +from configuration import Configuration + +@dataclass +class AssetProcessingContext: + source_rule: SourceRule + asset_rule: AssetRule + workspace_path: Path + engine_temp_dir: Path + output_base_path: Path + effective_supplier: Optional[str] + asset_metadata: Dict + processed_maps_details: Dict[str, Dict[str, Dict]] + merged_maps_details: Dict[str, Dict[str, Dict]] + files_to_process: List[FileRule] + loaded_data_cache: Dict + config_obj: Configuration + status_flags: Dict + incrementing_value: Optional[str] + sha5_value: Optional[str] \ No newline at end of file diff --git a/processing/pipeline/orchestrator.py b/processing/pipeline/orchestrator.py new file mode 100644 index 0000000..6396f38 --- /dev/null +++ b/processing/pipeline/orchestrator.py @@ -0,0 +1,133 @@ +from typing import List, Dict, Optional +from pathlib import Path +import shutil +import tempfile +import logging + +from configuration import Configuration +from rule_structure import SourceRule, AssetRule +from .asset_context import AssetProcessingContext +from .stages.base_stage import ProcessingStage + +log = logging.getLogger(__name__) + +class PipelineOrchestrator: + """ + Orchestrates the processing of assets based on source rules and a series of processing stages. + """ + + def __init__(self, config_obj: Configuration, stages: List[ProcessingStage]): + """ + Initializes the PipelineOrchestrator. + + Args: + config_obj: The main configuration object. + stages: A list of processing stages to be executed in order. + """ + self.config_obj: Configuration = config_obj + self.stages: List[ProcessingStage] = stages + + def process_source_rule( + self, + source_rule: SourceRule, + workspace_path: Path, + output_base_path: Path, + overwrite: bool, # Not used in this initial implementation, but part of the signature + incrementing_value: Optional[str], + sha5_value: Optional[str] # Corrected from sha5_value to sha256_value as per typical usage, assuming typo + ) -> Dict[str, List[str]]: + """ + Processes a single source rule, iterating through its asset rules and applying all stages. + + Args: + source_rule: The source rule to process. + workspace_path: The base path of the workspace. + output_base_path: The base path for output files. + overwrite: Whether to overwrite existing files (not fully implemented yet). + incrementing_value: An optional incrementing value for versioning or naming. + sha5_value: An optional SHA5 hash value for the asset (assuming typo, likely sha256). + + Returns: + A dictionary summarizing the processing status of assets. + """ + overall_status: Dict[str, List[str]] = { + "processed": [], + "skipped": [], + "failed": [], + } + engine_temp_dir_path: Optional[Path] = None # Initialize to None + + try: + # Create a temporary directory for this processing run if needed by any stage + # This temp dir is for the entire source_rule processing, not per asset. + # Individual stages might create their own sub-temp dirs if necessary. + temp_dir_path_str = tempfile.mkdtemp( + prefix="asset_processor_orchestrator_temp_", dir=self.config_obj.get_temp_directory_base() + ) + engine_temp_dir_path = Path(temp_dir_path_str) + log.debug(f"PipelineOrchestrator created temporary directory: {engine_temp_dir_path}") + + + for asset_rule in source_rule.assets: + log.debug(f"Orchestrator: Processing asset '{asset_rule.name}'") + context = AssetProcessingContext( + source_rule=source_rule, + asset_rule=asset_rule, + workspace_path=workspace_path, # This is the path to the source files (e.g. extracted archive) + engine_temp_dir=engine_temp_dir_path, # Pass the orchestrator's temp dir + output_base_path=output_base_path, + effective_supplier=None, # Will be set by SupplierDeterminationStage + asset_metadata={}, # Will be populated by stages + processed_maps_details={}, # Will be populated by stages + merged_maps_details={}, # Will be populated by stages + files_to_process=[], # Will be populated by FileRuleFilterStage + loaded_data_cache={}, # For image loading cache within this asset's processing + config_obj=self.config_obj, + status_flags={"skip_asset": False, "asset_failed": False}, # Initialize common flags + incrementing_value=incrementing_value, + sha256_value=sha5_value # Parameter name in context is sha256_value + ) + + for stage_idx, stage in enumerate(self.stages): + log.debug(f"Asset '{asset_rule.name}': Executing stage {stage_idx + 1}/{len(self.stages)}: {stage.__class__.__name__}") + try: + context = stage.execute(context) + except Exception as e: + log.error(f"Asset '{asset_rule.name}': Error during stage '{stage.__class__.__name__}': {e}", exc_info=True) + context.status_flags["asset_failed"] = True + context.asset_metadata["status"] = f"Failed: Error in stage {stage.__class__.__name__}" + context.asset_metadata["error_message"] = str(e) + break # Stop processing stages for this asset on error + + if context.status_flags.get("skip_asset"): + log.info(f"Asset '{asset_rule.name}': Skipped by stage '{stage.__class__.__name__}'. Reason: {context.status_flags.get('skip_reason', 'N/A')}") + break # Skip remaining stages for this asset + + # Refined status collection + if context.status_flags.get('skip_asset'): + overall_status["skipped"].append(asset_rule.name) + elif context.status_flags.get('asset_failed') or str(context.asset_metadata.get('status', '')).startswith("Failed"): + overall_status["failed"].append(asset_rule.name) + elif context.asset_metadata.get('status') == "Processed": + overall_status["processed"].append(asset_rule.name) + else: # Default or unknown state + log.warning(f"Asset '{asset_rule.name}': Unknown status after pipeline execution. Metadata status: '{context.asset_metadata.get('status')}'. Marking as failed.") + overall_status["failed"].append(f"{asset_rule.name} (Unknown Status: {context.asset_metadata.get('status')})") + log.debug(f"Asset '{asset_rule.name}' final status: {context.asset_metadata.get('status', 'N/A')}, Flags: {context.status_flags}") + + except Exception as e: + log.error(f"PipelineOrchestrator.process_source_rule failed: {e}", exc_info=True) + # Mark all remaining assets as failed if a top-level error occurs + processed_or_skipped_or_failed = set(overall_status["processed"] + overall_status["skipped"] + overall_status["failed"]) + for asset_rule in source_rule.assets: + if asset_rule.name not in processed_or_skipped_or_failed: + overall_status["failed"].append(f"{asset_rule.name} (Orchestrator Error)") + finally: + if engine_temp_dir_path and engine_temp_dir_path.exists(): + try: + log.debug(f"PipelineOrchestrator cleaning up temporary directory: {engine_temp_dir_path}") + shutil.rmtree(engine_temp_dir_path, ignore_errors=True) + except Exception as e: + log.error(f"Error cleaning up orchestrator temporary directory {engine_temp_dir_path}: {e}", exc_info=True) + + return overall_status \ No newline at end of file diff --git a/processing/pipeline/stages/alpha_extraction_to_mask.py b/processing/pipeline/stages/alpha_extraction_to_mask.py new file mode 100644 index 0000000..ca1ea38 --- /dev/null +++ b/processing/pipeline/stages/alpha_extraction_to_mask.py @@ -0,0 +1,175 @@ +import logging +import uuid +from pathlib import Path +from typing import List, Optional, Dict + +import numpy as np + +from .base_stage import ProcessingStage +from ..asset_context import AssetProcessingContext +from ...utils import image_processing_utils as ipu +from .....rule_structure import FileRule, TransformSettings +from .....utils.path_utils import sanitize_filename + +logger = logging.getLogger(__name__) + +class AlphaExtractionToMaskStage(ProcessingStage): + """ + Extracts an alpha channel from a suitable source map (e.g., Albedo, Diffuse) + to generate a MASK map if one is not explicitly defined. + """ + SUITABLE_SOURCE_MAP_TYPES = ["ALBEDO", "DIFFUSE", "BASE_COLOR"] # Map types likely to have alpha + + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: + logger.debug(f"Asset '{context.asset_rule.name}': Running AlphaExtractionToMaskStage.") + + if context.status_flags.get('skip_asset'): + logger.debug(f"Asset '{context.asset_rule.name}': Skipping due to 'skip_asset' flag.") + return context + + if not context.files_to_process or not context.processed_maps_details: + logger.debug( + f"Asset '{context.asset_rule.name}': Skipping alpha extraction - " + f"no files to process or no processed map details." + ) + return context + + # A. Check for Existing MASK Map + for file_rule in context.files_to_process: + if file_rule.map_type == "MASK": + logger.info( + f"Asset '{context.asset_rule.name}': MASK map already defined by FileRule " + f"'{file_rule.filename_pattern}'. Skipping alpha extraction." + ) + return context + + # B. Find Suitable Source Map with Alpha + source_map_details_for_alpha: Optional[Dict] = None + source_file_rule_id_for_alpha: Optional[str] = None + + for file_rule_id, details in context.processed_maps_details.items(): + if details.get('status') == 'Processed' and \ + details.get('map_type') in self.SUITABLE_SOURCE_MAP_TYPES: + try: + temp_path = Path(details['temp_processed_file']) + if not temp_path.exists(): + logger.warning( + f"Asset '{context.asset_rule.name}': Temp file {temp_path} for map " + f"{details['map_type']} (ID: {file_rule_id}) does not exist. Cannot check for alpha." + ) + continue + + # Load image header or minimal data to check for alpha if possible, + # otherwise load full image. ipu.load_image should handle this. + image_data = ipu.load_image(temp_path) + + if image_data is not None and image_data.ndim == 3 and image_data.shape[2] == 4: + source_map_details_for_alpha = details + source_file_rule_id_for_alpha = file_rule_id + logger.info( + f"Asset '{context.asset_rule.name}': Found potential source for alpha extraction: " + f"{temp_path} (MapType: {details['map_type']})" + ) + break + except Exception as e: + logger.warning( + f"Asset '{context.asset_rule.name}': Error checking alpha for {details.get('temp_processed_file', 'N/A')}: {e}" + ) + continue + + + if source_map_details_for_alpha is None or source_file_rule_id_for_alpha is None: + logger.info( + f"Asset '{context.asset_rule.name}': No suitable source map with alpha channel found " + f"for MASK extraction." + ) + return context + + # C. Extract Alpha Channel + source_image_path = Path(source_map_details_for_alpha['temp_processed_file']) + full_image_data = ipu.load_image(source_image_path) # Reload to ensure we have the original RGBA + + if full_image_data is None or not (full_image_data.ndim == 3 and full_image_data.shape[2] == 4): + logger.error( + f"Asset '{context.asset_rule.name}': Failed to reload or verify alpha channel from " + f"{source_image_path} for MASK extraction." + ) + return context + + alpha_channel: np.ndarray = full_image_data[:, :, 3] # Extract alpha (0-255) + + # D. Save New Temporary MASK Map + # Ensure the mask is a 2D grayscale image. If ipu.save_image expects 3 channels for grayscale, adapt. + # Assuming ipu.save_image can handle a 2D numpy array for a grayscale image. + if alpha_channel.ndim == 2: # Expected + pass + elif alpha_channel.ndim == 3 and alpha_channel.shape[2] == 1: # (H, W, 1) + alpha_channel = alpha_channel.squeeze(axis=2) + else: + logger.error( + f"Asset '{context.asset_rule.name}': Extracted alpha channel has unexpected dimensions: " + f"{alpha_channel.shape}. Cannot save." + ) + return context + + mask_temp_filename = ( + f"mask_from_alpha_{sanitize_filename(source_map_details_for_alpha['map_type'])}" + f"_{source_file_rule_id_for_alpha}{source_image_path.suffix}" + ) + mask_temp_path = context.engine_temp_dir / mask_temp_filename + + save_success = ipu.save_image(mask_temp_path, alpha_channel) + + if not save_success: + logger.error( + f"Asset '{context.asset_rule.name}': Failed to save extracted alpha mask to {mask_temp_path}." + ) + return context + + logger.info( + f"Asset '{context.asset_rule.name}': Extracted alpha and saved as new MASK map: {mask_temp_path}" + ) + + # E. Create New FileRule for the MASK and Update Context + new_mask_file_rule_id_obj = uuid.uuid4() + new_mask_file_rule_id_str = str(new_mask_file_rule_id_obj) # Use string for FileRule.id + new_mask_file_rule_id_hex = new_mask_file_rule_id_obj.hex # Use hex for dict key + + new_mask_file_rule = FileRule( + id=new_mask_file_rule_id_str, + map_type="MASK", + filename_pattern=mask_temp_path.name, # Pattern matches the generated temp file + item_type="MAP_COL", # Considered a collected map post-generation + active=True, + transform_settings=TransformSettings(), # Default transform settings + source_map_ids_for_generation=[source_file_rule_id_for_alpha] # Link to original source + # Ensure other necessary FileRule fields are defaulted or set if required + ) + + context.files_to_process.append(new_mask_file_rule) + + original_dims = source_map_details_for_alpha.get('original_dimensions') + if original_dims is None and full_image_data is not None: # Fallback if not in details + original_dims = (full_image_data.shape[1], full_image_data.shape[0]) + + + context.processed_maps_details[new_mask_file_rule_id_hex] = { + 'map_type': "MASK", + 'source_file': str(source_image_path), # Original RGBA map path + 'temp_processed_file': str(mask_temp_path), # Path to the new MASK map + 'original_dimensions': original_dims, # Dimensions of the source image + 'processed_dimensions': (alpha_channel.shape[1], alpha_channel.shape[0]), # Dimensions of MASK + 'status': 'Processed', # This map is now considered processed + 'notes': ( + f"Generated from alpha of {source_map_details_for_alpha['map_type']} " + f"(Source Rule ID: {source_file_rule_id_for_alpha})" + ), + 'file_rule_id': new_mask_file_rule_id_str # Link back to the new FileRule ID + } + + logger.info( + f"Asset '{context.asset_rule.name}': Added new FileRule for generated MASK " + f"(ID: {new_mask_file_rule_id_str}) and updated processed_maps_details." + ) + + return context \ No newline at end of file diff --git a/processing/pipeline/stages/asset_skip_logic.py b/processing/pipeline/stages/asset_skip_logic.py new file mode 100644 index 0000000..afb5b3c --- /dev/null +++ b/processing/pipeline/stages/asset_skip_logic.py @@ -0,0 +1,48 @@ +import logging +from ..base_stage import ProcessingStage +from ...asset_context import AssetProcessingContext + +class AssetSkipLogicStage(ProcessingStage): + """ + Processing stage to determine if an asset should be skipped based on various conditions. + """ + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: + """ + Executes the asset skip logic. + + Args: + context: The asset processing context. + + Returns: + The updated asset processing context. + """ + context.status_flags['skip_asset'] = False # Initialize/reset skip flag + + # 1. Check for Supplier Error + # Assuming 'supplier_error' might be set by a previous stage (e.g., SupplierDeterminationStage) + # or if effective_supplier is None after attempts to determine it. + if context.effective_supplier is None or context.status_flags.get('supplier_error', False): + logging.info(f"Asset '{context.asset_rule.name}': Skipping due to missing or invalid supplier.") + context.status_flags['skip_asset'] = True + context.status_flags['skip_reason'] = "Invalid or missing supplier" + return context + + # 2. Check asset_rule.process_status + if context.asset_rule.process_status == "SKIP": + logging.info(f"Asset '{context.asset_rule.name}': Skipping as per process_status 'SKIP'.") + context.status_flags['skip_asset'] = True + context.status_flags['skip_reason'] = "Process status set to SKIP" + return context + + if context.asset_rule.process_status == "PROCESSED" and \ + not context.config_obj.general_settings.overwrite_existing: + logging.info( + f"Asset '{context.asset_rule.name}': Skipping as it's already 'PROCESSED' " + f"and overwrite is disabled." + ) + context.status_flags['skip_asset'] = True + context.status_flags['skip_reason'] = "Already processed, overwrite disabled" + return context + + # If none of the above conditions are met, skip_asset remains False. + return context \ No newline at end of file diff --git a/processing/pipeline/stages/base_stage.py b/processing/pipeline/stages/base_stage.py new file mode 100644 index 0000000..321a0d4 --- /dev/null +++ b/processing/pipeline/stages/base_stage.py @@ -0,0 +1,22 @@ +from abc import ABC, abstractmethod + +from ..asset_context import AssetProcessingContext + + +class ProcessingStage(ABC): + """ + Abstract base class for a stage in the asset processing pipeline. + """ + + @abstractmethod + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: + """ + Executes the processing logic of this stage. + + Args: + context: The current asset processing context. + + Returns: + The updated asset processing context. + """ + pass \ No newline at end of file diff --git a/processing/pipeline/stages/file_rule_filter.py b/processing/pipeline/stages/file_rule_filter.py new file mode 100644 index 0000000..b7ae7c3 --- /dev/null +++ b/processing/pipeline/stages/file_rule_filter.py @@ -0,0 +1,80 @@ +import logging +import fnmatch +from typing import List, Set + +from ..base_stage import ProcessingStage +from ...asset_context import AssetProcessingContext +from .....rule_structure import FileRule + + +class FileRuleFilterStage(ProcessingStage): + """ + Determines which FileRules associated with an AssetRule should be processed. + Populates context.files_to_process, respecting FILE_IGNORE rules. + """ + + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: + """ + Executes the file rule filtering logic. + + Args: + context: The AssetProcessingContext for the current asset. + + Returns: + The modified AssetProcessingContext. + """ + if context.status_flags.get('skip_asset'): + logging.debug(f"Asset '{context.asset_rule.name}': Skipping FileRuleFilterStage due to 'skip_asset' flag.") + return context + + context.files_to_process: List[FileRule] = [] + ignore_patterns: Set[str] = set() + + # Step 1: Collect all FILE_IGNORE patterns + if context.asset_rule and context.asset_rule.file_rules: + for file_rule in context.asset_rule.file_rules: + if file_rule.item_type == "FILE_IGNORE" and file_rule.active: + ignore_patterns.add(file_rule.filename_pattern) + logging.debug( + f"Asset '{context.asset_rule.name}': Registering ignore pattern: '{file_rule.filename_pattern}'" + ) + else: + logging.debug(f"Asset '{context.asset_rule.name if context.asset_rule else 'Unknown'}': No file rules to process or asset_rule is None.") + # Still need to return context even if there are no rules + logging.info(f"Asset '{context.asset_rule.name if context.asset_rule else 'Unknown'}': 0 file rules queued for processing after filtering.") + return context + + + # Step 2: Filter and add processable FileRules + for file_rule in context.asset_rule.file_rules: + if not file_rule.active: + logging.debug( + f"Asset '{context.asset_rule.name}': Skipping inactive file rule '{file_rule.filename_pattern}'." + ) + continue + + if file_rule.item_type == "FILE_IGNORE": + # Already processed, skip. + continue + + is_ignored = False + for ignore_pat in ignore_patterns: + if fnmatch.fnmatch(file_rule.filename_pattern, ignore_pat): + is_ignored = True + logging.debug( + f"Asset '{context.asset_rule.name}': Skipping file rule '{file_rule.filename_pattern}' " + f"due to matching ignore pattern '{ignore_pat}'." + ) + break + + if not is_ignored: + context.files_to_process.append(file_rule) + logging.debug( + f"Asset '{context.asset_rule.name}': Adding file rule '{file_rule.filename_pattern}' " + f"(type: {file_rule.item_type}) to processing queue." + ) + + logging.info( + f"Asset '{context.asset_rule.name}': {len(context.files_to_process)} file rules queued for processing after filtering." + ) + return context \ No newline at end of file diff --git a/processing/pipeline/stages/gloss_to_rough_conversion.py b/processing/pipeline/stages/gloss_to_rough_conversion.py new file mode 100644 index 0000000..d99f06a --- /dev/null +++ b/processing/pipeline/stages/gloss_to_rough_conversion.py @@ -0,0 +1,156 @@ +import logging +from pathlib import Path +import numpy as np +from typing import List + +from .base_stage import ProcessingStage +from ..asset_context import AssetProcessingContext +from ...rule_structure import FileRule +from ..utils import image_processing_utils as ipu +from ...utils.path_utils import sanitize_filename + +logger = logging.getLogger(__name__) + +class GlossToRoughConversionStage(ProcessingStage): + """ + Processing stage to convert glossiness maps to roughness maps. + Iterates through FileRules, identifies GLOSS maps, loads their + corresponding temporary processed images, inverts them, and saves + them as new temporary ROUGHNESS maps. Updates the FileRule and + context.processed_maps_details accordingly. + """ + + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: + """ + Executes the gloss to roughness conversion logic. + + Args: + context: The AssetProcessingContext containing asset and processing details. + + Returns: + The updated AssetProcessingContext. + """ + if context.status_flags.get('skip_asset'): + logger.debug(f"Asset '{context.asset_rule.name}': Skipping GlossToRoughConversionStage due to skip_asset flag.") + return context + + if not context.files_to_process or not context.processed_maps_details: + logger.debug( + f"Asset '{context.asset_rule.name}': No files to process or processed_maps_details empty " + f"in GlossToRoughConversionStage. Skipping." + ) + return context + + new_files_to_process: List[FileRule] = [] + processed_a_gloss_map = False + + logger.info(f"Asset '{context.asset_rule.name}': Starting Gloss to Roughness Conversion Stage.") + + for idx, file_rule in enumerate(context.files_to_process): + if file_rule.map_type == "GLOSS": + map_detail_key = file_rule.id.hex + if map_detail_key not in context.processed_maps_details: + logger.warning( + f"Asset '{context.asset_rule.name}': GLOSS map '{file_rule.source_file_path}' " + f"(ID: {map_detail_key}) found in files_to_process but not in processed_maps_details. " + f"Adding original rule and skipping conversion for this map." + ) + new_files_to_process.append(file_rule) + continue + + map_details = context.processed_maps_details[map_detail_key] + + if map_details.get('status') != 'Processed' or 'temp_processed_file' not in map_details: + logger.warning( + f"Asset '{context.asset_rule.name}': GLOSS map '{file_rule.source_file_path}' " + f"(ID: {map_detail_key}) not successfully processed by previous stage or temp file missing. " + f"Status: {map_details.get('status')}. Adding original rule and skipping conversion." + ) + new_files_to_process.append(file_rule) + continue + + original_temp_path_str = map_details['temp_processed_file'] + original_temp_path = Path(original_temp_path_str) + + if not original_temp_path.exists(): + logger.error( + f"Asset '{context.asset_rule.name}': Temporary file {original_temp_path_str} for GLOSS map " + f"(ID: {map_detail_key}) does not exist. Adding original rule and skipping conversion." + ) + new_files_to_process.append(file_rule) + continue + + logger.debug(f"Asset '{context.asset_rule.name}': Processing GLOSS map {original_temp_path} for conversion.") + image_data = ipu.load_image(original_temp_path) + + if image_data is None: + logger.error( + f"Asset '{context.asset_rule.name}': Failed to load image data from {original_temp_path} " + f"for GLOSS map (ID: {map_detail_key}). Adding original rule and skipping conversion." + ) + new_files_to_process.append(file_rule) + continue + + # Perform Inversion + inverted_image_data: np.ndarray + if np.issubdtype(image_data.dtype, np.floating): + inverted_image_data = 1.0 - image_data + inverted_image_data = np.clip(inverted_image_data, 0.0, 1.0) # Ensure range for floats + logger.debug(f"Asset '{context.asset_rule.name}': Inverted float image data for {original_temp_path}.") + elif np.issubdtype(image_data.dtype, np.integer): + max_val = np.iinfo(image_data.dtype).max + inverted_image_data = max_val - image_data + logger.debug(f"Asset '{context.asset_rule.name}': Inverted integer image data (max_val: {max_val}) for {original_temp_path}.") + else: + logger.error( + f"Asset '{context.asset_rule.name}': Unsupported image data type {image_data.dtype} " + f"for GLOSS map {original_temp_path}. Cannot invert. Adding original rule." + ) + new_files_to_process.append(file_rule) + continue + + # Save New Temporary (Roughness) Map + # Using original_temp_path.suffix ensures we keep the format (e.g., .png, .exr) + new_temp_filename = f"rough_from_gloss_{sanitize_filename(file_rule.map_type)}_{file_rule.id.hex}{original_temp_path.suffix}" + new_temp_path = context.engine_temp_dir / new_temp_filename + + save_success = ipu.save_image(new_temp_path, inverted_image_data) + + if save_success: + logger.info( + f"Asset '{context.asset_rule.name}': Converted GLOSS map {original_temp_path} " + f"to ROUGHNESS map {new_temp_path}." + ) + + modified_file_rule = file_rule.model_copy(deep=True) + modified_file_rule.map_type = "ROUGHNESS" + + # Update context.processed_maps_details for the original file_rule.id.hex + context.processed_maps_details[map_detail_key]['temp_processed_file'] = str(new_temp_path) + context.processed_maps_details[map_detail_key]['original_map_type_before_conversion'] = "GLOSS" + context.processed_maps_details[map_detail_key]['notes'] = "Converted from GLOSS by GlossToRoughConversionStage" + + new_files_to_process.append(modified_file_rule) + processed_a_gloss_map = True + else: + logger.error( + f"Asset '{context.asset_rule.name}': Failed to save inverted ROUGHNESS map to {new_temp_path} " + f"for original GLOSS map (ID: {map_detail_key}). Adding original rule." + ) + new_files_to_process.append(file_rule) + else: # Not a gloss map + new_files_to_process.append(file_rule) + + context.files_to_process = new_files_to_process + + if processed_a_gloss_map: + logger.info( + f"Asset '{context.asset_rule.name}': Gloss to Roughness conversion stage successfully processed one or more maps and updated file list." + ) + else: + logger.debug( + f"Asset '{context.asset_rule.name}': No gloss maps were successfully converted in GlossToRoughConversionStage. " + f"File list for next stage contains original non-gloss maps and any gloss maps that failed conversion." + ) + + return context \ No newline at end of file diff --git a/processing/pipeline/stages/individual_map_processing.py b/processing/pipeline/stages/individual_map_processing.py new file mode 100644 index 0000000..72552c4 --- /dev/null +++ b/processing/pipeline/stages/individual_map_processing.py @@ -0,0 +1,245 @@ +import os +import logging +from pathlib import Path +from typing import Optional, Tuple, Dict + +import cv2 +import numpy as np + +from ..base_stage import ProcessingStage +from ..asset_context import AssetProcessingContext +from ....rule_structure import FileRule, TransformSettings +from ....utils.path_utils import sanitize_filename +from ...utils import image_processing_utils as ipu + +logger = logging.getLogger(__name__) + +class IndividualMapProcessingStage(ProcessingStage): + """ + Processes individual texture map files based on FileRules. + This stage finds the source file, loads it, applies transformations + (resize, color space), saves a temporary processed version, and updates + the AssetProcessingContext with details. + """ + + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: + """ + Executes the individual map processing logic. + """ + if context.status_flags.get('skip_asset', False): + logger.info(f"Asset {context.asset_id}: Skipping individual map processing due to skip_asset flag.") + return context + + if not hasattr(context, 'processed_maps_details') or context.processed_maps_details is None: + context.processed_maps_details = {} + logger.debug(f"Asset {context.asset_id}: Initialized processed_maps_details.") + + if not context.files_to_process: + logger.info(f"Asset {context.asset_id}: No files to process in this stage.") + return context + + source_base_path = Path(context.asset_rule.source_path) + if not source_base_path.is_dir(): + logger.error(f"Asset {context.asset_id}: Source path '{source_base_path}' is not a valid directory. Skipping individual map processing.") + context.status_flags['individual_map_processing_failed'] = True + # Potentially mark all file_rules as failed if source path is invalid + for file_rule in context.files_to_process: + if file_rule.item_type.startswith("MAP_"): # General check for map types + self._update_file_rule_status(context, file_rule.id.hex, 'Failed', details="Source path invalid") + return context + + for file_rule in context.files_to_process: + # Primarily focus on "MAP_COL", "MAP_NORM", "MAP_ROUGH", etc. + # For now, let's assume any item_type starting with "MAP_" is a candidate + # unless it's specifically handled by another stage (e.g., "MAP_GEN" might be). + # The prompt mentions "MAP_COL" primarily. + # Let's be a bit more specific for now, focusing on types that are typically direct file mappings. + # This can be refined based on how `item_type` is used for generated maps. + # For now, we'll process any `FileRule` that isn't explicitly a generated map type + # that would be handled *after* individual processing (e.g. a composite map). + # A simple check for now: + if not file_rule.item_type or not file_rule.item_type.startswith("MAP_") or file_rule.item_type == "MAP_GEN_COMPOSITE": # Example exclusion + logger.debug(f"Asset {context.asset_id}, FileRule {file_rule.id.hex} ({file_rule.map_type}): Skipping, item_type '{file_rule.item_type}' not targeted for individual processing.") + continue + + logger.info(f"Asset {context.asset_id}, FileRule {file_rule.id.hex} ({file_rule.map_type}): Starting individual processing.") + + # A. Find Source File + source_file_path = self._find_source_file(source_base_path, file_rule.filename_pattern, context.asset_id, file_rule.id.hex) + if not source_file_path: + logger.error(f"Asset {context.asset_id}, FileRule {file_rule.id.hex} ({file_rule.map_type}): Source file not found with pattern '{file_rule.filename_pattern}' in '{source_base_path}'.") + self._update_file_rule_status(context, file_rule.id.hex, 'Failed', map_type=file_rule.map_type, details="Source file not found") + continue + + # B. Load and Transform Image + image_data: Optional[np.ndarray] = ipu.load_image(str(source_file_path)) + if image_data is None: + logger.error(f"Asset {context.asset_id}, FileRule {file_rule.id.hex} ({file_rule.map_type}): Failed to load image from '{source_file_path}'.") + self._update_file_rule_status(context, file_rule.id.hex, 'Failed', map_type=file_rule.map_type, source_file=str(source_file_path), details="Image load failed") + continue + + original_height, original_width = image_data.shape[:2] + logger.debug(f"Asset {context.asset_id}, FileRule {file_rule.id.hex}: Loaded image '{source_file_path}' with dimensions {original_width}x{original_height}.") + + transform: TransformSettings = file_rule.transform_settings + + target_width, target_height = ipu.calculate_target_dimensions( + original_width, original_height, + transform.target_width, transform.target_height, + transform.resize_mode, + transform.ensure_pot, + transform.allow_upscale + ) + logger.debug(f"Asset {context.asset_id}, FileRule {file_rule.id.hex}: Original dims: ({original_width},{original_height}), Calculated target dims: ({target_width},{target_height})") + + processed_image_data = image_data.copy() # Start with a copy + + if (target_width, target_height) != (original_width, original_height): + logger.info(f"Asset {context.asset_id}, FileRule {file_rule.id.hex}: Resizing from ({original_width},{original_height}) to ({target_width},{target_height}).") + # Map resize_filter string to cv2 interpolation constant + interpolation_map = { + "NEAREST": cv2.INTER_NEAREST, + "LINEAR": cv2.INTER_LINEAR, + "CUBIC": cv2.INTER_CUBIC, + "AREA": cv2.INTER_AREA, # Good for downscaling + "LANCZOS4": cv2.INTER_LANCZOS4 + } + interpolation = interpolation_map.get(transform.resize_filter.upper(), cv2.INTER_AREA) # Default to INTER_AREA + processed_image_data = ipu.resize_image(processed_image_data, target_width, target_height, interpolation=interpolation) + if processed_image_data is None: # Should not happen if resize_image handles errors, but good practice + logger.error(f"Asset {context.asset_id}, FileRule {file_rule.id.hex} ({file_rule.map_type}): Failed to resize image.") + self._update_file_rule_status(context, file_rule.id.hex, 'Failed', map_type=file_rule.map_type, source_file=str(source_file_path), original_dimensions=(original_width, original_height), details="Image resize failed") + continue + + + # Color Space Conversion (simplified) + # Assuming ipu.load_image loads as BGR if color. + # This needs more robust handling of source color profiles if they are known. + if transform.color_profile_management and transform.target_color_profile == "RGB": + if len(processed_image_data.shape) == 3 and processed_image_data.shape[2] == 3: # Check if it's a color image + logger.info(f"Asset {context.asset_id}, FileRule {file_rule.id.hex}: Converting BGR to RGB.") + processed_image_data = ipu.convert_bgr_to_rgb(processed_image_data) + elif len(processed_image_data.shape) == 3 and processed_image_data.shape[2] == 4: # Check for BGRA + logger.info(f"Asset {context.asset_id}, FileRule {file_rule.id.hex}: Converting BGRA to RGBA.") + processed_image_data = ipu.convert_bgra_to_rgba(processed_image_data) + + + # C. Save Temporary Processed Map + # Ensure engine_temp_dir exists (orchestrator should handle this, but good to be safe) + if not context.engine_temp_dir.exists(): + try: + context.engine_temp_dir.mkdir(parents=True, exist_ok=True) + logger.info(f"Asset {context.asset_id}: Created engine_temp_dir at '{context.engine_temp_dir}'") + except OSError as e: + logger.error(f"Asset {context.asset_id}: Failed to create engine_temp_dir '{context.engine_temp_dir}': {e}") + self._update_file_rule_status(context, file_rule.id.hex, 'Failed', map_type=file_rule.map_type, source_file=str(source_file_path), details="Failed to create temp directory") + continue # Or potentially fail the whole asset processing here + + temp_filename_suffix = Path(source_file_path).suffix + # Use a more descriptive name if possible, including map_type + safe_map_type = sanitize_filename(file_rule.map_type if file_rule.map_type else "unknown_map") + temp_output_filename = f"processed_{safe_map_type}_{file_rule.id.hex}{temp_filename_suffix}" + temp_output_path = context.engine_temp_dir / temp_output_filename + + # Consider output_format_settings from transform if they apply here + # For now, save_image handles basic saving. + # Example: cv2.imwrite params for quality for JPG, compression for PNG + save_params = [] + if transform.output_format_settings: + if temp_filename_suffix.lower() in ['.jpg', '.jpeg']: + quality = transform.output_format_settings.get('quality', 95) + save_params = [cv2.IMWRITE_JPEG_QUALITY, quality] + elif temp_filename_suffix.lower() == '.png': + compression = transform.output_format_settings.get('compression_level', 3) # 0-9, 3 is default + save_params = [cv2.IMWRITE_PNG_COMPRESSION, compression] + # Add more formats as needed (e.g., EXR, TIFF) + + save_success = ipu.save_image(str(temp_output_path), processed_image_data, params=save_params) + + if not save_success: + logger.error(f"Asset {context.asset_id}, FileRule {file_rule.id.hex} ({file_rule.map_type}): Failed to save temporary image to '{temp_output_path}'.") + self._update_file_rule_status( + context, file_rule.id.hex, 'Failed', + map_type=file_rule.map_type, + source_file=str(source_file_path), + original_dimensions=(original_width, original_height), + processed_dimensions=(processed_image_data.shape[1], processed_image_data.shape[0]) if processed_image_data is not None else None, + details="Temporary image save failed" + ) + continue + + logger.info(f"Asset {context.asset_id}, FileRule {file_rule.id.hex} ({file_rule.map_type}): Successfully processed and saved temporary map to '{temp_output_path}'.") + + # D. Update Context + self._update_file_rule_status( + context, file_rule.id.hex, 'Processed', + map_type=file_rule.map_type, + source_file=str(source_file_path), + temp_processed_file=str(temp_output_path), + original_dimensions=(original_width, original_height), + processed_dimensions=(processed_image_data.shape[1], processed_image_data.shape[0]), + details="Successfully processed" + ) + + # Optional: Update context.asset_metadata['processed_files'] + if 'processed_files' not in context.asset_metadata: + context.asset_metadata['processed_files'] = [] + context.asset_metadata['processed_files'].append({ + 'file_rule_id': file_rule.id.hex, + 'path': str(temp_output_path), + 'type': 'temporary_map', + 'map_type': file_rule.map_type + }) + + + logger.info(f"Asset {context.asset_id}: Finished individual map processing stage.") + return context + + def _find_source_file(self, base_path: Path, pattern: str, asset_id: str, file_rule_id_hex: str) -> Optional[Path]: + """ + Finds a single source file matching the pattern within the base_path. + Adapts logic from ProcessingEngine._find_source_file. + """ + if not pattern: + logger.warning(f"Asset {asset_id}, FileRule {file_rule_id_hex}: Empty filename pattern provided.") + return None + + try: + # Using rglob for potentially nested structures, though original might have been simpler. + # If pattern is exact filename, it will also work. + # If pattern is a glob, it will search. + matched_files = list(base_path.rglob(pattern)) + + if not matched_files: + logger.debug(f"Asset {asset_id}, FileRule {file_rule_id_hex}: No files found matching pattern '{pattern}' in '{base_path}' (recursive).") + # Try non-recursive if rglob fails and pattern might be for top-level + matched_files_non_recursive = list(base_path.glob(pattern)) + if matched_files_non_recursive: + logger.debug(f"Asset {asset_id}, FileRule {file_rule_id_hex}: Found {len(matched_files_non_recursive)} files non-recursively. Using first: {matched_files_non_recursive[0]}") + return matched_files_non_recursive[0] + return None + + if len(matched_files) > 1: + logger.warning(f"Asset {asset_id}, FileRule {file_rule_id_hex}: Multiple files ({len(matched_files)}) found for pattern '{pattern}' in '{base_path}'. Using the first one: {matched_files[0]}. Files: {matched_files}") + + return matched_files[0] + + except Exception as e: + logger.error(f"Asset {asset_id}, FileRule {file_rule_id_hex}: Error searching for file with pattern '{pattern}' in '{base_path}': {e}") + return None + + def _update_file_rule_status(self, context: AssetProcessingContext, file_rule_id_hex: str, status: str, **kwargs): + """Helper to update processed_maps_details for a file_rule.""" + if file_rule_id_hex not in context.processed_maps_details: + context.processed_maps_details[file_rule_id_hex] = {} + + context.processed_maps_details[file_rule_id_hex]['status'] = status + for key, value in kwargs.items(): + context.processed_maps_details[file_rule_id_hex][key] = value + + # Ensure essential keys are present even on failure, if known + if 'map_type' not in context.processed_maps_details[file_rule_id_hex] and 'map_type' in kwargs: + context.processed_maps_details[file_rule_id_hex]['map_type'] = kwargs['map_type'] + + + logger.debug(f"Asset {context.asset_id}, FileRule {file_rule_id_hex}: Status updated to '{status}'. Details: {kwargs}") \ No newline at end of file diff --git a/processing/pipeline/stages/map_merging.py b/processing/pipeline/stages/map_merging.py new file mode 100644 index 0000000..6e0fd0f --- /dev/null +++ b/processing/pipeline/stages/map_merging.py @@ -0,0 +1,310 @@ +import logging +from pathlib import Path +from typing import Dict, Optional, List, Tuple + +import numpy as np +import cv2 # For potential direct cv2 operations if ipu doesn't cover all merge needs + +from ..base_stage import ProcessingStage +from ...asset_context import AssetProcessingContext +from ....rule_structure import FileRule, MergeSettings, MergeInputChannel +from ....utils.path_utils import sanitize_filename +from ...utils import image_processing_utils as ipu + + +logger = logging.getLogger(__name__) + +class MapMergingStage(ProcessingStage): + """ + Merges individually processed maps based on MAP_MERGE rules. + This stage performs operations like channel packing. + """ + + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: + """ + Executes the map merging logic. + + Args: + context: The asset processing context. + + Returns: + The updated asset processing context. + """ + if context.status_flags.get('skip_asset'): + logger.info(f"Skipping map merging for asset {context.asset_name} as skip_asset flag is set.") + return context + + if not hasattr(context, 'merged_maps_details'): + context.merged_maps_details = {} + + if not hasattr(context, 'processed_maps_details'): + logger.warning(f"Asset {context.asset_name}: 'processed_maps_details' not found in context. Cannot perform map merging.") + return context + + if not context.files_to_process: + logger.info(f"Asset {context.asset_name}: No files_to_process defined. Skipping map merging.") + return context + + logger.info(f"Starting MapMergingStage for asset: {context.asset_name}") + + for merge_rule in context.files_to_process: + if not isinstance(merge_rule, FileRule) or merge_rule.item_type != "MAP_MERGE": + continue + + if not merge_rule.merge_settings: + logger.error(f"Asset {context.asset_name}, Rule ID {merge_rule.id.hex}: Merge rule for map_type '{merge_rule.map_type}' is missing merge_settings. Skipping this merge.") + context.merged_maps_details[merge_rule.id.hex] = { + 'map_type': merge_rule.map_type, + 'status': 'Failed', + 'reason': 'Missing merge_settings in FileRule.' + } + continue + + merge_settings: MergeSettings = merge_rule.merge_settings + output_map_type = merge_rule.map_type + rule_id_hex = merge_rule.id.hex + logger.info(f"Processing MAP_MERGE rule for '{output_map_type}' (ID: {rule_id_hex})") + + loaded_input_maps: Dict[str, np.ndarray] = {} + input_map_paths: Dict[str, str] = {} + target_dims: Optional[Tuple[int, int]] = None # width, height + all_inputs_valid = True + + # A. Load Input Maps for Merging + if not merge_settings.input_maps: + logger.warning(f"Asset {context.asset_name}, Rule ID {rule_id_hex}: No input_maps defined in merge_settings for '{output_map_type}'. Skipping this merge.") + context.merged_maps_details[rule_id_hex] = { + 'map_type': output_map_type, + 'status': 'Failed', + 'reason': 'No input_maps defined in merge_settings.' + } + continue + + for input_map_config in merge_settings.input_maps: + input_rule_id_hex = input_map_config.file_rule_id.hex + processed_detail = context.processed_maps_details.get(input_rule_id_hex) + + if not processed_detail or processed_detail.get('status') != 'Processed': + error_msg = f"Input map (Rule ID: {input_rule_id_hex}) for merge rule '{output_map_type}' (Rule ID: {rule_id_hex}) not found or not processed. Details: {processed_detail}" + logger.error(error_msg) + all_inputs_valid = False + context.merged_maps_details[rule_id_hex] = { + 'map_type': output_map_type, + 'status': 'Failed', + 'reason': f"Input map {input_rule_id_hex} not processed or missing." + } + break + + temp_processed_file_path = Path(processed_detail['temp_processed_file']) + if not temp_processed_file_path.exists(): + error_msg = f"Input map file {temp_processed_file_path} for merge rule '{output_map_type}' (Rule ID: {rule_id_hex}) does not exist." + logger.error(error_msg) + all_inputs_valid = False + context.merged_maps_details[rule_id_hex] = { + 'map_type': output_map_type, + 'status': 'Failed', + 'reason': f"Input map file {temp_processed_file_path} not found." + } + break + + try: + image_data = ipu.load_image(temp_processed_file_path) + except Exception as e: + logger.error(f"Error loading image {temp_processed_file_path} for merge rule '{output_map_type}' (Rule ID: {rule_id_hex}): {e}") + all_inputs_valid = False + context.merged_maps_details[rule_id_hex] = { + 'map_type': output_map_type, + 'status': 'Failed', + 'reason': f"Error loading input image {temp_processed_file_path}." + } + break + + if image_data is None: + logger.error(f"Failed to load image data from {temp_processed_file_path} for merge rule '{output_map_type}' (Rule ID: {rule_id_hex}).") + all_inputs_valid = False + context.merged_maps_details[rule_id_hex] = { + 'map_type': output_map_type, + 'status': 'Failed', + 'reason': f"Failed to load image data from {temp_processed_file_path}." + } + break + + loaded_input_maps[input_rule_id_hex] = image_data + input_map_paths[input_rule_id_hex] = str(temp_processed_file_path) + + current_dims = (image_data.shape[1], image_data.shape[0]) # width, height + if target_dims is None: + target_dims = current_dims + logger.debug(f"Merge rule '{output_map_type}' (ID: {rule_id_hex}): Set target dimensions to {target_dims} from first input {temp_processed_file_path}.") + elif current_dims != target_dims: + logger.warning(f"Input map {temp_processed_file_path} for merge rule '{output_map_type}' (ID: {rule_id_hex}) has dimensions {current_dims}, but target is {target_dims}. Resizing.") + try: + image_data = ipu.resize_image(image_data, target_dims[0], target_dims[1]) + if image_data is None: + raise ValueError("Resize operation returned None.") + loaded_input_maps[input_rule_id_hex] = image_data + except Exception as e: + logger.error(f"Failed to resize image {temp_processed_file_path} for merge rule '{output_map_type}' (ID: {rule_id_hex}): {e}") + all_inputs_valid = False + context.merged_maps_details[rule_id_hex] = { + 'map_type': output_map_type, + 'status': 'Failed', + 'reason': f"Failed to resize input image {temp_processed_file_path}." + } + break + + if not all_inputs_valid: + # Failure already logged and recorded in context.merged_maps_details + logger.warning(f"Skipping merge for '{output_map_type}' (ID: {rule_id_hex}) due to invalid inputs.") + continue + + if target_dims is None: # Should not happen if all_inputs_valid is true and there was at least one input map + logger.error(f"Merge rule '{output_map_type}' (ID: {rule_id_hex}): Target dimensions not determined despite valid inputs. This indicates an issue with input map loading or an empty input_maps list that wasn't caught.") + context.merged_maps_details[rule_id_hex] = { + 'map_type': output_map_type, + 'status': 'Failed', + 'reason': 'Target dimensions could not be determined.' + } + continue + + # B. Perform Merge Operation + try: + if merge_settings.output_channels == 1: + merged_image = np.zeros((target_dims[1], target_dims[0]), dtype=np.uint8) + else: + merged_image = np.zeros((target_dims[1], target_dims[0], merge_settings.output_channels), dtype=np.uint8) + except Exception as e: + logger.error(f"Error creating empty merged image for '{output_map_type}' (ID: {rule_id_hex}) with dims {target_dims} and {merge_settings.output_channels} channels: {e}") + context.merged_maps_details[rule_id_hex] = { + 'map_type': output_map_type, + 'status': 'Failed', + 'reason': f'Error creating output image canvas: {e}' + } + continue + + merge_op_failed = False + for input_map_config in merge_settings.input_maps: + source_image = loaded_input_maps[input_map_config.file_rule_id.hex] + source_channel_index = input_map_config.source_channel + target_channel_index = input_map_config.target_channel + + source_data = None + if source_image.ndim == 2: # Grayscale + source_data = source_image + elif source_image.ndim == 3: # Multi-channel (e.g. RGB, RGBA) + if source_channel_index >= source_image.shape[2]: + logger.error(f"Merge rule '{output_map_type}' (ID: {rule_id_hex}): Source channel index {source_channel_index} out of bounds for source image with shape {source_image.shape} (from Rule ID {input_map_config.file_rule_id.hex}).") + merge_op_failed = True + break + source_data = source_image[:, :, source_channel_index] + else: + logger.error(f"Merge rule '{output_map_type}' (ID: {rule_id_hex}): Source image (from Rule ID {input_map_config.file_rule_id.hex}) has unexpected dimensions: {source_image.ndim}. Shape: {source_image.shape}") + merge_op_failed = True + break + + if source_data is None: # Should be caught by previous checks + logger.error(f"Merge rule '{output_map_type}' (ID: {rule_id_hex}): Failed to extract source_data for unknown reasons from input {input_map_config.file_rule_id.hex}.") + merge_op_failed = True + break + + # Assign to target channel + try: + if merged_image.ndim == 2: # Output is grayscale + if merge_settings.output_channels != 1: + logger.error(f"Merge rule '{output_map_type}' (ID: {rule_id_hex}): Mismatch - merged_image is 2D but output_channels is {merge_settings.output_channels}.") + merge_op_failed = True + break + merged_image = source_data # Overwrites if multiple inputs map to grayscale; consider blending or specific logic if needed + elif merged_image.ndim == 3: # Output is multi-channel + if target_channel_index >= merged_image.shape[2]: + logger.error(f"Merge rule '{output_map_type}' (ID: {rule_id_hex}): Target channel index {target_channel_index} out of bounds for merged image with shape {merged_image.shape}.") + merge_op_failed = True + break + merged_image[:, :, target_channel_index] = source_data + else: # Should not happen + logger.error(f"Merge rule '{output_map_type}' (ID: {rule_id_hex}): Merged image has unexpected dimensions: {merged_image.ndim}. Shape: {merged_image.shape}") + merge_op_failed = True + break + except Exception as e: + logger.error(f"Error assigning source data to target channel for '{output_map_type}' (ID: {rule_id_hex}): {e}. Source shape: {source_data.shape}, Target channel: {target_channel_index}, Merged image shape: {merged_image.shape}") + merge_op_failed = True + break + + if input_map_config.invert_source_channel: + if merged_image.ndim == 2: + merged_image = 255 - merged_image # Assumes uint8 + elif merged_image.ndim == 3: + # Ensure we are not inverting an alpha channel if that's not desired, + # but current spec inverts the target channel data. + merged_image[:, :, target_channel_index] = 255 - merged_image[:, :, target_channel_index] + + # input_map_config.default_value_if_missing: + # This was handled by all_inputs_valid check for file presence. + # If a channel is missing from a multi-channel source, that's an error in source_channel_index. + # If a file is entirely missing and a default color/value is needed for the *output channel*, + # that would be a different logic, perhaps pre-filling merged_image. + # For now, we assume if an input map is specified, it must be present and valid. + + if merge_op_failed: + logger.error(f"Merge operation failed for '{output_map_type}' (ID: {rule_id_hex}).") + context.merged_maps_details[rule_id_hex] = { + 'map_type': output_map_type, + 'status': 'Failed', + 'reason': 'Error during channel packing/merge operation.' + } + continue + + # C. Save Temporary Merged Map + # Default to PNG, or use format from merge_settings if available (future enhancement) + output_format = getattr(merge_settings, 'output_format', 'png').lower() + if output_format not in ['png', 'jpg', 'jpeg', 'tif', 'tiff', 'exr']: # Add more as ipu supports + logger.warning(f"Unsupported output_format '{output_format}' in merge_settings for '{output_map_type}' (ID: {rule_id_hex}). Defaulting to PNG.") + output_format = 'png' + + temp_merged_filename = f"merged_{sanitize_filename(output_map_type)}_{rule_id_hex}.{output_format}" + + if not context.engine_temp_dir: + logger.error(f"Asset {context.asset_name}: engine_temp_dir is not set. Cannot save merged map.") + context.merged_maps_details[rule_id_hex] = { + 'map_type': output_map_type, + 'status': 'Failed', + 'reason': 'engine_temp_dir not set in context.' + } + continue + + temp_merged_path = context.engine_temp_dir / temp_merged_filename + + try: + save_success = ipu.save_image(temp_merged_path, merged_image) + except Exception as e: + logger.error(f"Error saving merged image {temp_merged_path} for '{output_map_type}' (ID: {rule_id_hex}): {e}") + save_success = False + + if not save_success: + logger.error(f"Failed to save temporary merged map to {temp_merged_path} for '{output_map_type}' (ID: {rule_id_hex}).") + context.merged_maps_details[rule_id_hex] = { + 'map_type': output_map_type, + 'status': 'Failed', + 'reason': f'Failed to save merged image to {temp_merged_path}.' + } + continue + + logger.info(f"Successfully merged and saved '{output_map_type}' (ID: {rule_id_hex}) to {temp_merged_path}") + + # D. Update Context + context.merged_maps_details[rule_id_hex] = { + 'map_type': output_map_type, + 'temp_merged_file': str(temp_merged_path), + 'input_map_ids_used': [mc.file_rule_id.hex for mc in merge_settings.input_maps], + 'input_map_files_used': input_map_paths, # Dict[rule_id_hex, path_str] + 'merged_dimensions': target_dims, # (width, height) + 'status': 'Processed', + 'file_rule_id': rule_id_hex # For easier reverse lookup if needed + } + + # Optional: Update context.asset_metadata['processed_files'] or similar + # This might be better handled by a later stage that finalizes files. + # For now, merged_maps_details is the primary record. + + logger.info(f"Finished MapMergingStage for asset: {context.asset_name}. Merged maps: {len(context.merged_maps_details)}") + return context \ No newline at end of file diff --git a/processing/pipeline/stages/metadata_finalization_save.py b/processing/pipeline/stages/metadata_finalization_save.py new file mode 100644 index 0000000..d18bfc4 --- /dev/null +++ b/processing/pipeline/stages/metadata_finalization_save.py @@ -0,0 +1,119 @@ +import datetime +import json +import logging +from pathlib import Path +from typing import Any, Dict + +from ..asset_context import AssetProcessingContext +from .base_stage import ProcessingStage +from ....utils.path_utils import generate_path_from_pattern + + +logger = logging.getLogger(__name__) + +class MetadataFinalizationAndSaveStage(ProcessingStage): + """ + This stage finalizes the asset_metadata (e.g., setting processing end time, + final status) and saves it as a JSON file. + """ + + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: + """ + Finalizes metadata, determines output path, and saves the metadata JSON file. + """ + if not hasattr(context, 'asset_metadata') or not context.asset_metadata: + if context.status_flags.get('skip_asset'): + logger.info( + f"Asset '{context.asset_rule.name if hasattr(context, 'asset_rule') and context.asset_rule else 'Unknown'}': " + f"Skipped before metadata initialization. No metadata file will be saved." + ) + else: + logger.warning( + f"Asset '{context.asset_rule.name if hasattr(context, 'asset_rule') and context.asset_rule else 'Unknown'}': " + f"asset_metadata not initialized. Skipping metadata finalization and save." + ) + return context + + # Check Skip Flag + if context.status_flags.get('skip_asset'): + context.asset_metadata['status'] = "Skipped" + context.asset_metadata['processing_end_time'] = datetime.datetime.now().isoformat() + context.asset_metadata['notes'] = context.status_flags.get('skip_reason', 'Skipped early in pipeline') + logger.info( + f"Asset '{context.asset_rule.name}': Marked as skipped. Reason: {context.asset_metadata['notes']}" + ) + # Assuming we save metadata for skipped assets if it was initialized. + # If not, the logic to skip saving would be here or before path generation. + + # A. Finalize Metadata + context.asset_metadata['processing_end_time'] = datetime.datetime.now().isoformat() + + # Determine final status (if not already set to Skipped) + if context.asset_metadata.get('status') != "Skipped": + has_errors = any( + context.status_flags.get(error_flag) + for error_flag in ['file_processing_error', 'merge_error', 'critical_error'] # Added critical_error + ) + if has_errors: + context.asset_metadata['status'] = "Failed" + else: + context.asset_metadata['status'] = "Processed" + + # Add details of processed and merged maps + context.asset_metadata['processed_map_details'] = getattr(context, 'processed_maps_details', {}) + context.asset_metadata['merged_map_details'] = getattr(context, 'merged_maps_details', {}) + + # (Optional) Add a list of all temporary files + context.asset_metadata['temporary_files'] = getattr(context, 'temporary_files', []) + + # B. Determine Metadata Output Path + # Ensure asset_rule and source_rule exist before accessing their names + asset_name = context.asset_rule.name if hasattr(context, 'asset_rule') and context.asset_rule else "unknown_asset" + source_rule_name = context.source_rule.name if hasattr(context, 'source_rule') and context.source_rule else "unknown_source" + + metadata_filename = f"{asset_name}_metadata.json" + output_path_pattern = context.asset_rule.output_path_pattern if hasattr(context, 'asset_rule') and context.asset_rule else "" + + # Handle potential missing sha5_value, defaulting to None or an empty string + sha_value = getattr(context, 'sha5_value', getattr(context, 'sha_value', None)) + + + full_output_path = generate_path_from_pattern( + base_path=str(context.output_base_path), # Ensure base_path is a string + pattern=output_path_pattern, + asset_name=asset_name, + map_type="metadata", # Special map_type for metadata + filename=metadata_filename, + source_rule_name=source_rule_name, + incrementing_value=getattr(context, 'incrementing_value', None), + sha_value=sha_value # Changed from sha5_value to sha_value for more generality + ) + metadata_save_path = Path(full_output_path) + + # C. Save Metadata File + try: + metadata_save_path.parent.mkdir(parents=True, exist_ok=True) + + def make_serializable(data: Any) -> Any: + if isinstance(data, Path): + return str(data) + if isinstance(data, datetime.datetime): # Ensure datetime is serializable + return data.isoformat() + if isinstance(data, dict): + return {k: make_serializable(v) for k, v in data.items()} + if isinstance(data, list): + return [make_serializable(i) for i in data] + return data + + serializable_metadata = make_serializable(context.asset_metadata) + + with open(metadata_save_path, 'w') as f: + json.dump(serializable_metadata, f, indent=4) + logger.info(f"Asset '{asset_name}': Metadata saved to {metadata_save_path}") + context.asset_metadata['metadata_file_path'] = str(metadata_save_path) + except Exception as e: + logger.error(f"Asset '{asset_name}': Failed to save metadata to {metadata_save_path}. Error: {e}") + context.asset_metadata['status'] = "Failed (Metadata Save Error)" + context.status_flags['metadata_save_error'] = True + + return context \ No newline at end of file diff --git a/processing/pipeline/stages/metadata_initialization.py b/processing/pipeline/stages/metadata_initialization.py new file mode 100644 index 0000000..4d5fbf5 --- /dev/null +++ b/processing/pipeline/stages/metadata_initialization.py @@ -0,0 +1,163 @@ +import datetime +import logging + +from ..base_stage import ProcessingStage +from ...asset_context import AssetProcessingContext # Adjusted import path assuming asset_context is in processing.pipeline +# If AssetProcessingContext is directly under 'processing', the import would be: +# from ...asset_context import AssetProcessingContext +# Based on the provided file structure, asset_context.py is in processing/pipeline/ +# So, from ...asset_context import AssetProcessingContext is likely incorrect. +# It should be: from ..asset_context import AssetProcessingContext +# Correcting this based on typical Python package structure and the location of base_stage.py + +# Re-evaluating import based on common structure: +# If base_stage.py is in processing/pipeline/stages/ +# and asset_context.py is in processing/pipeline/ +# then the import for AssetProcessingContext from metadata_initialization.py (in stages) would be: +# from ..asset_context import AssetProcessingContext + +# Let's assume the following structure for clarity: +# processing/ +# L-- pipeline/ +# L-- __init__.py +# L-- asset_context.py +# L-- base_stage.py (Mistake here, base_stage is in stages, so it's ..base_stage) +# L-- stages/ +# L-- __init__.py +# L-- metadata_initialization.py +# L-- base_stage.py (Corrected: base_stage.py is here) + +# Corrected imports based on the plan and typical structure: +# base_stage.py is in processing/pipeline/stages/ +# asset_context.py is in processing/pipeline/ + +# from ..base_stage import ProcessingStage # This would mean base_stage is one level up from stages (i.e. in pipeline) +# The plan says: from ..base_stage import ProcessingStage +# This implies that metadata_initialization.py is in a subdirectory of where base_stage.py is. +# However, the file path for metadata_initialization.py is processing/pipeline/stages/metadata_initialization.py +# And base_stage.py is listed as processing/pipeline/stages/base_stage.py in the open tabs. +# So, the import should be: +# from .base_stage import ProcessingStage + +# AssetProcessingContext is at processing/pipeline/asset_context.py +# So from processing/pipeline/stages/metadata_initialization.py, it would be: +# from ..asset_context import AssetProcessingContext + +# Final check on imports based on instructions: +# `from ..base_stage import ProcessingStage` -> This means base_stage.py is in `processing/pipeline/` +# `from ...asset_context import AssetProcessingContext` -> This means asset_context.py is in `processing/` +# Let's verify the location of these files from the environment details. +# processing/pipeline/asset_context.py +# processing/pipeline/stages/base_stage.py +# +# So, from processing/pipeline/stages/metadata_initialization.py: +# To import ProcessingStage from processing/pipeline/stages/base_stage.py: +# from .base_stage import ProcessingStage +# To import AssetProcessingContext from processing/pipeline/asset_context.py: +# from ..asset_context import AssetProcessingContext + +# The instructions explicitly state: +# `from ..base_stage import ProcessingStage` +# `from ...asset_context import AssetProcessingContext` +# This implies a different structure than what seems to be in the file tree. +# I will follow the explicit import instructions from the task. +# This means: +# base_stage.py is expected at `processing/pipeline/base_stage.py` +# asset_context.py is expected at `processing/asset_context.py` + +# Given the file tree: +# processing/pipeline/asset_context.py +# processing/pipeline/stages/base_stage.py +# The imports in `processing/pipeline/stages/metadata_initialization.py` should be: +# from .base_stage import ProcessingStage +# from ..asset_context import AssetProcessingContext + +# I will use the imports that align with the provided file structure. + +from .base_stage import ProcessingStage +from ..asset_context import AssetProcessingContext + + +logger = logging.getLogger(__name__) + +class MetadataInitializationStage(ProcessingStage): + """ + Initializes metadata structures within the AssetProcessingContext. + This stage sets up asset_metadata, processed_maps_details, and + merged_maps_details. + """ + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: + """ + Executes the metadata initialization logic. + + Args: + context: The AssetProcessingContext for the current asset. + + Returns: + The modified AssetProcessingContext. + """ + if context.status_flags.get('skip_asset', False): + logger.debug(f"Asset '{context.asset_rule.name if context.asset_rule else 'Unknown'}': Skipping metadata initialization as 'skip_asset' is True.") + return context + + logger.debug(f"Asset '{context.asset_rule.name}': Initializing metadata.") + + context.asset_metadata = {} + context.processed_maps_details = {} + context.merged_maps_details = {} + + # Populate Initial asset_metadata + if context.asset_rule: + context.asset_metadata['asset_name'] = context.asset_rule.name + context.asset_metadata['asset_id'] = str(context.asset_rule.id) + context.asset_metadata['source_path'] = str(context.asset_rule.source_path) + context.asset_metadata['output_path_pattern'] = context.asset_rule.output_path_pattern + context.asset_metadata['tags'] = list(context.asset_rule.tags) if context.asset_rule.tags else [] + context.asset_metadata['custom_fields'] = dict(context.asset_rule.custom_fields) if context.asset_rule.custom_fields else {} + else: + # Handle cases where asset_rule might be None, though typically it should be set + logger.warning("AssetRule is not set in context during metadata initialization.") + context.asset_metadata['asset_name'] = "Unknown Asset" + context.asset_metadata['asset_id'] = "N/A" + context.asset_metadata['source_path'] = "N/A" + context.asset_metadata['output_path_pattern'] = "N/A" + context.asset_metadata['tags'] = [] + context.asset_metadata['custom_fields'] = {} + + + if context.source_rule: + context.asset_metadata['source_rule_name'] = context.source_rule.name + context.asset_metadata['source_rule_id'] = str(context.source_rule.id) + else: + logger.warning("SourceRule is not set in context during metadata initialization.") + context.asset_metadata['source_rule_name'] = "Unknown Source Rule" + context.asset_metadata['source_rule_id'] = "N/A" + + context.asset_metadata['effective_supplier'] = context.effective_supplier + 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 + 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 + + if context.incrementing_value is not None: + context.asset_metadata['incrementing_value'] = context.incrementing_value + + # The plan mentions sha5_value, which is likely a typo for sha256 or similar. + # Implementing as 'sha5_value' per instructions, but noting the potential typo. + if hasattr(context, 'sha5_value') and context.sha5_value is not None: # Check attribute existence + context.asset_metadata['sha5_value'] = context.sha5_value + elif hasattr(context, 'sha256_value') and context.sha256_value is not None: # Fallback if sha5 was a typo + logger.debug("sha5_value not found, using sha256_value if available for metadata.") + context.asset_metadata['sha256_value'] = context.sha256_value + + + logger.info(f"Asset '{context.asset_metadata.get('asset_name', 'Unknown')}': Metadata initialized.") + # Example of how you might log the full metadata for debugging: + # logger.debug(f"Initialized metadata: {context.asset_metadata}") + + return context \ No newline at end of file diff --git a/processing/pipeline/stages/normal_map_green_channel.py b/processing/pipeline/stages/normal_map_green_channel.py new file mode 100644 index 0000000..ca7984b --- /dev/null +++ b/processing/pipeline/stages/normal_map_green_channel.py @@ -0,0 +1,154 @@ +import logging +import numpy as np +from pathlib import Path +from typing import List + +from ..base_stage import ProcessingStage +from ...asset_context import AssetProcessingContext +from .....rule_structure import FileRule +from ...utils import image_processing_utils as ipu +from .....utils.path_utils import sanitize_filename + +logger = logging.getLogger(__name__) + +class NormalMapGreenChannelStage(ProcessingStage): + """ + Processing stage to invert the green channel of normal maps if configured. + This is often needed when converting between DirectX (Y-) and OpenGL (Y+) normal map formats. + """ + + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: + """ + Identifies NORMAL maps, checks configuration for green channel inversion, + performs inversion if needed, saves a new temporary file, and updates + the AssetProcessingContext. + """ + if context.status_flags.get('skip_asset'): + logger.debug(f"Asset '{context.asset_rule.name}': Skipping NormalMapGreenChannelStage due to skip_asset flag.") + return context + + if not context.files_to_process or not context.processed_maps_details: + logger.debug( + f"Asset '{context.asset_rule.name}': No files to process or processed_maps_details empty in NormalMapGreenChannelStage. Skipping." + ) + return context + + new_files_to_process: List[FileRule] = [] + processed_a_normal_map = False + + for file_rule in context.files_to_process: + if file_rule.map_type == "NORMAL": + # Check configuration for inversion + # Assuming a global setting for now. + # This key should exist in the Configuration object's general_settings. + should_invert = context.config_obj.general_settings.get('invert_normal_map_green_channel_globally', False) + + if not should_invert: + logger.debug( + f"Asset '{context.asset_rule.name}': Normal map green channel inversion not enabled globally. " + f"Skipping for {file_rule.filename_pattern} (ID: {file_rule.id.hex})." + ) + new_files_to_process.append(file_rule) + continue + + # Get the temporary processed file path + map_details = context.processed_maps_details.get(file_rule.id.hex) + if not map_details or map_details.get('status') != 'Processed' or not map_details.get('temp_processed_file'): + logger.warning( + f"Asset '{context.asset_rule.name}': Normal map {file_rule.filename_pattern} (ID: {file_rule.id.hex}) " + f"not found in processed_maps_details or not marked as 'Processed'. Cannot invert green channel." + ) + new_files_to_process.append(file_rule) + continue + + original_temp_path = Path(map_details['temp_processed_file']) + if not original_temp_path.exists(): + logger.error( + f"Asset '{context.asset_rule.name}': Temporary file {original_temp_path} for normal map " + f"{file_rule.filename_pattern} (ID: {file_rule.id.hex}) does not exist. Cannot invert green channel." + ) + new_files_to_process.append(file_rule) + continue + + image_data = ipu.load_image(original_temp_path) + + if image_data is None: + logger.error( + f"Asset '{context.asset_rule.name}': Failed to load image from {original_temp_path} " + f"for normal map {file_rule.filename_pattern} (ID: {file_rule.id.hex})." + ) + new_files_to_process.append(file_rule) + continue + + if image_data.ndim != 3 or image_data.shape[2] < 2: # Must have at least R, G channels + logger.error( + f"Asset '{context.asset_rule.name}': Image {original_temp_path} for normal map " + f"{file_rule.filename_pattern} (ID: {file_rule.id.hex}) is not a valid RGB/normal map " + f"(ndim={image_data.ndim}, channels={image_data.shape[2] if image_data.ndim == 3 else 'N/A'}) " + f"for green channel inversion." + ) + new_files_to_process.append(file_rule) + continue + + # Perform Green Channel Inversion + modified_image_data = image_data.copy() + try: + if np.issubdtype(modified_image_data.dtype, np.floating): + modified_image_data[:, :, 1] = 1.0 - modified_image_data[:, :, 1] + elif np.issubdtype(modified_image_data.dtype, np.integer): + max_val = np.iinfo(modified_image_data.dtype).max + modified_image_data[:, :, 1] = max_val - modified_image_data[:, :, 1] + else: + logger.error( + f"Asset '{context.asset_rule.name}': Unsupported image data type " + f"{modified_image_data.dtype} for normal map {original_temp_path}. Cannot invert green channel." + ) + new_files_to_process.append(file_rule) + continue + except IndexError: + logger.error( + f"Asset '{context.asset_rule.name}': Image {original_temp_path} for normal map " + f"{file_rule.filename_pattern} (ID: {file_rule.id.hex}) does not have a green channel (index 1) " + f"or has unexpected dimensions ({modified_image_data.shape}). Cannot invert." + ) + new_files_to_process.append(file_rule) + continue + + + # Save New Temporary (Modified Normal) Map + new_temp_filename = f"normal_g_inv_{sanitize_filename(file_rule.map_type)}_{file_rule.id.hex}{original_temp_path.suffix}" + new_temp_path = context.engine_temp_dir / new_temp_filename + + save_success = ipu.save_image(new_temp_path, modified_image_data) + + if save_success: + logger.info( + f"Asset '{context.asset_rule.name}': Inverted green channel for NORMAL map " + f"{original_temp_path.name}, saved to {new_temp_path.name}." + ) + # Update processed_maps_details + context.processed_maps_details[file_rule.id.hex]['temp_processed_file'] = str(new_temp_path) + current_notes = context.processed_maps_details[file_rule.id.hex].get('notes', '') + context.processed_maps_details[file_rule.id.hex]['notes'] = \ + f"{current_notes}; Green channel inverted by NormalMapGreenChannelStage".strip('; ') + + new_files_to_process.append(file_rule) # Add original rule, it now points to modified data + processed_a_normal_map = True + else: + logger.error( + f"Asset '{context.asset_rule.name}': Failed to save inverted normal map to {new_temp_path} " + f"for original {original_temp_path.name}." + ) + new_files_to_process.append(file_rule) # Add original rule, as processing failed + else: + # Not a normal map, just pass it through + new_files_to_process.append(file_rule) + + context.files_to_process = new_files_to_process + + if processed_a_normal_map: + logger.info(f"Asset '{context.asset_rule.name}': NormalMapGreenChannelStage processed relevant normal maps.") + else: + logger.debug(f"Asset '{context.asset_rule.name}': No normal maps found or processed in NormalMapGreenChannelStage.") + + return context \ No newline at end of file diff --git a/processing/pipeline/stages/output_organization.py b/processing/pipeline/stages/output_organization.py new file mode 100644 index 0000000..358d3e1 --- /dev/null +++ b/processing/pipeline/stages/output_organization.py @@ -0,0 +1,155 @@ +import logging +import shutil +from pathlib import Path +from typing import List, Dict, Optional + +from ..base_stage import ProcessingStage +from ...asset_context import AssetProcessingContext +from ....utils.path_utils import generate_path_from_pattern, sanitize_filename +from ....config import FileRule, MergeRule # Assuming these are needed for type hints if not directly in context + + +logger = logging.getLogger(__name__) + +class OutputOrganizationStage(ProcessingStage): + """ + Organizes output files by copying temporary processed files to their final destinations. + """ + + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: + """ + Copies temporary processed and merged files to their final output locations + based on path patterns and updates AssetProcessingContext. + """ + logger.debug(f"Asset '{context.asset_rule.name}': Starting output organization stage.") + + if context.status_flags.get('skip_asset'): + logger.info(f"Asset '{context.asset_rule.name}': Output organization skipped as 'skip_asset' is True.") + return context + + current_status = context.asset_metadata.get('status', '') + if current_status.startswith("Failed") or current_status == "Skipped": + logger.info(f"Asset '{context.asset_rule.name}': Output organization skipped due to prior status: {current_status}.") + return context + + final_output_files: List[str] = [] + # Ensure config_obj and general_settings are present, provide default for overwrite_existing if not + overwrite_existing = False + if context.config_obj and hasattr(context.config_obj, 'general_settings'): + overwrite_existing = context.config_obj.general_settings.overwrite_existing + else: + logger.warning(f"Asset '{context.asset_rule.name}': config_obj.general_settings not found, defaulting overwrite_existing to False.") + + + # A. Organize Processed Individual Maps + if context.processed_maps_details: + logger.debug(f"Asset '{context.asset_rule.name}': Organizing {len(context.processed_maps_details)} processed individual map(s).") + for file_rule_id, details in context.processed_maps_details.items(): + if details.get('status') != 'Processed' or not details.get('temp_processed_file'): + logger.debug(f"Asset '{context.asset_rule.name}': Skipping file_rule_id '{file_rule_id}' due to status '{details.get('status')}' or missing temp file.") + continue + + temp_file_path = Path(details['temp_processed_file']) + map_type = details['map_type'] + + output_filename = f"{context.asset_rule.name}_{sanitize_filename(map_type)}{temp_file_path.suffix}" + if context.asset_rule and context.asset_rule.file_rules: + current_file_rule: Optional[FileRule] = next( + (fr for fr in context.asset_rule.file_rules if fr.id == file_rule_id), None + ) + if current_file_rule and current_file_rule.output_filename_pattern: + output_filename = current_file_rule.output_filename_pattern + + try: + final_path_str = generate_path_from_pattern( + base_path=str(context.output_base_path), + pattern=context.asset_rule.output_path_pattern, + asset_name=context.asset_rule.name, + map_type=map_type, + filename=output_filename, + source_rule_name=context.source_rule.name if context.source_rule else "DefaultSource", + incrementing_value=str(context.incrementing_value) if context.incrementing_value is not None else None, + sha5_value=context.sha5_value + ) + final_path = Path(final_path_str) + final_path.parent.mkdir(parents=True, exist_ok=True) + + if final_path.exists() and not overwrite_existing: + logger.info(f"Asset '{context.asset_rule.name}': Output file {final_path} exists and overwrite is disabled. Skipping copy.") + else: + shutil.copy2(temp_file_path, final_path) + logger.info(f"Asset '{context.asset_rule.name}': Copied {temp_file_path} to {final_path}") + final_output_files.append(str(final_path)) + + context.processed_maps_details[file_rule_id]['final_output_path'] = str(final_path) + context.processed_maps_details[file_rule_id]['status'] = 'Organized' # Or some other status indicating completion + + except Exception as e: + logger.error(f"Asset '{context.asset_rule.name}': Failed to copy {temp_file_path} to {final_path_str if 'final_path_str' in locals() else 'unknown destination'} for file_rule_id '{file_rule_id}'. Error: {e}", exc_info=True) + context.status_flags['output_organization_error'] = True + context.asset_metadata['status'] = "Failed (Output Organization Error)" + # Optionally update status in details as well + context.processed_maps_details[file_rule_id]['status'] = 'Organization Failed' + else: + logger.debug(f"Asset '{context.asset_rule.name}': No processed individual maps to organize.") + + # B. Organize Merged Maps + if context.merged_maps_details: + logger.debug(f"Asset '{context.asset_rule.name}': Organizing {len(context.merged_maps_details)} merged map(s).") + for merge_rule_id, details in context.merged_maps_details.items(): + if details.get('status') != 'Processed' or not details.get('temp_merged_file'): + logger.debug(f"Asset '{context.asset_rule.name}': Skipping merge_rule_id '{merge_rule_id}' due to status '{details.get('status')}' or missing temp file.") + continue + + temp_file_path = Path(details['temp_merged_file']) + map_type = details['map_type'] # This is the output_map_type of the merge rule + + output_filename = f"{context.asset_rule.name}_{sanitize_filename(map_type)}{temp_file_path.suffix}" + if context.asset_rule and context.asset_rule.merge_rules: + current_merge_rule: Optional[MergeRule] = next( + (mr for mr in context.asset_rule.merge_rules if mr.id == merge_rule_id), None + ) + if current_merge_rule and current_merge_rule.output_filename_pattern: + output_filename = current_merge_rule.output_filename_pattern + + try: + final_path_str = generate_path_from_pattern( + base_path=str(context.output_base_path), + pattern=context.asset_rule.output_path_pattern, + asset_name=context.asset_rule.name, + map_type=map_type, + filename=output_filename, + source_rule_name=context.source_rule.name if context.source_rule else "DefaultSource", + incrementing_value=str(context.incrementing_value) if context.incrementing_value is not None else None, + sha5_value=context.sha5_value + ) + final_path = Path(final_path_str) + final_path.parent.mkdir(parents=True, exist_ok=True) + + if final_path.exists() and not overwrite_existing: + logger.info(f"Asset '{context.asset_rule.name}': Output file {final_path} exists and overwrite is disabled. Skipping copy for merged map.") + else: + shutil.copy2(temp_file_path, final_path) + logger.info(f"Asset '{context.asset_rule.name}': Copied merged map {temp_file_path} to {final_path}") + final_output_files.append(str(final_path)) + + context.merged_maps_details[merge_rule_id]['final_output_path'] = str(final_path) + context.merged_maps_details[merge_rule_id]['status'] = 'Organized' + + except Exception as e: + logger.error(f"Asset '{context.asset_rule.name}': Failed to copy merged map {temp_file_path} to {final_path_str if 'final_path_str' in locals() else 'unknown destination'} for merge_rule_id '{merge_rule_id}'. Error: {e}", exc_info=True) + context.status_flags['output_organization_error'] = True + context.asset_metadata['status'] = "Failed (Output Organization Error)" + context.merged_maps_details[merge_rule_id]['status'] = 'Organization Failed' + else: + logger.debug(f"Asset '{context.asset_rule.name}': No merged maps to organize.") + + context.asset_metadata['final_output_files'] = final_output_files + + if context.status_flags.get('output_organization_error'): + logger.error(f"Asset '{context.asset_rule.name}': Output organization encountered errors. Status: {context.asset_metadata['status']}") + else: + logger.info(f"Asset '{context.asset_rule.name}': Output organization complete. {len(final_output_files)} files placed.") + + logger.debug(f"Asset '{context.asset_rule.name}': Output organization stage finished.") + return context \ No newline at end of file diff --git a/processing/pipeline/stages/supplier_determination.py b/processing/pipeline/stages/supplier_determination.py new file mode 100644 index 0000000..ff60722 --- /dev/null +++ b/processing/pipeline/stages/supplier_determination.py @@ -0,0 +1,61 @@ +import logging + +from .base_stage import ProcessingStage +from ..asset_context import AssetProcessingContext + +class SupplierDeterminationStage(ProcessingStage): + """ + Determines the effective supplier for an asset based on asset and source rules. + """ + + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: + """ + Determines and validates the effective supplier for the asset. + + Args: + context: The asset processing context. + + Returns: + The updated asset processing context. + """ + effective_supplier = None + logger = logging.getLogger(__name__) # Using a logger specific to this module + + # 1. Check asset_rule.supplier_override + if context.asset_rule and context.asset_rule.supplier_override: + effective_supplier = context.asset_rule.supplier_override + logger.debug(f"Asset '{context.asset_rule.name}': Supplier override found: '{effective_supplier}'.") + + # 2. If not overridden, check source_rule.supplier + if not effective_supplier and context.source_rule and context.source_rule.supplier: + effective_supplier = context.source_rule.supplier + logger.debug(f"Asset '{context.asset_rule.name if context.asset_rule else 'Unknown'}': Source rule supplier found: '{effective_supplier}'.") + + # 3. Validation + if not effective_supplier: + asset_name = context.asset_rule.name if context.asset_rule else "Unknown Asset" + logger.error(f"Asset '{asset_name}': No supplier defined in asset rule or source rule.") + context.effective_supplier = None + if 'status_flags' not in context: # Ensure status_flags exists + context.status_flags = {} + context.status_flags['supplier_error'] = True + elif context.config_obj and effective_supplier not in context.config_obj.suppliers: + asset_name = context.asset_rule.name if context.asset_rule else "Unknown Asset" + logger.warning( + f"Asset '{asset_name}': Supplier '{effective_supplier}' not found in global supplier configuration. " + f"Available: {list(context.config_obj.suppliers.keys()) if context.config_obj.suppliers else 'None'}" + ) + context.effective_supplier = None + if 'status_flags' not in context: # Ensure status_flags exists + context.status_flags = {} + context.status_flags['supplier_error'] = True + else: + context.effective_supplier = effective_supplier + asset_name = context.asset_rule.name if context.asset_rule else "Unknown Asset" + logger.info(f"Asset '{asset_name}': Effective supplier set to '{effective_supplier}'.") + # Optionally clear the error flag if previously set and now resolved, though current logic doesn't show this path. + # if 'status_flags' in context and 'supplier_error' in context.status_flags: + # del context.status_flags['supplier_error'] + + + return context \ No newline at end of file diff --git a/processing/utils/__init__.py b/processing/utils/__init__.py new file mode 100644 index 0000000..5f3ceb7 --- /dev/null +++ b/processing/utils/__init__.py @@ -0,0 +1 @@ +# This file makes the 'utils' directory a Python package. \ No newline at end of file diff --git a/processing/utils/image_processing_utils.py b/processing/utils/image_processing_utils.py new file mode 100644 index 0000000..46768a8 --- /dev/null +++ b/processing/utils/image_processing_utils.py @@ -0,0 +1,357 @@ +import cv2 +import numpy as np +from pathlib import Path +import math +from typing import Optional, Union, List, Tuple, Dict + +# --- Basic Power-of-Two Utilities --- + +def is_power_of_two(n: int) -> bool: + """Checks if a number is a power of two.""" + return (n > 0) and (n & (n - 1) == 0) + +def get_nearest_pot(value: int) -> int: + """Finds the nearest power of two to the given value.""" + if value <= 0: + return 1 # POT must be positive, return 1 as a fallback + if is_power_of_two(value): + return value + + lower_pot = 1 << (value.bit_length() - 1) + upper_pot = 1 << value.bit_length() + + if (value - lower_pot) < (upper_pot - value): + return lower_pot + else: + return upper_pot + +# --- Dimension Calculation --- + +def calculate_target_dimensions( + original_width: int, + original_height: int, + target_width: Optional[int] = None, + target_height: Optional[int] = None, + resize_mode: str = "fit", # e.g., "fit", "stretch", "max_dim_pot" + ensure_pot: bool = False, + allow_upscale: bool = False, + target_max_dim_for_pot_mode: Optional[int] = None # Specific for "max_dim_pot" +) -> Tuple[int, int]: + """ + Calculates target dimensions based on various modes and constraints. + + Args: + original_width: Original width of the image. + original_height: Original height of the image. + target_width: Desired target width. + target_height: Desired target height. + resize_mode: + - "fit": Scales to fit within target_width/target_height, maintaining aspect ratio. + Requires at least one of target_width or target_height. + - "stretch": Scales to exactly target_width and target_height, ignoring aspect ratio. + Requires both target_width and target_height. + - "max_dim_pot": Scales to fit target_max_dim_for_pot_mode while maintaining aspect ratio, + then finds nearest POT for each dimension. Requires target_max_dim_for_pot_mode. + ensure_pot: If True, final dimensions will be adjusted to the nearest power of two. + allow_upscale: If False, dimensions will not exceed original dimensions unless ensure_pot forces it. + target_max_dim_for_pot_mode: Max dimension to use when resize_mode is "max_dim_pot". + + Returns: + A tuple (new_width, new_height). + """ + if original_width <= 0 or original_height <= 0: + # Fallback for invalid original dimensions + fallback_dim = 1 + if ensure_pot: + if target_width and target_height: + fallback_dim = get_nearest_pot(max(target_width, target_height, 1)) + elif target_width: + fallback_dim = get_nearest_pot(target_width) + elif target_height: + fallback_dim = get_nearest_pot(target_height) + elif target_max_dim_for_pot_mode: + fallback_dim = get_nearest_pot(target_max_dim_for_pot_mode) + else: # Default POT if no target given + fallback_dim = 256 + return (fallback_dim, fallback_dim) + return (target_width or 1, target_height or 1) + + + w, h = original_width, original_height + + if resize_mode == "max_dim_pot": + if target_max_dim_for_pot_mode is None: + raise ValueError("target_max_dim_for_pot_mode must be provided for 'max_dim_pot' resize_mode.") + + # Logic adapted from old processing_engine.calculate_target_dimensions + ratio = w / h + if ratio > 1: # Width is dominant + scaled_w = target_max_dim_for_pot_mode + scaled_h = max(1, round(scaled_w / ratio)) + else: # Height is dominant or square + scaled_h = target_max_dim_for_pot_mode + scaled_w = max(1, round(scaled_h * ratio)) + + # Upscale check for this mode is implicitly handled by target_max_dim + # If ensure_pot is true (as it was in the original logic), it's applied here + # For this mode, ensure_pot is effectively always true for the final step + w = get_nearest_pot(scaled_w) + h = get_nearest_pot(scaled_h) + return int(w), int(h) + + elif resize_mode == "fit": + if target_width is None and target_height is None: + raise ValueError("At least one of target_width or target_height must be provided for 'fit' mode.") + + if target_width and target_height: + ratio_orig = w / h + ratio_target = target_width / target_height + if ratio_orig > ratio_target: # Original is wider than target aspect + w_new = target_width + h_new = max(1, round(w_new / ratio_orig)) + else: # Original is taller or same aspect + h_new = target_height + w_new = max(1, round(h_new * ratio_orig)) + elif target_width: + w_new = target_width + h_new = max(1, round(w_new / (w / h))) + else: # target_height is not None + h_new = target_height + w_new = max(1, round(h_new * (w / h))) + w, h = w_new, h_new + + elif resize_mode == "stretch": + if target_width is None or target_height is None: + raise ValueError("Both target_width and target_height must be provided for 'stretch' mode.") + w, h = target_width, target_height + + else: + raise ValueError(f"Unsupported resize_mode: {resize_mode}") + + if not allow_upscale: + if w > original_width: w = original_width + if h > original_height: h = original_height + + if ensure_pot: + w = get_nearest_pot(w) + h = get_nearest_pot(h) + # Re-check upscale if POT adjustment made it larger than original and not allowed + if not allow_upscale: + if w > original_width: w = get_nearest_pot(original_width) # Get closest POT to original + if h > original_height: h = get_nearest_pot(original_height) + + + return int(max(1, w)), int(max(1, h)) + + +# --- Image Statistics --- + +def calculate_image_stats(image_data: np.ndarray) -> Optional[Dict]: + """ + Calculates min, max, mean for a given numpy image array. + Handles grayscale and multi-channel images. Converts to float64 for calculation. + Normalizes uint8/uint16 data to 0-1 range before calculating stats. + """ + if image_data is None: + return None + try: + data_float = image_data.astype(np.float64) + + if image_data.dtype == np.uint16: + data_float /= 65535.0 + elif image_data.dtype == np.uint8: + data_float /= 255.0 + + stats = {} + if len(data_float.shape) == 2: # Grayscale (H, W) + stats["min"] = float(np.min(data_float)) + stats["max"] = float(np.max(data_float)) + stats["mean"] = float(np.mean(data_float)) + elif len(data_float.shape) == 3: # Color (H, W, C) + stats["min"] = [float(v) for v in np.min(data_float, axis=(0, 1))] + stats["max"] = [float(v) for v in np.max(data_float, axis=(0, 1))] + stats["mean"] = [float(v) for v in np.mean(data_float, axis=(0, 1))] + else: + return None # Unsupported shape + return stats + except Exception: + return {"error": "Error calculating image stats"} + +# --- Aspect Ratio String --- + +def normalize_aspect_ratio_change(original_width: int, original_height: int, resized_width: int, resized_height: int, decimals: int = 2) -> str: + """ + Calculates the aspect ratio change string (e.g., "EVEN", "X133"). + """ + if original_width <= 0 or original_height <= 0: + return "InvalidInput" + if resized_width <= 0 or resized_height <= 0: + return "InvalidResize" + + width_change_percentage = ((resized_width - original_width) / original_width) * 100 + height_change_percentage = ((resized_height - original_height) / original_height) * 100 + + normalized_width_change = width_change_percentage / 100 + normalized_height_change = height_change_percentage / 100 + + normalized_width_change = min(max(normalized_width_change + 1, 0), 2) + normalized_height_change = min(max(normalized_height_change + 1, 0), 2) + + epsilon = 1e-9 + if abs(normalized_width_change) < epsilon and abs(normalized_height_change) < epsilon: + closest_value_to_one = 1.0 + elif abs(normalized_width_change) < epsilon: + closest_value_to_one = abs(normalized_height_change) + elif abs(normalized_height_change) < epsilon: + closest_value_to_one = abs(normalized_width_change) + else: + closest_value_to_one = min(abs(normalized_width_change), abs(normalized_height_change)) + + scale_factor = 1 / (closest_value_to_one + epsilon) if abs(closest_value_to_one) < epsilon else 1 / closest_value_to_one + + scaled_normalized_width_change = scale_factor * normalized_width_change + scaled_normalized_height_change = scale_factor * normalized_height_change + + output_width = round(scaled_normalized_width_change, decimals) + output_height = round(scaled_normalized_height_change, decimals) + + if abs(output_width - 1.0) < epsilon: output_width = 1 + if abs(output_height - 1.0) < epsilon: output_height = 1 + + if abs(output_width - output_height) < epsilon: # Handles original square or aspect maintained + output = "EVEN" + elif output_width != 1 and abs(output_height - 1.0) < epsilon : # Width changed, height maintained relative to width + output = f"X{str(output_width).replace('.', '')}" + elif output_height != 1 and abs(output_width - 1.0) < epsilon: # Height changed, width maintained relative to height + output = f"Y{str(output_height).replace('.', '')}" + else: # Both changed relative to each other + output = f"X{str(output_width).replace('.', '')}Y{str(output_height).replace('.', '')}" + return output + +# --- Image Loading, Conversion, Resizing --- + +def load_image(image_path: Union[str, Path], read_flag: int = cv2.IMREAD_UNCHANGED) -> Optional[np.ndarray]: + """Loads an image from the specified path.""" + try: + img = cv2.imread(str(image_path), read_flag) + if img is None: + # print(f"Warning: Failed to load image: {image_path}") # Optional: for debugging utils + return None + return img + except Exception: # as e: + # print(f"Error loading image {image_path}: {e}") # Optional: for debugging utils + return None + +def convert_bgr_to_rgb(image: np.ndarray) -> np.ndarray: + """Converts an image from BGR to RGB color space.""" + if image is None or len(image.shape) < 3: + return image # Return as is if not a color image or None + + if image.shape[2] == 4: # BGRA + return cv2.cvtColor(image, cv2.COLOR_BGRA2RGB) + elif image.shape[2] == 3: # BGR + return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + return image # Return as is if not 3 or 4 channels + +def convert_rgb_to_bgr(image: np.ndarray) -> np.ndarray: + """Converts an image from RGB to BGR color space.""" + if image is None or len(image.shape) < 3 or image.shape[2] != 3: # Only for 3-channel RGB + return image # Return as is if not a 3-channel color image or None + return cv2.cvtColor(image, cv2.COLOR_RGB2BGR) + + +def resize_image(image: np.ndarray, target_width: int, target_height: int, interpolation: Optional[int] = None) -> np.ndarray: + """Resizes an image to target_width and target_height.""" + if image is None: + raise ValueError("Cannot resize a None image.") + if target_width <= 0 or target_height <= 0: + raise ValueError("Target width and height must be positive.") + + original_height, original_width = image.shape[:2] + + if interpolation is None: + # Default interpolation: Lanczos for downscaling, Cubic for upscaling/same + if (target_width * target_height) < (original_width * original_height): + interpolation = cv2.INTER_LANCZOS4 + else: + interpolation = cv2.INTER_CUBIC + + return cv2.resize(image, (target_width, target_height), interpolation=interpolation) + +# --- Image Saving --- + +def save_image( + image_path: Union[str, Path], + image_data: np.ndarray, + output_format: Optional[str] = None, # e.g. "png", "jpg", "exr" + output_dtype_target: Optional[np.dtype] = None, # e.g. np.uint8, np.uint16, np.float16 + params: Optional[List[int]] = None, + convert_to_bgr_before_save: bool = True # True for most formats except EXR +) -> bool: + """ + Saves image data to a file. Handles data type and color space conversions. + + Args: + image_path: Path to save the image. + image_data: NumPy array of the image. + output_format: Desired output format (e.g., 'png', 'jpg'). If None, derived from extension. + output_dtype_target: Target NumPy dtype for saving (e.g., np.uint8, np.uint16). + If None, tries to use image_data.dtype or a sensible default. + params: OpenCV imwrite parameters (e.g., [cv2.IMWRITE_JPEG_QUALITY, 90]). + convert_to_bgr_before_save: If True and image is 3-channel, converts RGB to BGR. + Set to False for formats like EXR that expect RGB. + + Returns: + True if saving was successful, False otherwise. + """ + if image_data is None: + return False + + img_to_save = image_data.copy() + path_obj = Path(image_path) + path_obj.parent.mkdir(parents=True, exist_ok=True) + + # 1. Data Type Conversion + if output_dtype_target is not None: + if output_dtype_target == np.uint8 and img_to_save.dtype != np.uint8: + if img_to_save.dtype == np.uint16: img_to_save = (img_to_save.astype(np.float32) / 65535.0 * 255.0).astype(np.uint8) + elif img_to_save.dtype in [np.float16, np.float32, np.float64]: img_to_save = (np.clip(img_to_save, 0.0, 1.0) * 255.0).astype(np.uint8) + else: img_to_save = img_to_save.astype(np.uint8) + elif output_dtype_target == np.uint16 and img_to_save.dtype != np.uint16: + if img_to_save.dtype == np.uint8: img_to_save = (img_to_save.astype(np.float32) / 255.0 * 65535.0).astype(np.uint16) # More accurate + elif img_to_save.dtype in [np.float16, np.float32, np.float64]: img_to_save = (np.clip(img_to_save, 0.0, 1.0) * 65535.0).astype(np.uint16) + else: img_to_save = img_to_save.astype(np.uint16) + elif output_dtype_target == np.float16 and img_to_save.dtype != np.float16: + if img_to_save.dtype == np.uint16: img_to_save = (img_to_save.astype(np.float32) / 65535.0).astype(np.float16) + elif img_to_save.dtype == np.uint8: img_to_save = (img_to_save.astype(np.float32) / 255.0).astype(np.float16) + elif img_to_save.dtype in [np.float32, np.float64]: img_to_save = img_to_save.astype(np.float16) + # else: cannot convert to float16 easily + elif output_dtype_target == np.float32 and img_to_save.dtype != np.float32: + if img_to_save.dtype == np.uint16: img_to_save = (img_to_save.astype(np.float32) / 65535.0) + elif img_to_save.dtype == np.uint8: img_to_save = (img_to_save.astype(np.float32) / 255.0) + elif img_to_save.dtype == np.float16: img_to_save = img_to_save.astype(np.float32) + + + # 2. Color Space Conversion (RGB -> BGR) + # Typically, OpenCV expects BGR for formats like PNG, JPG. EXR usually expects RGB. + # The `convert_to_bgr_before_save` flag controls this. + # If output_format is exr, this should generally be False. + current_format = output_format if output_format else path_obj.suffix.lower().lstrip('.') + + if convert_to_bgr_before_save and current_format != 'exr': + if len(img_to_save.shape) == 3 and img_to_save.shape[2] == 3: + img_to_save = convert_rgb_to_bgr(img_to_save) + # BGRA is handled by OpenCV imwrite for PNGs, no explicit conversion needed if saving as RGBA. + # If it's 4-channel and not PNG/TIFF with alpha, it might need stripping or specific handling. + # For simplicity, this function assumes 3-channel RGB input if BGR conversion is active. + + # 3. Save Image + try: + if params: + cv2.imwrite(str(path_obj), img_to_save, params) + else: + cv2.imwrite(str(path_obj), img_to_save) + return True + except Exception: # as e: + # print(f"Error saving image {path_obj}: {e}") # Optional: for debugging utils + return False \ No newline at end of file diff --git a/processing_engine.py b/processing_engine.py index f4d0812..779f6f8 100644 --- a/processing_engine.py +++ b/processing_engine.py @@ -5,12 +5,8 @@ import math import shutil import tempfile import logging -import json -import re -import time from pathlib import Path from typing import List, Dict, Tuple, Optional, Set -from collections import defaultdict # Attempt to import image processing libraries try: @@ -23,21 +19,13 @@ except ImportError: cv2 = None np = None -# Attempt to import OpenEXR - Check if needed for advanced EXR flags/types -try: - import OpenEXR - import Imath - _HAS_OPENEXR = True -except ImportError: - _HAS_OPENEXR = False - # Log this information - basic EXR might still work via OpenCV - logging.debug("Optional 'OpenEXR' python package not found. EXR saving relies on OpenCV's built-in support.") try: from configuration import Configuration, ConfigurationError from rule_structure import SourceRule, AssetRule, FileRule - from utils.path_utils import generate_path_from_pattern + from utils.path_utils import generate_path_from_pattern, sanitize_filename + from utils import image_processing_utils as ipu # Added import except ImportError: print("ERROR: Cannot import Configuration or rule_structure classes.") print("Ensure configuration.py and rule_structure.py are in the same directory or Python path.") @@ -49,6 +37,20 @@ except ImportError: # Use logger defined in main.py (or configure one here if run standalone) + +from processing.pipeline.orchestrator import PipelineOrchestrator +# from processing.pipeline.asset_context import AssetProcessingContext # AssetProcessingContext is used by the orchestrator +from processing.pipeline.stages.supplier_determination import SupplierDeterminationStage +from processing.pipeline.stages.asset_skip_logic import AssetSkipLogicStage +from processing.pipeline.stages.metadata_initialization import MetadataInitializationStage +from processing.pipeline.stages.file_rule_filter import FileRuleFilterStage +from processing.pipeline.stages.gloss_to_rough_conversion import GlossToRoughConversionStage +from processing.pipeline.stages.alpha_extraction_to_mask import AlphaExtractionToMaskStage +from processing.pipeline.stages.normal_map_green_channel import NormalMapGreenChannelStage +from processing.pipeline.stages.individual_map_processing import IndividualMapProcessingStage +from processing.pipeline.stages.map_merging import MapMergingStage +from processing.pipeline.stages.metadata_finalization_save import MetadataFinalizationAndSaveStage +from processing.pipeline.stages.output_organization import OutputOrganizationStage log = logging.getLogger(__name__) # Basic config if logger hasn't been set up elsewhere (e.g., during testing) if not log.hasHandlers(): @@ -60,183 +62,7 @@ class ProcessingEngineError(Exception): """Custom exception for errors during processing engine operations.""" pass -# --- Helper Functions (Moved from AssetProcessor or kept static) --- - -def _is_power_of_two(n: int) -> bool: - """Checks if a number is a power of two.""" - return (n > 0) and (n & (n - 1) == 0) - -def get_nearest_pot(value: int) -> int: - """Finds the nearest power of two to the given value.""" - if value <= 0: - return 1 # Or raise error, POT must be positive - if _is_power_of_two(value): - return value - - # Calculate the powers of two below and above the value - lower_pot = 1 << (value.bit_length() - 1) - upper_pot = 1 << value.bit_length() - - # Determine which power of two is closer - if (value - lower_pot) < (upper_pot - value): - return lower_pot - else: - return upper_pot - -def calculate_target_dimensions(orig_w, orig_h, target_max_dim) -> tuple[int, int]: - """ - Calculates target dimensions by first scaling to fit target_max_dim - while maintaining aspect ratio, then finding the nearest power-of-two - value for each resulting dimension (Stretch/Squash to POT). - """ - if orig_w <= 0 or orig_h <= 0: - # Fallback to target_max_dim if original dimensions are invalid - pot_dim = get_nearest_pot(target_max_dim) - log.warning(f"Invalid original dimensions ({orig_w}x{orig_h}). Falling back to nearest POT of target_max_dim: {pot_dim}x{pot_dim}") - return (pot_dim, pot_dim) - - # Step 1: Calculate intermediate dimensions maintaining aspect ratio - ratio = orig_w / orig_h - if ratio > 1: # Width is dominant - scaled_w = target_max_dim - scaled_h = max(1, round(scaled_w / ratio)) - else: # Height is dominant or square - scaled_h = target_max_dim - scaled_w = max(1, round(scaled_h * ratio)) - - # Step 2: Find the nearest power of two for each scaled dimension - pot_w = get_nearest_pot(scaled_w) - pot_h = get_nearest_pot(scaled_h) - - log.debug(f"POT Calc: Orig=({orig_w}x{orig_h}), MaxDim={target_max_dim} -> Scaled=({scaled_w}x{scaled_h}) -> POT=({pot_w}x{pot_h})") - - return int(pot_w), int(pot_h) - -def _calculate_image_stats(image_data: np.ndarray) -> dict | None: - """ - Calculates min, max, mean for a given numpy image array. - Handles grayscale and multi-channel images. Converts to float64 for calculation. - """ - if image_data is None: - log.warning("Attempted to calculate stats on None image data.") - return None - if np is None: - log.error("Numpy not available for stats calculation.") - return None - try: - # Use float64 for calculations to avoid potential overflow/precision issues - data_float = image_data.astype(np.float64) - - # Normalize data_float based on original dtype before calculating stats - if image_data.dtype == np.uint16: - log.debug("Stats calculation: Normalizing uint16 data to 0-1 range.") - data_float /= 65535.0 - elif image_data.dtype == np.uint8: - log.debug("Stats calculation: Normalizing uint8 data to 0-1 range.") - data_float /= 255.0 - # Assuming float inputs are already in 0-1 range or similar - - log.debug(f"Stats calculation: data_float dtype: {data_float.dtype}, shape: {data_float.shape}") - # Log a few sample values to check range after normalization - if data_float.size > 0: - sample_values = data_float.flatten()[:10] # Get first 10 values - log.debug(f"Stats calculation: Sample values (first 10) after normalization: {sample_values.tolist()}") - - - if len(data_float.shape) == 2: # Grayscale (H, W) - min_val = float(np.min(data_float)) - max_val = float(np.max(data_float)) - mean_val = float(np.mean(data_float)) - stats = {"min": min_val, "max": max_val, "mean": mean_val} - log.debug(f"Calculated Grayscale Stats: Min={min_val:.4f}, Max={max_val:.4f}, Mean={mean_val:.4f}") - elif len(data_float.shape) == 3: # Color (H, W, C) - channels = data_float.shape[2] - min_val = [float(v) for v in np.min(data_float, axis=(0, 1))] - max_val = [float(v) for v in np.max(data_float, axis=(0, 1))] - mean_val = [float(v) for v in np.mean(data_float, axis=(0, 1))] - # Assume data is RGB order after potential conversion in _load_and_transform_source - stats = {"min": min_val, "max": max_val, "mean": mean_val} - log.debug(f"Calculated {channels}-Channel Stats (RGB order): Min={min_val}, Max={max_val}, Mean={mean_val}") - else: - log.warning(f"Cannot calculate stats for image with unsupported shape {data_float.shape}") - return None - return stats - except Exception as e: - log.error(f"Error calculating image stats: {e}", exc_info=True) # Log exception info - return {"error": str(e)} - -def _sanitize_filename(name: str) -> str: - """Removes or replaces characters invalid for filenames/directory names.""" - if not isinstance(name, str): name = str(name) - name = re.sub(r'[^\w.\-]+', '_', name) # Allow alphanumeric, underscore, hyphen, dot - name = re.sub(r'_+', '_', name) - name = name.strip('_') - if not name: name = "invalid_name" - return name - -def _normalize_aspect_ratio_change(original_width, original_height, resized_width, resized_height, decimals=2): - """ - Calculates the aspect ratio change string (e.g., "EVEN", "X133"). - Returns the string representation. - """ - if original_width <= 0 or original_height <= 0: - log.warning("Cannot calculate aspect ratio change with zero original dimensions.") - return "InvalidInput" - - # Avoid division by zero if resize resulted in zero dimensions (shouldn't happen with checks) - if resized_width <= 0 or resized_height <= 0: - log.warning("Cannot calculate aspect ratio change with zero resized dimensions.") - return "InvalidResize" - - # Original logic from user feedback - width_change_percentage = ((resized_width - original_width) / original_width) * 100 - height_change_percentage = ((resized_height - original_height) / original_height) * 100 - - normalized_width_change = width_change_percentage / 100 - normalized_height_change = height_change_percentage / 100 - - normalized_width_change = min(max(normalized_width_change + 1, 0), 2) - normalized_height_change = min(max(normalized_height_change + 1, 0), 2) - - # Handle potential zero division if one dimension change is exactly -100% (normalized to 0) - # If both are 0, aspect ratio is maintained. If one is 0, the other dominates. - if normalized_width_change == 0 and normalized_height_change == 0: - closest_value_to_one = 1.0 # Avoid division by zero, effectively scale_factor = 1 - elif normalized_width_change == 0: - closest_value_to_one = abs(normalized_height_change) - elif normalized_height_change == 0: - closest_value_to_one = abs(normalized_width_change) - else: - closest_value_to_one = min(abs(normalized_width_change), abs(normalized_height_change)) - - # Add a small epsilon to avoid division by zero if closest_value_to_one is extremely close to 0 - epsilon = 1e-9 - scale_factor = 1 / (closest_value_to_one + epsilon) if abs(closest_value_to_one) < epsilon else 1 / closest_value_to_one - - scaled_normalized_width_change = scale_factor * normalized_width_change - scaled_normalized_height_change = scale_factor * normalized_height_change - - output_width = round(scaled_normalized_width_change, decimals) - output_height = round(scaled_normalized_height_change, decimals) - - # Convert to int if exactly 1.0 after rounding - if abs(output_width - 1.0) < epsilon: output_width = 1 - if abs(output_height - 1.0) < epsilon: output_height = 1 - - # Determine output string - if original_width == original_height or abs(output_width - output_height) < epsilon: - output = "EVEN" - elif output_width != 1 and output_height == 1: - output = f"X{str(output_width).replace('.', '')}" - elif output_height != 1 and output_width == 1: - output = f"Y{str(output_height).replace('.', '')}" - else: - # Both changed relative to each other - output = f"X{str(output_width).replace('.', '')}Y{str(output_height).replace('.', '')}" - - log.debug(f"Aspect ratio change calculated: Orig=({original_width}x{original_height}), Resized=({resized_width}x{resized_height}) -> String='{output}'") - return output - +# Helper functions moved to processing.utils.image_processing_utils # --- Processing Engine Class --- class ProcessingEngine: @@ -262,6 +88,27 @@ class ProcessingEngine: self.temp_dir: Path | None = None # Path to the temporary working directory for a process run self.loaded_data_cache: dict = {} # Cache for loaded/resized data within a single process call + # --- Pipeline Orchestrator Setup --- + self.stages = [ + SupplierDeterminationStage(), + AssetSkipLogicStage(), + MetadataInitializationStage(), + FileRuleFilterStage(), + GlossToRoughConversionStage(), + AlphaExtractionToMaskStage(), + NormalMapGreenChannelStage(), + IndividualMapProcessingStage(), + MapMergingStage(), + MetadataFinalizationAndSaveStage(), + OutputOrganizationStage(), + ] + try: + self.pipeline_orchestrator = PipelineOrchestrator(config_obj=self.config_obj, stages=self.stages) + log.info("PipelineOrchestrator initialized successfully in ProcessingEngine.") + except Exception as e: + log.error(f"Failed to initialize PipelineOrchestrator in ProcessingEngine: {e}", exc_info=True) + self.pipeline_orchestrator = None # Ensure it's None if init fails + log.debug("ProcessingEngine initialized.") @@ -312,111 +159,21 @@ class ProcessingEngine: try: self.temp_dir = Path(tempfile.mkdtemp(prefix=self.config_obj.temp_dir_prefix)) log.debug(f"Created temporary workspace for engine: {self.temp_dir}") - # --- Loop through each asset defined in the SourceRule --- - for asset_rule in source_rule.assets: - asset_name = asset_rule.asset_name - log.info(f"--- Processing asset: '{asset_name}' ---") - asset_processed = False - asset_skipped = False - asset_failed = False - temp_metadata_path_asset = None # Track metadata file for this asset - - try: - # --- Determine Effective Supplier (Override > Identifier > Fallback) --- - effective_supplier = source_rule.supplier_override # Prioritize override - if effective_supplier is None: - effective_supplier = source_rule.supplier_identifier # Fallback to original identifier - if not effective_supplier: # Check if still None or empty - log.warning(f"Asset '{asset_name}': Supplier identifier missing from rule and override. Using fallback 'UnknownSupplier'.") - effective_supplier = "UnknownSupplier" # Final fallback - - log.debug(f"Asset '{asset_name}': Effective supplier determined as '{effective_supplier}' (Override: '{source_rule.supplier_override}', Original: '{source_rule.supplier_identifier}')") - - # --- Skip Check (using effective supplier) --- - supplier_sanitized = _sanitize_filename(effective_supplier) - asset_name_sanitized = _sanitize_filename(asset_name) - final_dir = output_base_path / supplier_sanitized / asset_name_sanitized - metadata_file_path = final_dir / self.config_obj.metadata_filename # Metadata filename still comes from config - - log.debug(f"Checking for existing output/overwrite at: {final_dir} (using effective supplier: '{effective_supplier}')") - - if not overwrite and final_dir.exists(): - log.info(f"Output directory found for asset '{asset_name_sanitized}' (Supplier: '{effective_supplier}') and overwrite is False. Skipping.") - overall_status["skipped"].append(asset_name) - asset_skipped = True - continue # Skip to the next asset - - elif overwrite and final_dir.exists(): - log.warning(f"Output directory exists for '{asset_name_sanitized}' (Supplier: '{effective_supplier}') and overwrite is True. Removing existing directory: {final_dir}") - try: - shutil.rmtree(final_dir) - except Exception as rm_err: - raise ProcessingEngineError(f"Failed to remove existing output directory {final_dir} during overwrite: {rm_err}") from rm_err - - # --- Prepare Asset Metadata --- - # Start with common metadata from the rule, add asset name - current_asset_metadata = asset_rule.common_metadata.copy() - current_asset_metadata["asset_name"] = asset_name - # Use the EFFECTIVE supplier here - current_asset_metadata["supplier_name"] = effective_supplier - # Add other fields that will be populated - current_asset_metadata["maps_present"] = [] - current_asset_metadata["merged_maps"] = [] - current_asset_metadata["shader_features"] = [] - current_asset_metadata["source_files_in_extra"] = [] - current_asset_metadata["image_stats_1k"] = {} - current_asset_metadata["map_details"] = {} - current_asset_metadata["aspect_ratio_change_string"] = "N/A" - current_asset_metadata["merged_map_channel_stats"] = {} - - # --- Process Individual Maps --- - processed_maps_details_asset, image_stats_asset, aspect_ratio_change_string_asset = self._process_individual_maps( - asset_rule=asset_rule, - workspace_path=workspace_path, # Use the workspace path received by process() (contains prepared files) - current_asset_metadata=current_asset_metadata # Pass mutable dict - ) - # Update metadata with results (stats and aspect ratio are updated directly in current_asset_metadata by the method) - # map_details are also updated directly in current_asset_metadata - - # --- Merge Maps --- - merged_maps_details_asset = self._merge_maps( - asset_rule=asset_rule, - workspace_path=workspace_path, - processed_maps_details_asset=processed_maps_details_asset, # Needed to find resolutions - current_asset_metadata=current_asset_metadata # Pass mutable dict for stats - ) - - # --- Generate Metadata --- - # Pass effective_supplier instead of the whole source_rule - temp_metadata_path_asset = self._generate_metadata_file( - effective_supplier=effective_supplier, # Pass the determined supplier - asset_rule=asset_rule, - current_asset_metadata=current_asset_metadata, # Pass the populated dict - processed_maps_details_asset=processed_maps_details_asset, - merged_maps_details_asset=merged_maps_details_asset - ) - - # --- Organize Output --- - # Pass effective_supplier instead of source_rule.supplier_identifier - self._organize_output_files( - asset_rule=asset_rule, - workspace_path=workspace_path, # Pass the original workspace path - supplier_identifier=effective_supplier, # Pass the determined supplier - output_base_path=output_base_path, # Pass output path - processed_maps_details_asset=processed_maps_details_asset, - merged_maps_details_asset=merged_maps_details_asset, - temp_metadata_info=temp_metadata_path_asset - ) - - log.info(f"--- Asset '{asset_name}' processed successfully (Supplier: {effective_supplier}). ---") - overall_status["processed"].append(asset_name) - asset_processed = True - - except Exception as asset_err: - log.error(f"--- Failed processing asset '{asset_name}': {asset_err} ---", exc_info=True) - overall_status["failed"].append(asset_name) - asset_failed = True - # Continue to the next asset + # --- NEW PIPELINE ORCHESTRATOR LOGIC --- + if hasattr(self, 'pipeline_orchestrator') and self.pipeline_orchestrator: + log.info("Processing source rule using PipelineOrchestrator.") + overall_status = self.pipeline_orchestrator.process_source_rule( + source_rule=source_rule, + workspace_path=workspace_path, # This is the path to the source files (e.g. extracted archive) + output_base_path=output_base_path, + overwrite=overwrite, + incrementing_value=self.current_incrementing_value, + sha5_value=self.current_sha5_value + ) + else: + log.error(f"PipelineOrchestrator not available for SourceRule '{source_rule.input_path}'. Marking all {len(source_rule.assets)} assets as failed.") + for asset_rule in source_rule.assets: + overall_status["failed"].append(asset_rule.asset_name) log.info(f"ProcessingEngine finished. Summary: {overall_status}") return overall_status @@ -446,1243 +203,3 @@ class ProcessingEngine: log.error(f"Failed to remove engine temporary workspace {self.temp_dir}: {e}", exc_info=True) self.loaded_data_cache = {} # Clear cache after cleanup - def _get_ftd_key_from_override(self, override_string: str) -> Optional[str]: - """ - Attempts to derive a base FILE_TYPE_DEFINITIONS key from an override string - which might have a variant suffix (e.g., "MAP_COL-1" -> "MAP_COL"). - """ - if not override_string: # Handle empty or None override_string - return None - if override_string in self.config_obj.FILE_TYPE_DEFINITIONS: - return override_string - - # Regex to remove trailing suffixes like -, -, _ - # e.g., "MAP_COL-1" -> "MAP_COL", "MAP_ROUGH_variantA" -> "MAP_ROUGH" - base_candidate = re.sub(r"(-[\w\d]+|_[\w\d]+)$", "", override_string) - if base_candidate in self.config_obj.FILE_TYPE_DEFINITIONS: - return base_candidate - - return None - - def _get_map_variant_suffix(self, map_identifier: str, base_ftd_key: str) -> str: - """ - Extracts a variant suffix (e.g., "-1", "_variantA") from a map_identifier - if the base_ftd_key is a prefix of it and the suffix indicates a variant. - Example: map_identifier="MAP_COL-1", base_ftd_key="MAP_COL" -> returns "-1" - map_identifier="MAP_COL_variant", base_ftd_key="MAP_COL" -> returns "_variant" - map_identifier="MAP_COL", base_ftd_key="MAP_COL" -> returns "" - """ - if not base_ftd_key: # Ensure base_ftd_key is not empty - return "" - if map_identifier.startswith(base_ftd_key): - suffix = map_identifier[len(base_ftd_key):] - # Ensure suffix looks like a variant (starts with - or _) or is empty - if not suffix or suffix.startswith(('-', '_')): - return suffix - return "" # Default to no suffix - - def _get_base_map_type(self, map_identifier: str) -> str: - """ - Gets the base standard type (e.g., "COL") from a map identifier (e.g., "MAP_COL-1", "COL-1"), - or returns the identifier itself if it's a merged type (e.g., "NRMRGH") or not resolvable to a standard type. - """ - if not map_identifier: # Handle empty or None map_identifier - return "" - - # Try to get FTD key from "MAP_COL-1" -> "MAP_COL" or "MAP_COL" -> "MAP_COL" - ftd_key = self._get_ftd_key_from_override(map_identifier) - if ftd_key: - definition = self.config_obj.FILE_TYPE_DEFINITIONS.get(ftd_key) - if definition and definition.get("standard_type"): # Check if standard_type exists and is not empty - return definition["standard_type"] # Returns "COL" - - # If map_identifier was like "COL-1" or "ROUGH" (a standard_type itself, possibly with suffix) - # Strip suffix and check if the base is a known standard_type - # Regex to get the initial part of the string composed of uppercase letters and underscores - base_candidate_match = re.match(r"([A-Z_]+)", map_identifier.upper()) - if base_candidate_match: - potential_std_type = base_candidate_match.group(1) - for _, definition_val in self.config_obj.FILE_TYPE_DEFINITIONS.items(): - if definition_val.get("standard_type") == potential_std_type: - return potential_std_type # Found "COL" - - # If it's a merged map type (e.g., "NRMRGH"), it won't be in FTDs as a key or standard_type. - # Check if it's one of the output_map_types from MAP_MERGE_RULES. - for rule in self.config_obj.map_merge_rules: - if rule.get("output_map_type") == map_identifier: - return map_identifier # Return "NRMRGH" as is - - # Fallback: return the original identifier, uppercased. - log.debug(f"_get_base_map_type: Could not determine standard base for '{map_identifier}'. Returning as is (uppercase).") - return map_identifier.upper() - - def _load_and_transform_source(self, source_path_abs: Path, map_type: str, target_resolution_key: str, is_gloss_source: bool) -> Tuple[Optional[np.ndarray], Optional[np.dtype]]: - """ - Loads a source image file, performs initial prep (BGR->RGB, Gloss->Rough if applicable), - resizes it to the target resolution, and caches the result. - Uses static configuration from self.config_obj. - - Args: - source_path_abs: Absolute path to the source file in the workspace. - map_type: The item_type_override (e.g., "MAP_NRM", "MAP_ROUGH-1"). - target_resolution_key: The key for the target resolution (e.g., "4K"). - is_gloss_source: Boolean indicating if this source should be treated as gloss for inversion (if map_type is ROUGH). - - Returns: - Tuple containing: - - Resized NumPy array (float32 for gloss-inverted, original type otherwise) or None if loading/processing fails. - - Original source NumPy dtype or None if loading fails. - """ - if cv2 is None or np is None: - log.error("OpenCV or NumPy not available for image loading.") - return None, None - - cache_key = (source_path_abs, target_resolution_key) # Use absolute path for cache key - if cache_key in self.loaded_data_cache: - log.debug(f"CACHE HIT: Returning cached data for {source_path_abs.name} at {target_resolution_key}") - return self.loaded_data_cache[cache_key] # Return tuple (image_data, source_dtype) - - log.debug(f"CACHE MISS: Loading and transforming {source_path_abs.name} for {target_resolution_key} (map_type: {map_type})") - img_prepared = None - source_dtype = None - - try: - # --- 1. Load Source Image --- - # Determine read flag based on is_grayscale from FTD - ftd_key = self._get_ftd_key_from_override(map_type) # map_type is item_type_override - is_map_grayscale = False - standard_type_for_checks = None # For MASK check - - if ftd_key: - ftd_definition = self.config_obj.FILE_TYPE_DEFINITIONS.get(ftd_key, {}) - is_map_grayscale = ftd_definition.get("is_grayscale", False) - standard_type_for_checks = ftd_definition.get("standard_type") - log.debug(f"For map_type '{map_type}' (FTD key '{ftd_key}'), is_grayscale: {is_map_grayscale}, standard_type: {standard_type_for_checks}") - else: - log.warning(f"Could not determine FTD key for map_type '{map_type}' to check is_grayscale. Assuming not grayscale.") - - read_flag = cv2.IMREAD_GRAYSCALE if is_map_grayscale else cv2.IMREAD_UNCHANGED - - # Special case for MASK: always load unchanged first to check alpha - if standard_type_for_checks == 'MASK': - log.debug(f"Map type '{map_type}' (standard_type 'MASK') will be loaded with IMREAD_UNCHANGED for alpha check.") - read_flag = cv2.IMREAD_UNCHANGED - - log.debug(f"Loading source {source_path_abs.name} with flag: {'GRAYSCALE' if read_flag == cv2.IMREAD_GRAYSCALE else 'UNCHANGED'}") - img_loaded = cv2.imread(str(source_path_abs), read_flag) - if img_loaded is None: - raise ProcessingEngineError(f"Failed to load image file: {source_path_abs.name} with flag {read_flag}") - source_dtype = img_loaded.dtype - log.debug(f"Loaded source {source_path_abs.name}, dtype: {source_dtype}, shape: {img_loaded.shape}") - - # --- 2. Initial Preparation (BGR->RGB, Gloss Inversion, MASK handling) --- - img_prepared = img_loaded # Start with loaded image - - # MASK Handling (Extract alpha or convert) - Do this BEFORE general color conversions - if standard_type_for_checks == 'MASK': - log.debug(f"Processing as MASK type for {source_path_abs.name}.") - shape = img_prepared.shape - if len(shape) == 3 and shape[2] == 4: # BGRA or RGBA (OpenCV loads BGRA) - log.debug("MASK processing: Extracting alpha channel (4-channel source).") - img_prepared = img_prepared[:, :, 3] # Extract alpha - elif len(shape) == 3 and shape[2] == 3: # BGR or RGB - log.debug("MASK processing: Converting 3-channel source to Grayscale.") - img_prepared = cv2.cvtColor(img_prepared, cv2.COLOR_BGR2GRAY if read_flag != cv2.IMREAD_GRAYSCALE else cv2.COLOR_RGB2GRAY) # If loaded UNCHANGED and 3-channel, assume BGR - elif len(shape) == 2: - log.debug("MASK processing: Source is already grayscale.") - else: - log.warning(f"MASK processing: Unexpected source shape {shape}. Cannot reliably extract mask.") - img_prepared = None # Cannot process - else: - # BGR -> RGB conversion (only for 3/4-channel images not loaded as grayscale) - if len(img_prepared.shape) == 3 and img_prepared.shape[2] >= 3 and read_flag != cv2.IMREAD_GRAYSCALE: - log.debug(f"Converting loaded image from BGR to RGB for {source_path_abs.name}.") - if img_prepared.shape[2] == 4: # BGRA -> RGBA (then to RGB) - img_prepared = cv2.cvtColor(img_prepared, cv2.COLOR_BGRA2RGB) # OpenCV BGRA to RGB - else: # BGR -> RGB - img_prepared = cv2.cvtColor(img_prepared, cv2.COLOR_BGR2RGB) - elif len(img_prepared.shape) == 2: - log.debug(f"Image {source_path_abs.name} is grayscale or loaded as such, no BGR->RGB conversion needed.") - - if img_prepared is None: raise ProcessingEngineError("Image data is None after MASK/Color prep.") - - # Gloss -> Roughness Inversion (if map_type is ROUGH and is_gloss_source is True) - # This is triggered by the new filename logic in _process_individual_maps - if standard_type_for_checks == 'ROUGH' and is_gloss_source: - log.info(f"Performing filename-triggered Gloss->Roughness inversion for {source_path_abs.name} (map_type: {map_type})") - if len(img_prepared.shape) == 3: - log.debug("Gloss Inversion: Converting 3-channel image to grayscale before inversion.") - img_prepared = cv2.cvtColor(img_prepared, cv2.COLOR_RGB2GRAY) # Should be RGB at this point if 3-channel - - stats_before = _calculate_image_stats(img_prepared) - log.debug(f"Gloss Inversion: Image stats BEFORE inversion: {stats_before}") - - if source_dtype == np.uint16: - img_float = 1.0 - (img_prepared.astype(np.float32) / 65535.0) - elif source_dtype == np.uint8: - img_float = 1.0 - (img_prepared.astype(np.float32) / 255.0) - else: # Assuming float input is already 0-1 range - img_float = 1.0 - img_prepared.astype(np.float32) - - img_prepared = np.clip(img_float, 0.0, 1.0) # Result is float32 - - stats_after = _calculate_image_stats(img_prepared) - log.debug(f"Gloss Inversion: Image stats AFTER inversion (float32): {stats_after}") - log.debug(f"Inverted gloss map stored as float32 for ROUGH, original dtype: {source_dtype}") - - # Ensure data is float32/uint8/uint16 for resizing compatibility - if isinstance(img_prepared, np.ndarray) and img_prepared.dtype not in [np.uint8, np.uint16, np.float32, np.float16]: - log.warning(f"Converting unexpected dtype {img_prepared.dtype} to float32 before resizing for {source_path_abs.name}.") - img_prepared = img_prepared.astype(np.float32) - - # --- 3. Resize --- - if img_prepared is None: raise ProcessingEngineError(f"Image data is None after initial prep for {source_path_abs.name}.") - orig_h, orig_w = img_prepared.shape[:2] - # Get resolutions from static config - target_dim_px = self.config_obj.image_resolutions.get(target_resolution_key) - if not target_dim_px: - raise ProcessingEngineError(f"Target resolution key '{target_resolution_key}' not found in config.") - - # Avoid upscaling check (using static config) - max_original_dimension = max(orig_w, orig_h) - if target_dim_px > max_original_dimension: - log.warning(f"Target dimension {target_dim_px}px is larger than original {max_original_dimension}px for {source_path_abs.name}. Skipping resize for {target_resolution_key}.") - # Store None in cache for this specific resolution to avoid retrying - self.loaded_data_cache[cache_key] = (None, source_dtype) - return None, source_dtype # Indicate resize was skipped - - if orig_w <= 0 or orig_h <= 0: - raise ProcessingEngineError(f"Invalid original dimensions ({orig_w}x{orig_h}) for {source_path_abs.name}.") - - target_w, target_h = calculate_target_dimensions(orig_w, orig_h, target_dim_px) - interpolation = cv2.INTER_LANCZOS4 if (target_w * target_h) < (orig_w * orig_h) else cv2.INTER_CUBIC - log.debug(f"Resizing {source_path_abs.name} from ({orig_w}x{orig_h}) to ({target_w}x{target_h}) for {target_resolution_key}") - img_resized = cv2.resize(img_prepared, (target_w, target_h), interpolation=interpolation) - - # --- 4. Cache and Return --- - # Keep resized dtype unless it was gloss-inverted (which is float32) - final_data_to_cache = img_resized - # Ensure gloss-inverted maps are float32 - if standard_type_for_checks == 'ROUGH' and is_gloss_source and final_data_to_cache.dtype != np.float32: - log.debug(f"Ensuring gloss-inverted ROUGH map ({map_type}) is float32.") - final_data_to_cache = final_data_to_cache.astype(np.float32) - - log.debug(f"CACHING result for {cache_key}. Shape: {final_data_to_cache.shape}, Dtype: {final_data_to_cache.dtype}") - self.loaded_data_cache[cache_key] = (final_data_to_cache, source_dtype) - return final_data_to_cache, source_dtype - - except Exception as e: - log.error(f"Error in _load_and_transform_source for {source_path_abs.name} at {target_resolution_key}: {e}", exc_info=True) - # Cache None to prevent retrying on error for this specific key - self.loaded_data_cache[cache_key] = (None, None) - return None, None - - - def _save_image(self, image_data: np.ndarray, supplier_name: str, asset_name: str, current_map_identifier: str, resolution_key: str, source_info: dict, output_bit_depth_rule: str) -> Optional[Dict]: - """ - Handles saving an image NumPy array to a temporary file within the engine's temp_dir using token-based path generation. - Uses static configuration from self.config_obj for formats, quality, etc. - The 'maptype' token for the filename is derived based on standard_type and variants. - - Args: - image_data: NumPy array containing the image data to save. - supplier_name: The effective supplier name for the asset. - asset_name: The name of the asset. - current_map_identifier: The map type being saved (e.g., "MAP_COL", "MAP_ROUGH-1", "NRMRGH"). This is item_type_override or merged map type. - resolution_key: The resolution key (e.g., "4K"). - source_info: Dictionary containing details about the source(s). - output_bit_depth_rule: Rule for determining output bit depth. - - Returns: - A dictionary containing details of the saved file or None if saving failed. - """ - if cv2 is None or np is None: - log.error("OpenCV or NumPy not available for image saving.") - return None - if image_data is None: - log.error(f"Cannot save image for {current_map_identifier} ({resolution_key}): image_data is None.") - return None - if not self.temp_dir or not self.temp_dir.exists(): - log.error(f"Cannot save image for {current_map_identifier} ({resolution_key}): Engine temp_dir is invalid.") - return None - - try: - h, w = image_data.shape[:2] - current_dtype = image_data.dtype - log.debug(f"Saving {current_map_identifier} ({resolution_key}) for asset '{asset_name}'. Input shape: {image_data.shape}, dtype: {current_dtype}") - - config = self.config_obj - primary_fmt_16, fallback_fmt_16 = config.get_16bit_output_formats() - fmt_8bit_config = config.get_8bit_output_format() - threshold = config.resolution_threshold_for_jpg - force_lossless_map_types = config.force_lossless_map_types - jpg_quality = config.jpg_quality - png_compression_level = config._core_settings.get('PNG_COMPRESSION_LEVEL', 6) - image_resolutions = config.image_resolutions - output_directory_pattern = config.output_directory_pattern - output_filename_pattern = config.output_filename_pattern - - # --- 1. Determine Output Bit Depth --- - source_bpc = source_info.get('source_bit_depth', 8) - max_input_bpc = source_info.get('max_input_bit_depth', source_bpc) - output_dtype_target, output_bit_depth = np.uint8, 8 - - if output_bit_depth_rule == 'force_8bit': output_dtype_target, output_bit_depth = np.uint8, 8 - elif output_bit_depth_rule == 'force_16bit': output_dtype_target, output_bit_depth = np.uint16, 16 - elif output_bit_depth_rule == 'respect': - if source_bpc == 16: output_dtype_target, output_bit_depth = np.uint16, 16 - elif output_bit_depth_rule == 'respect_inputs': - if max_input_bpc == 16: output_dtype_target, output_bit_depth = np.uint16, 16 - else: - log.warning(f"Unknown output_bit_depth_rule '{output_bit_depth_rule}'. Defaulting to 8-bit.") - output_dtype_target, output_bit_depth = np.uint8, 8 - log.debug(f"Target output bit depth: {output_bit_depth}-bit for {current_map_identifier}") - - # --- 2. Determine Output Format --- - output_format, output_ext, save_params, needs_float16 = "", "", [], False - # Use the (potentially suffixed) standard_type for lossless check - base_standard_type_for_lossless_check = self._get_base_map_type(current_map_identifier) # "COL", "NRM", "DISP-Detail" -> "DISP" - - # Check if the pure standard type (without suffix) is in force_lossless_map_types - pure_standard_type = self._get_ftd_key_from_override(base_standard_type_for_lossless_check) # Get FTD key if possible - std_type_from_ftd = None - if pure_standard_type and pure_standard_type in self.config_obj.FILE_TYPE_DEFINITIONS: - std_type_from_ftd = self.config_obj.FILE_TYPE_DEFINITIONS[pure_standard_type].get("standard_type") - - # Use std_type_from_ftd if available and non-empty, else base_standard_type_for_lossless_check - check_type_for_lossless = std_type_from_ftd if std_type_from_ftd else base_standard_type_for_lossless_check - - force_lossless = check_type_for_lossless in force_lossless_map_types - original_extension = source_info.get('original_extension', '.png') - involved_extensions = source_info.get('involved_extensions', {original_extension}) - target_dim_px = image_resolutions.get(resolution_key, 0) - - if force_lossless: - log.debug(f"Format forced to lossless for map type '{current_map_identifier}' (checked as '{check_type_for_lossless}').") - if output_bit_depth == 16: - output_format = primary_fmt_16 - if output_format.startswith("exr"): output_ext, needs_float16 = ".exr", True; save_params.extend([cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF]) - else: output_format = fallback_fmt_16 if fallback_fmt_16 == "png" else "png"; output_ext = ".png"; save_params.extend([cv2.IMWRITE_PNG_COMPRESSION, png_compression_level]) - else: output_format, output_ext = "png", ".png"; save_params = [cv2.IMWRITE_PNG_COMPRESSION, png_compression_level] - elif output_bit_depth == 8 and target_dim_px >= threshold: - output_format = 'jpg'; output_ext = '.jpg'; save_params.extend([cv2.IMWRITE_JPEG_QUALITY, jpg_quality]) - else: - highest_format_str = 'jpg' - if '.exr' in involved_extensions: highest_format_str = 'exr' - elif '.tif' in involved_extensions: highest_format_str = 'tif' - elif '.png' in involved_extensions: highest_format_str = 'png' - - if highest_format_str == 'exr': - if output_bit_depth == 16: output_format, output_ext, needs_float16 = "exr", ".exr", True; save_params.extend([cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF]) - else: output_format, output_ext = "png", ".png"; save_params.extend([cv2.IMWRITE_PNG_COMPRESSION, png_compression_level]) - elif highest_format_str == 'tif' or highest_format_str == 'png': - if output_bit_depth == 16: - output_format = primary_fmt_16 - if output_format.startswith("exr"): output_ext, needs_float16 = ".exr", True; save_params.extend([cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF]) - else: output_format = "png"; output_ext = ".png"; save_params.extend([cv2.IMWRITE_PNG_COMPRESSION, png_compression_level]) - else: output_format, output_ext = "png", ".png"; save_params.extend([cv2.IMWRITE_PNG_COMPRESSION, png_compression_level]) - else: - output_format = fmt_8bit_config; output_ext = f".{output_format}" - if output_format == "png": save_params.extend([cv2.IMWRITE_PNG_COMPRESSION, png_compression_level]) - elif output_format == "jpg": save_params.extend([cv2.IMWRITE_JPEG_QUALITY, jpg_quality]) - - if output_format == "jpg" and output_bit_depth == 16: - log.warning(f"Output format JPG, but target 16-bit. Forcing 8-bit for {current_map_identifier}.") - output_dtype_target, output_bit_depth = np.uint8, 8 - log.debug(f"Determined save format for {current_map_identifier}: {output_format}, ext: {output_ext}, bit_depth: {output_bit_depth}") - - # --- 3. Final Data Type Conversion --- - img_to_save = image_data.copy() - if output_dtype_target == np.uint8 and img_to_save.dtype != np.uint8: - if img_to_save.dtype == np.uint16: img_to_save = (img_to_save.astype(np.float32) / 65535.0 * 255.0).astype(np.uint8) - elif img_to_save.dtype in [np.float16, np.float32]: img_to_save = (np.clip(img_to_save, 0.0, 1.0) * 255.0).astype(np.uint8) - else: img_to_save = img_to_save.astype(np.uint8) - elif output_dtype_target == np.uint16 and img_to_save.dtype != np.uint16: - if img_to_save.dtype == np.uint8: img_to_save = img_to_save.astype(np.uint16) * 257 - elif img_to_save.dtype in [np.float16, np.float32]: img_to_save = (np.clip(img_to_save, 0.0, 1.0) * 65535.0).astype(np.uint16) - else: img_to_save = img_to_save.astype(np.uint16) - if needs_float16 and img_to_save.dtype != np.float16: - if img_to_save.dtype == np.uint16: img_to_save = (img_to_save.astype(np.float32) / 65535.0).astype(np.float16) - elif img_to_save.dtype == np.uint8: img_to_save = (img_to_save.astype(np.float32) / 255.0).astype(np.float16) - elif img_to_save.dtype == np.float32: img_to_save = img_to_save.astype(np.float16) - else: log.warning(f"Cannot convert {img_to_save.dtype} to float16 for EXR save."); return None - - img_save_final = img_to_save - if len(img_to_save.shape) == 3 and img_to_save.shape[2] == 3 and not output_format.startswith("exr"): - try: img_save_final = cv2.cvtColor(img_to_save, cv2.COLOR_RGB2BGR) - except Exception as cvt_err: log.error(f"RGB->BGR conversion failed for {current_map_identifier}: {cvt_err}. Saving original."); - - filename_map_type_token: str - is_merged_map = any(rule.get("output_map_type") == current_map_identifier for rule in self.config_obj.map_merge_rules) - - if is_merged_map: - filename_map_type_token = current_map_identifier # e.g., "NRMRGH" - else: - base_ftd_key = self._get_ftd_key_from_override(current_map_identifier) # e.g., "MAP_COL" - if base_ftd_key: - definition = self.config_obj.FILE_TYPE_DEFINITIONS.get(base_ftd_key) - if definition and "standard_type" in definition: - standard_type_alias = definition["standard_type"] # e.g., "COL" - if standard_type_alias: # Ensure not empty - variant_suffix = self._get_map_variant_suffix(current_map_identifier, base_ftd_key) # e.g., "-1" or "" - if standard_type_alias in self.config_obj.respect_variant_map_types: - filename_map_type_token = standard_type_alias + variant_suffix # e.g., "COL-1" - else: - filename_map_type_token = standard_type_alias # e.g., "COL" - else: - log.warning(f"Empty standard_type for FTD key '{base_ftd_key}'. Using identifier '{current_map_identifier}' for maptype token.") - filename_map_type_token = current_map_identifier - else: - log.warning(f"No definition or standard_type for FTD key '{base_ftd_key}'. Using identifier '{current_map_identifier}' for maptype token.") - filename_map_type_token = current_map_identifier - else: - log.warning(f"Could not derive FTD key from '{current_map_identifier}'. Using it directly for maptype token.") - filename_map_type_token = current_map_identifier - - log.debug(f"Filename maptype token for '{current_map_identifier}' is '{filename_map_type_token}'") - - # --- 6. Construct Path using Token Pattern & Save --- - token_data = { - "supplier": _sanitize_filename(supplier_name), - "assetname": _sanitize_filename(asset_name), - "maptype": filename_map_type_token, - "resolution": resolution_key, - "width": w, "height": h, - "bitdepth": output_bit_depth, - "ext": output_ext.lstrip('.') - } - if hasattr(self, 'current_incrementing_value') and self.current_incrementing_value is not None: - token_data['incrementingvalue'] = self.current_incrementing_value - if hasattr(self, 'current_sha5_value') and self.current_sha5_value is not None: - token_data['sha5'] = self.current_sha5_value - - try: - relative_dir_path_str = generate_path_from_pattern(output_directory_pattern, token_data) - filename_str = generate_path_from_pattern(output_filename_pattern, token_data) - full_relative_path_str = str(Path(relative_dir_path_str) / filename_str) - except Exception as path_gen_err: - log.error(f"Failed to generate output path for {current_map_identifier} with data {token_data}: {path_gen_err}", exc_info=True) - return None - - output_path_temp = self.temp_dir / full_relative_path_str - log.debug(f"Attempting to save {current_map_identifier} to temporary path: {output_path_temp}") - - try: - output_path_temp.parent.mkdir(parents=True, exist_ok=True) - except Exception as mkdir_err: - log.error(f"Failed to create temporary directory {output_path_temp.parent}: {mkdir_err}", exc_info=True) - return None - - saved_successfully = False - actual_format_saved = output_format - try: - cv2.imwrite(str(output_path_temp), img_save_final, save_params) - saved_successfully = True - log.info(f" > Saved {current_map_identifier} ({resolution_key}, {output_bit_depth}-bit) as {output_format}") - except Exception as save_err: - log.error(f"Save failed ({output_format}) for {current_map_identifier} {resolution_key}: {save_err}") - if output_bit_depth == 16 and output_format.startswith("exr") and fallback_fmt_16 != output_format and fallback_fmt_16 == "png": - log.warning(f"Attempting fallback PNG save for {current_map_identifier} {resolution_key}") - actual_format_saved = "png"; output_ext = ".png" - # Regenerate path with .png extension for fallback - token_data_fallback = token_data.copy() - token_data_fallback["ext"] = "png" - try: - # Regenerate directory and filename separately for fallback - relative_dir_path_str_fb = generate_path_from_pattern(output_directory_pattern, token_data_fallback) - filename_str_fb = generate_path_from_pattern(output_filename_pattern, token_data_fallback) - full_relative_path_str_fb = str(Path(relative_dir_path_str_fb) / filename_str_fb) - output_path_temp = self.temp_dir / full_relative_path_str_fb # Update temp path for fallback - output_path_temp.parent.mkdir(parents=True, exist_ok=True) - except Exception as path_gen_err_fb: - log.error(f"Failed to generate fallback PNG path: {path_gen_err_fb}", exc_info=True) - return None - - save_params_fallback = [cv2.IMWRITE_PNG_COMPRESSION, png_compression_level] - img_fallback = None; target_fallback_dtype = np.uint16 - - if img_to_save.dtype == np.float16: - img_scaled = np.clip(img_to_save.astype(np.float32) * 65535.0, 0, 65535) - img_fallback = img_scaled.astype(target_fallback_dtype) - elif img_to_save.dtype == target_fallback_dtype: img_fallback = img_to_save - else: log.error(f"Cannot convert {img_to_save.dtype} for PNG fallback."); return None - - img_fallback_save_final = img_fallback - is_3_channel_fallback = len(img_fallback.shape) == 3 and img_fallback.shape[2] == 3 - if is_3_channel_fallback: # PNG is non-EXR - log.debug(f"Converting RGB to BGR for fallback PNG save {current_map_identifier} ({resolution_key})") - try: img_fallback_save_final = cv2.cvtColor(img_fallback, cv2.COLOR_RGB2BGR) - except Exception as cvt_err_fb: log.error(f"Failed RGB->BGR conversion for fallback PNG: {cvt_err_fb}. Saving original."); - - try: - cv2.imwrite(str(output_path_temp), img_fallback_save_final, save_params_fallback) - saved_successfully = True - log.info(f" > Saved {current_map_identifier} ({resolution_key}) using fallback PNG") - except Exception as fallback_err: - log.error(f"Fallback PNG save failed for {current_map_identifier} {resolution_key}: {fallback_err}", exc_info=True) - else: - log.error(f"No suitable fallback available or applicable for failed save of {current_map_identifier} ({resolution_key}) as {output_format}.") - - - # --- 6. Return Result --- - if saved_successfully: - # Return the full relative path string generated by the patterns - final_relative_path_str = full_relative_path_str_fb if actual_format_saved == "png" and output_format.startswith("exr") else full_relative_path_str - return { - "path": final_relative_path_str, # Store relative path string - "resolution": resolution_key, - "width": w, "height": h, - "bit_depth": output_bit_depth, - "format": actual_format_saved - } - else: - return None # Indicate save failure - - except Exception as e: - log.error(f"Unexpected error in _save_image for {current_map_identifier} ({resolution_key}): {e}", exc_info=True) - return None - - - def _process_individual_maps(self, asset_rule: AssetRule, workspace_path: Path, current_asset_metadata: Dict) -> Tuple[Dict[str, Dict[str, Dict]], Dict[str, Dict], str]: - """ - Processes, resizes, and saves individual map files for a specific asset - based on the provided AssetRule and static configuration. - - Args: - asset_rule: The AssetRule object containing file rules for this asset. - workspace_path: Path to the directory containing the source files. - current_asset_metadata: Mutable metadata dictionary for the current asset (updated directly). - - Returns: - Tuple containing: - - processed_maps_details_asset: Dict mapping map_type to resolution details. - - image_stats_asset: Dict mapping map_type to calculated image statistics (also added to current_asset_metadata). - - aspect_ratio_change_string_asset: String indicating aspect ratio change (also added to current_asset_metadata). - """ - if not self.temp_dir: raise ProcessingEngineError("Engine workspace (temp_dir) not setup.") - asset_name = asset_rule.asset_name - log.info(f"Processing individual map files for asset '{asset_name}'...") - - # Initialize results specific to this asset - processed_maps_details_asset: Dict[str, Dict[str, Dict]] = defaultdict(dict) - image_stats_asset: Dict[str, Dict] = {} # Local dict for stats - map_details_asset: Dict[str, Dict] = {} # Store details like source bit depth, gloss inversion - aspect_ratio_change_string_asset: str = "N/A" - - # --- Settings retrieval from static config --- - resolutions = self.config_obj.image_resolutions - stats_res_key = self.config_obj.calculate_stats_resolution - stats_target_dim = resolutions.get(stats_res_key) - if not stats_target_dim: log.warning(f"Stats resolution key '{stats_res_key}' not found in config. Stats skipped for '{asset_name}'.") - base_name = asset_name # Use the asset name from the rule - - # --- Aspect Ratio Calculation Setup --- - first_map_rule_for_aspect = next((fr for fr in asset_rule.files if fr.item_type_override is not None and fr.item_type_override != "EXTRA"), None) # Exclude EXTRA - orig_w_aspect, orig_h_aspect = None, None - if first_map_rule_for_aspect: - first_res_key = next(iter(resolutions)) # Use first resolution key - source_path_abs = workspace_path / first_map_rule_for_aspect.file_path - temp_img_for_dims, _ = self._load_and_transform_source( - source_path_abs, - first_map_rule_for_aspect.item_type_override, - first_res_key, - is_gloss_source=False # Not relevant for dimension check - # self.loaded_data_cache is used internally by the method - ) - if temp_img_for_dims is not None: - orig_h_aspect, orig_w_aspect = temp_img_for_dims.shape[:2] - log.debug(f"Got original dimensions ({orig_w_aspect}x{orig_h_aspect}) for aspect ratio calculation from {first_map_rule_for_aspect.file_path}") - else: - log.warning(f"Could not load image {first_map_rule_for_aspect.file_path} to get original dimensions for aspect ratio.") - else: - log.warning("No map files found in AssetRule, cannot calculate aspect ratio string.") - - - # --- Process Each Individual Map defined in the AssetRule --- - for file_rule in asset_rule.files: - should_skip = ( - file_rule.item_type_override is None or - file_rule.item_type_override == "EXTRA" or - getattr(file_rule, 'skip_processing', False) or - file_rule.item_type == "FILE_IGNORE" # Consolidated check: Use item_type for base classification - ) - if should_skip: - skip_reason = [] - if file_rule.item_type_override is None: skip_reason.append("No ItemTypeOverride") - if file_rule.item_type_override == "EXTRA": skip_reason.append("Explicitly EXTRA type") - if getattr(file_rule, 'skip_processing', False): skip_reason.append("SkipProcessing flag set") - if file_rule.item_type == "FILE_IGNORE": skip_reason.append("ItemType is FILE_IGNORE") - - log.debug(f"Skipping individual processing for {file_rule.file_path} ({', '.join(skip_reason)})") - continue # Skip to the next file_rule - - # --- Proceed with processing for this file_rule --- - source_path_rel = Path(file_rule.file_path) # Ensure it's a Path object - # IMPORTANT: Use the ENGINE's workspace_path (self.temp_dir) for loading, - # as individual maps should have been copied there by the caller (ProcessingTask) - # Correction: _process_individual_maps receives the *engine's* temp_dir as workspace_path - source_path_abs = workspace_path / source_path_rel - # Store original rule-based type and gloss flag - original_item_type_override = file_rule.item_type_override - # original_is_gloss_source_context removed as it's part of deprecated logic - - # --- New gloss map filename logic --- - filename_str = source_path_rel.name - is_filename_gloss_map = "map_gloss" in filename_str.lower() - - effective_map_type_for_processing = original_item_type_override - effective_is_gloss_source_for_load = False # Default to False, new filename logic will set to True if applicable - map_was_retagged_from_filename_gloss = False - - if is_filename_gloss_map: - log.info(f"-- Asset '{asset_name}': Filename '{filename_str}' contains 'MAP_GLOSS'. Applying new gloss handling. Original type from rule: '{original_item_type_override}'.") - effective_is_gloss_source_for_load = True # Force inversion if type becomes ROUGH (handled by filename logic below) - map_was_retagged_from_filename_gloss = True - - # Attempt to retag original_item_type_override from GLOSS to ROUGH, preserving MAP_ prefix case and suffix - if original_item_type_override and "gloss" in original_item_type_override.lower(): - match = re.match(r"(MAP_)(GLOSS)((?:[-_]\w+)*)", original_item_type_override, re.IGNORECASE) - if match: - prefix = match.group(1) # e.g., "MAP_" - suffix = match.group(3) if match.group(3) else "" # e.g., "-variant1_detail" or "" - effective_map_type_for_processing = f"{prefix}ROUGH{suffix}" - log.debug(f"Retagged filename gloss: original FTD key '{original_item_type_override}' to '{effective_map_type_for_processing}' for processing.") - else: - log.warning(f"Filename gloss '{original_item_type_override}' matched 'gloss' but not the expected 'MAP_GLOSS' pattern for precise retagging. Defaulting to 'MAP_ROUGH'.") - effective_map_type_for_processing = "MAP_ROUGH" - else: - # If original_item_type_override was None or didn't contain "gloss" (e.g., file was untyped but filename had MAP_GLOSS) - log.debug(f"Filename '{filename_str}' identified as gloss, but original type override ('{original_item_type_override}') was not GLOSS-specific. Setting type to 'MAP_ROUGH' for processing.") - effective_map_type_for_processing = "MAP_ROUGH" - # --- End of new gloss map filename logic --- - - log.debug(f"DEBUG POST-RETAG: effective_map_type_for_processing='{effective_map_type_for_processing}' for file '{source_path_rel.name}'") - original_extension = source_path_rel.suffix.lower() # Get from path - - log.info(f"-- Asset '{asset_name}': Processing Individual Map: {effective_map_type_for_processing} (Source: {source_path_rel.name}, EffectiveIsGlossSourceForLoad: {effective_is_gloss_source_for_load}, OriginalRuleItemType: {original_item_type_override}) --") - - current_map_details = {} # Old "derived_from_gloss_context" removed - if map_was_retagged_from_filename_gloss: - current_map_details["derived_from_gloss_filename"] = True - current_map_details["original_item_type_override_before_gloss_filename_retag"] = original_item_type_override - current_map_details["effective_item_type_override_after_gloss_filename_retag"] = effective_map_type_for_processing - source_bit_depth_found = None # Track if we've found the bit depth for this map type - - try: - # --- Loop through target resolutions from static config --- - for res_key, target_dim_px in resolutions.items(): - log.debug(f"Processing {effective_map_type_for_processing} for resolution: {res_key}...") - - # --- 1. Load and Transform Source (using helper + cache) --- - # This now only runs for files that have an item_type_override - img_resized, source_dtype = self._load_and_transform_source( - source_path_abs=source_path_abs, - map_type=effective_map_type_for_processing, # Use effective type - target_resolution_key=res_key, - is_gloss_source=effective_is_gloss_source_for_load # Pass the flag determined by filename logic - # self.loaded_data_cache is used internally - ) - - if img_resized is None: - # This warning now correctly indicates a failure for a map we *intended* to process - log.warning(f"Failed to load/transform source map {source_path_rel} (processed as {effective_map_type_for_processing}) for {res_key}. Skipping resolution.") - continue # Skip this resolution - - # Store source bit depth once found - if source_dtype is not None and source_bit_depth_found is None: - source_bit_depth_found = 16 if source_dtype == np.uint16 else (8 if source_dtype == np.uint8 else 8) # Default non-uint to 8 - current_map_details["source_bit_depth"] = source_bit_depth_found - log.debug(f"Stored source bit depth for {effective_map_type_for_processing}: {source_bit_depth_found}") - - # --- 2. Calculate Stats (if applicable) --- - if res_key == stats_res_key and stats_target_dim: - log.debug(f"Calculating stats for {effective_map_type_for_processing} using {res_key} image...") - stats = _calculate_image_stats(img_resized) - if stats: image_stats_asset[effective_map_type_for_processing] = stats # Store locally first - else: log.warning(f"Stats calculation failed for {effective_map_type_for_processing} at {res_key}.") - - # --- 3. Calculate Aspect Ratio Change String (once per asset) --- - if aspect_ratio_change_string_asset == "N/A" and orig_w_aspect is not None and orig_h_aspect is not None: - target_w_aspect, target_h_aspect = img_resized.shape[1], img_resized.shape[0] # Use current resized dims - try: - aspect_string = _normalize_aspect_ratio_change(orig_w_aspect, orig_h_aspect, target_w_aspect, target_h_aspect) - aspect_ratio_change_string_asset = aspect_string - log.debug(f"Stored aspect ratio change string using {res_key}: '{aspect_string}'") - except Exception as aspect_err: - log.error(f"Failed to calculate aspect ratio change string using {res_key}: {aspect_err}", exc_info=True) - aspect_ratio_change_string_asset = "Error" - elif aspect_ratio_change_string_asset == "N/A": - aspect_ratio_change_string_asset = "Unknown" # Set to unknown if original dims failed - - # --- 4. Save Image (using helper) --- - source_info = { - 'original_extension': original_extension, - 'source_bit_depth': source_bit_depth_found or 8, # Use found depth or default - 'involved_extensions': {original_extension} # Only self for individual maps - } - # Get bit depth rule solely from the static configuration using the correct method signature - bit_depth_rule = self.config_obj.get_bit_depth_rule(effective_map_type_for_processing) # Use effective type - - # Determine the map_type to use for saving (use effective_map_type_for_processing) - save_map_type_for_filename = effective_map_type_for_processing - # If effective_map_type_for_processing is None, this file shouldn't be saved as an individual map. - # This case should ideally be caught by the skip logic earlier, but adding a check here for safety. - if save_map_type_for_filename is None: - log.warning(f"Skipping save for {file_rule.file_path}: effective_map_type_for_processing is None.") - continue # Skip saving this file - - # Get supplier name from metadata (set in process method) - supplier_name = current_asset_metadata.get("supplier_name", "UnknownSupplier") - - save_result = self._save_image( - image_data=img_resized, - supplier_name=supplier_name, - asset_name=base_name, - current_map_identifier=save_map_type_for_filename, # Pass the effective map type to be saved - resolution_key=res_key, - source_info=source_info, - output_bit_depth_rule=bit_depth_rule - ) - - # --- 5. Store Result --- - if save_result: - processed_maps_details_asset.setdefault(effective_map_type_for_processing, {})[res_key] = save_result - # Update overall map detail (e.g., final format) if needed - current_map_details["output_format"] = save_result.get("format") - else: - log.error(f"Failed to save {effective_map_type_for_processing} at {res_key}.") - processed_maps_details_asset.setdefault(effective_map_type_for_processing, {})[f'error_{res_key}'] = "Save failed" - - - except Exception as map_proc_err: - log.error(f"Failed processing map {effective_map_type_for_processing} from {source_path_rel.name}: {map_proc_err}", exc_info=True) - processed_maps_details_asset.setdefault(effective_map_type_for_processing, {})['error'] = str(map_proc_err) - - # Store collected details for this map type (using effective_map_type_for_processing as the key) - map_details_asset[effective_map_type_for_processing] = current_map_details - - # --- Final Metadata Updates --- - # Update the passed-in current_asset_metadata dictionary directly - current_asset_metadata["map_details"] = map_details_asset - current_asset_metadata["image_stats_1k"] = image_stats_asset # Add collected stats - current_asset_metadata["aspect_ratio_change_string"] = aspect_ratio_change_string_asset # Add collected aspect string - - log.info(f"Finished processing individual map files for asset '{asset_name}'.") - # Return details needed for organization, stats and aspect ratio are updated in-place - return processed_maps_details_asset, image_stats_asset, aspect_ratio_change_string_asset - - - def _merge_maps(self, asset_rule: AssetRule, workspace_path: Path, processed_maps_details_asset: Dict[str, Dict[str, Dict]], current_asset_metadata: Dict) -> Dict[str, Dict[str, Dict]]: - """ - Merges channels from different source maps for a specific asset based on static - merge rules in configuration, using explicit file paths from the AssetRule. - - Args: - asset_rule: The AssetRule object containing file rules for this asset. - workspace_path: Path to the directory containing the source files. - processed_maps_details_asset: Details of processed maps (used to find common resolutions). - current_asset_metadata: Mutable metadata dictionary for the current asset (updated for stats). - - - Returns: - Dict[str, Dict[str, Dict]]: Details of the merged maps created for this asset. - """ - if not self.temp_dir: raise ProcessingEngineError("Engine workspace (temp_dir) not setup.") - asset_name = asset_rule.asset_name - # Get merge rules from static config - merge_rules = self.config_obj.map_merge_rules - log.info(f"Asset '{asset_name}': Applying {len(merge_rules)} map merging rule(s) from static config...") - - # Initialize results for this asset - merged_maps_details_asset: Dict[str, Dict[str, Dict]] = defaultdict(dict) - - for rule_index, rule in enumerate(merge_rules): - output_map_type = rule.get("output_map_type") - inputs_mapping = rule.get("inputs") # e.g., {"R": "AO", "G": "ROUGH", "B": "METAL"} - defaults = rule.get("defaults", {}) - rule_bit_depth = rule.get("output_bit_depth", "respect_inputs") - - if not output_map_type or not inputs_mapping: - log.warning(f"Asset '{asset_name}': Skipping static merge rule #{rule_index+1}: Missing 'output_map_type' or 'inputs'. Rule: {rule}") - continue - - log.info(f"-- Asset '{asset_name}': Applying merge rule for '{output_map_type}' --") - - # --- Find required SOURCE FileRules within the AssetRule --- - required_input_file_rules: Dict[str, FileRule] = {} # map_type -> FileRule - possible_to_find_sources = True - input_types_needed = set(inputs_mapping.values()) # e.g., {"AO", "ROUGH", "METAL"} - - for input_type in input_types_needed: - found_rule_for_type = False - # Search in the asset_rule's files - for file_rule in asset_rule.files: - # Check if the file_rule's item_type_override matches the required input type - item_override = getattr(file_rule, 'item_type_override', None) - item_base_type = getattr(file_rule, 'item_type', None) # Get base type for ignore check - - # Check if override matches the required input type AND the base type is not FILE_IGNORE - if item_override == input_type and item_base_type != "FILE_IGNORE": - # Found a valid match based on item_type_override and not ignored - required_input_file_rules[input_type] = file_rule - found_rule_for_type = True - # Update log message (see step 2) - log.debug(f"Found source FileRule for merge input '{input_type}': {file_rule.file_path} (ItemTypeOverride: {item_override}, ItemType: {item_base_type})") - break # Take the first valid match found - if not found_rule_for_type: - log.warning(f"Asset '{asset_name}': Required source FileRule for input map type '{input_type}' not found in AssetRule. Cannot perform merge for '{output_map_type}'.") - possible_to_find_sources = False - break - - if not possible_to_find_sources: - continue # Skip this merge rule - - # --- Determine common resolutions based on *processed* maps --- - # This still seems the most reliable way to know which sizes are actually available - possible_resolutions_per_input: List[Set[str]] = [] - resolutions_config = self.config_obj.image_resolutions # Static config - - for input_type in input_types_needed: - # Find the corresponding processed map details (might be ROUGH-1, ROUGH-2 etc.) - processed_details_for_input = None - input_file_rule = required_input_file_rules.get(input_type) - if input_file_rule: - processed_details_for_input = processed_maps_details_asset.get(input_file_rule.item_type_override) # Use the correct attribute - - if processed_details_for_input: - res_keys = {res for res, details in processed_details_for_input.items() if isinstance(details, dict) and 'error' not in details} - if not res_keys: - log.warning(f"Asset '{asset_name}': Input map type '{input_type}' (using {input_file_rule.item_type_override if input_file_rule else 'N/A'}) for merge rule '{output_map_type}' has no successfully processed resolutions.") # Use item_type_override - possible_resolutions_per_input = [] # Invalidate if any input has no resolutions - break - possible_resolutions_per_input.append(res_keys) - else: - # If the input map wasn't processed individually (used_for_merge_only=True) - # Assume all configured resolutions are potentially available. Loading will handle skips. - log.debug(f"Input map type '{input_type}' for merge rule '{output_map_type}' might not have been processed individually. Assuming all configured resolutions possible.") - possible_resolutions_per_input.append(set(resolutions_config.keys())) - - - if not possible_resolutions_per_input: - log.warning(f"Asset '{asset_name}': Cannot determine common resolutions for '{output_map_type}'. Skipping rule.") - continue - - common_resolutions = set.intersection(*possible_resolutions_per_input) - - if not common_resolutions: - log.warning(f"Asset '{asset_name}': No common resolutions found among required inputs {input_types_needed} for merge rule '{output_map_type}'. Skipping rule.") - continue - log.debug(f"Asset '{asset_name}': Common resolutions for '{output_map_type}': {common_resolutions}") - - # --- Loop through common resolutions --- - res_order = {k: resolutions_config[k] for k in common_resolutions if k in resolutions_config} - if not res_order: - log.warning(f"Asset '{asset_name}': Common resolutions {common_resolutions} do not match config. Skipping merge for '{output_map_type}'.") - continue - - sorted_res_keys = sorted(res_order.keys(), key=lambda k: res_order[k], reverse=True) - base_name = asset_name # Use current asset's name - - for current_res_key in sorted_res_keys: - log.debug(f"Asset '{asset_name}': Merging '{output_map_type}' for resolution: {current_res_key}") - try: - loaded_inputs_data = {} # map_type -> loaded numpy array - source_info_for_save = {'involved_extensions': set(), 'max_input_bit_depth': 8} - - # --- Load required SOURCE maps using helper --- - possible_to_load = True - target_channels = list(inputs_mapping.keys()) # e.g., ['R', 'G', 'B'] - - for map_type_needed in input_types_needed: # e.g., {"AO", "ROUGH", "METAL"} - file_rule = required_input_file_rules.get(map_type_needed) - if not file_rule: - log.error(f"Internal Error: FileRule missing for '{map_type_needed}' during merge load.") - possible_to_load = False; break - - source_path_rel_str = file_rule.file_path # Keep original string if needed - source_path_rel = Path(source_path_rel_str) # Convert to Path object - source_path_abs = workspace_path / source_path_rel - original_ext = source_path_rel.suffix.lower() # Now works on Path object - source_info_for_save['involved_extensions'].add(original_ext) - - # Determine if this specific source for merge should be treated as gloss - # based on its filename, aligning with the new primary rule. - filename_str_for_merge_input = source_path_rel.name - is_gloss_for_merge_input = "map_gloss" in filename_str_for_merge_input.lower() - if is_gloss_for_merge_input: - log.debug(f"Merge input '{filename_str_for_merge_input}' for '{map_type_needed}' identified as gloss by filename. Will pass is_gloss_source=True.") - - log.debug(f"Loading source '{source_path_rel}' for merge input '{map_type_needed}' at {current_res_key} (is_gloss_for_merge_input: {is_gloss_for_merge_input})") - img_resized, source_dtype = self._load_and_transform_source( - source_path_abs=source_path_abs, - map_type=file_rule.item_type_override, # Use the specific type override from rule (e.g., ROUGH-1) - target_resolution_key=current_res_key, - is_gloss_source=is_gloss_for_merge_input # Pass determined gloss state - # self.loaded_data_cache used internally - ) - - if img_resized is None: - log.warning(f"Asset '{asset_name}': Failed to load/transform source '{source_path_rel}' for merge input '{map_type_needed}' at {current_res_key}. Skipping resolution.") - possible_to_load = False; break - - loaded_inputs_data[map_type_needed] = img_resized # Store by base type (AO, ROUGH) - - # Track max source bit depth - if source_dtype == np.uint16: - source_info_for_save['max_input_bit_depth'] = max(source_info_for_save['max_input_bit_depth'], 16) - # Add other dtype checks if needed - - if not possible_to_load: continue - - # --- Calculate Stats for ROUGH source if used and at stats resolution --- - stats_res_key = self.config_obj.calculate_stats_resolution - if current_res_key == stats_res_key: - log.debug(f"Asset '{asset_name}': Checking for ROUGH source stats for '{output_map_type}' at {stats_res_key}") - for target_channel, source_map_type in inputs_mapping.items(): - if source_map_type == 'ROUGH' and source_map_type in loaded_inputs_data: - log.debug(f"Asset '{asset_name}': Calculating stats for ROUGH source (mapped to channel '{target_channel}') for '{output_map_type}' at {stats_res_key}") - rough_image_data = loaded_inputs_data[source_map_type] - rough_stats = _calculate_image_stats(rough_image_data) - if rough_stats: - # Update the mutable metadata dict passed in - stats_dict = current_asset_metadata.setdefault("merged_map_channel_stats", {}).setdefault(output_map_type, {}).setdefault(target_channel, {}) - stats_dict[stats_res_key] = rough_stats - log.debug(f"Asset '{asset_name}': Stored ROUGH stats for '{output_map_type}' channel '{target_channel}' at {stats_res_key}: {rough_stats}") - else: - log.warning(f"Asset '{asset_name}': Failed to calculate ROUGH stats for '{output_map_type}' channel '{target_channel}' at {stats_res_key}.") - - - # --- Determine dimensions --- - first_map_type = next(iter(loaded_inputs_data)) - h, w = loaded_inputs_data[first_map_type].shape[:2] - num_target_channels = len(target_channels) - - # --- Prepare and Merge Channels --- - merged_channels_float32 = [] - for target_channel in target_channels: # e.g., 'R', 'G', 'B' - source_map_type = inputs_mapping.get(target_channel) # e.g., "AO", "ROUGH", "METAL" - channel_data_float32 = None - - if source_map_type and source_map_type in loaded_inputs_data: - img_input = loaded_inputs_data[source_map_type] # Get the loaded NumPy array - - # Ensure input is float32 0-1 range for merging - if img_input.dtype == np.uint16: img_float = img_input.astype(np.float32) / 65535.0 - elif img_input.dtype == np.uint8: img_float = img_input.astype(np.float32) / 255.0 - elif img_input.dtype == np.float16: img_float = img_input.astype(np.float32) # Assume float16 is 0-1 - else: img_float = img_input.astype(np.float32) # Assume other floats are 0-1 - - num_source_channels = img_float.shape[2] if len(img_float.shape) == 3 else 1 - - # Extract the correct channel - if num_source_channels >= 3: - if target_channel == 'R': channel_data_float32 = img_float[:, :, 0] - elif target_channel == 'G': channel_data_float32 = img_float[:, :, 1] - elif target_channel == 'B': channel_data_float32 = img_float[:, :, 2] - elif target_channel == 'A' and num_source_channels == 4: channel_data_float32 = img_float[:, :, 3] - else: log.warning(f"Target channel '{target_channel}' invalid for 3/4 channel source '{source_map_type}'.") - elif num_source_channels == 1 or len(img_float.shape) == 2: - # If source is grayscale, use it for R, G, B, or A target channels - channel_data_float32 = img_float.reshape(h, w) - else: - log.warning(f"Unexpected shape {img_float.shape} for source '{source_map_type}'.") - - # Apply default if channel data couldn't be extracted - if channel_data_float32 is None: - default_val = defaults.get(target_channel) - if default_val is None: - raise ProcessingEngineError(f"Missing input/default for target channel '{target_channel}' in merge rule '{output_map_type}'.") - log.debug(f"Using default value {default_val} for target channel '{target_channel}' in '{output_map_type}'.") - channel_data_float32 = np.full((h, w), float(default_val), dtype=np.float32) - - merged_channels_float32.append(channel_data_float32) - - if not merged_channels_float32 or len(merged_channels_float32) != num_target_channels: - raise ProcessingEngineError(f"Channel count mismatch during merge for '{output_map_type}'. Expected {num_target_channels}, got {len(merged_channels_float32)}.") - - merged_image_float32 = cv2.merge(merged_channels_float32) - log.debug(f"Merged channels for '{output_map_type}' ({current_res_key}). Result shape: {merged_image_float32.shape}, dtype: {merged_image_float32.dtype}") - - # --- Save Merged Map using Helper --- - # Get supplier name from metadata (set in process method) - supplier_name = current_asset_metadata.get("supplier_name", "UnknownSupplier") - - save_result = self._save_image( - image_data=merged_image_float32, - supplier_name=supplier_name, - asset_name=base_name, - current_map_identifier=output_map_type, # Merged map type - resolution_key=current_res_key, - source_info=source_info_for_save, - output_bit_depth_rule=rule_bit_depth - ) - - # --- Record details locally --- - if save_result: - merged_maps_details_asset[output_map_type][current_res_key] = save_result - else: - log.error(f"Asset '{asset_name}': Failed to save merged map '{output_map_type}' at resolution '{current_res_key}'.") - merged_maps_details_asset.setdefault(output_map_type, {})[f'error_{current_res_key}'] = "Save failed via helper" - - - except Exception as merge_res_err: - log.error(f"Asset '{asset_name}': Failed merging '{output_map_type}' at resolution '{current_res_key}': {merge_res_err}", exc_info=True) - # Store error locally for this asset - merged_maps_details_asset.setdefault(output_map_type, {})[f'error_{current_res_key}'] = str(merge_res_err) - - log.info(f"Asset '{asset_name}': Finished applying map merging rules.") - # Return the details for this asset - return merged_maps_details_asset - - - def _generate_metadata_file(self, effective_supplier: str, asset_rule: AssetRule, current_asset_metadata: Dict, processed_maps_details_asset: Dict[str, Dict[str, Dict]], merged_maps_details_asset: Dict[str, Dict[str, Dict]]) -> Tuple[Path, str]: - """ - Gathers metadata for a specific asset based on the AssetRule and processing results, - and writes it to a temporary JSON file in the engine's temp_dir using separate directory/filename patterns. - - Args: - effective_supplier: The supplier name to use (override or original). - asset_rule: The AssetRule object for this asset. - current_asset_metadata: Base metadata dictionary (already contains name, category, archetype, stats, aspect ratio, map_details). - processed_maps_details_asset: Details of processed maps for this asset. - merged_maps_details_asset: Details of merged maps for this asset. - - Returns: - Tuple[Path, str]: A tuple containing the relative directory Path object and the filename string within the temp_dir. - """ - if not self.temp_dir: raise ProcessingEngineError("Engine workspace (temp_dir) not setup.") - asset_name = asset_rule.asset_name - if not asset_name: - log.warning("Asset name missing during metadata generation, file may be incomplete or incorrectly named.") - asset_name = "UnknownAsset_Metadata" # Fallback for filename - - log.info(f"Generating metadata file for asset '{asset_name}' (Supplier: {effective_supplier})...") - - # Start with the base metadata passed in (already contains name, category, archetype, stats, aspect, map_details) - final_metadata = current_asset_metadata.copy() - final_metadata["category"] = asset_rule.asset_type # Ensure standardized asset type is in metadata - - # Use the effective supplier passed as argument - final_metadata["supplier_name"] = effective_supplier # Already determined in process() - - # Populate map resolution details from processing results - final_metadata["processed_map_resolutions"] = {} - for map_type, res_dict in processed_maps_details_asset.items(): - keys = [res for res, d in res_dict.items() if isinstance(d, dict) and 'error' not in d] - if keys: final_metadata["processed_map_resolutions"][map_type] = sorted(keys) - - final_metadata["merged_map_resolutions"] = {} - for map_type, res_dict in merged_maps_details_asset.items(): - keys = [res for res, d in res_dict.items() if isinstance(d, dict) and 'error' not in d] - if keys: final_metadata["merged_map_resolutions"][map_type] = sorted(keys) - - # Determine maps present based on successful processing for this asset - final_metadata["maps_present"] = sorted(list(processed_maps_details_asset.keys())) - final_metadata["merged_maps"] = sorted(list(merged_maps_details_asset.keys())) - - # Determine shader features based on this asset's maps and rules - features = set() - map_details_asset = final_metadata.get("map_details", {}) # Get from metadata dict - for map_type, details in map_details_asset.items(): # map_type here is item_type_override like "MAP_COL-1" - base_standard_type = self._get_base_map_type(map_type) # Should give "COL" - # Check standard feature types - if base_standard_type in ["SSS", "FUZZ", "MASK", "TRANSMISSION", "EMISSION", "CLEARCOAT"]: - features.add(base_standard_type) - if details.get("derived_from_gloss"): features.add("InvertedGloss") - # Check if any resolution was saved as 16-bit - res_details = processed_maps_details_asset.get(map_type, {}) - if any(res_info.get("bit_depth") == 16 for res_info in res_details.values() if isinstance(res_info, dict)): - features.add(f"16bit_{base_standard_type}") - # Check merged maps for 16-bit output - for map_type, res_dict in merged_maps_details_asset.items(): # map_type here is "NRMRGH" - base_standard_type = self._get_base_map_type(map_type) # Should give "NRMRGH" - if any(res_info.get("bit_depth") == 16 for res_info in res_dict.values() if isinstance(res_info, dict)): - features.add(f"16bit_{base_standard_type}") - - final_metadata["shader_features"] = sorted(list(features)) - - # Determine source files in this asset's Extra folder based on FileRule category - source_files_in_extra_set = set() - for file_rule in asset_rule.files: - if file_rule.item_type_override is None: # Assume files without an assigned type are extra/ignored/unmatched - source_files_in_extra_set.add(str(file_rule.file_path)) - final_metadata["source_files_in_extra"] = sorted(list(source_files_in_extra_set)) - - # Add processing info - final_metadata["_processing_info"] = { - "preset_used": self.config_obj.preset_name, # Preset name comes from the engine's config - "timestamp_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), - "input_source": effective_supplier, # Use the effective supplier - } - - # Sort lists just before writing - for key in ["maps_present", "merged_maps", "shader_features", "source_files_in_extra"]: - if key in final_metadata and isinstance(final_metadata[key], list): final_metadata[key].sort() - - # --- Generate Path and Save --- - # Get the new separate patterns from config - output_directory_pattern = self.config_obj.output_directory_pattern - output_filename_pattern = self.config_obj.output_filename_pattern - metadata_filename_base = self.config_obj.metadata_filename # e.g., "metadata.json" - metadata_ext = Path(metadata_filename_base).suffix.lstrip('.') or 'json' - metadata_maptype = Path(metadata_filename_base).stem # Use filename stem as maptype token - - token_data = { - "supplier": _sanitize_filename(effective_supplier), - "assetname": _sanitize_filename(asset_name), - "maptype": metadata_maptype, - "resolution": "meta", - "width": 0, - "height": 0, - "bitdepth": 0, - "ext": metadata_ext - } - if hasattr(self, 'current_incrementing_value') and self.current_incrementing_value is not None: - token_data['incrementingvalue'] = self.current_incrementing_value - if hasattr(self, 'current_sha5_value') and self.current_sha5_value is not None: - token_data['sha5'] = self.current_sha5_value - log.debug(f"Token data for _generate_metadata_file path generation: {token_data}") # DEBUG LOG - - - try: - # Generate directory and filename separately - relative_dir_path_str = generate_path_from_pattern(output_directory_pattern, token_data) - filename_str = generate_path_from_pattern(output_filename_pattern, token_data) - # Combine for the full temporary path - full_relative_path_str = str(Path(relative_dir_path_str) / filename_str) - relative_dir_path = Path(relative_dir_path_str) # Keep the directory Path object - except Exception as path_gen_err: - log.error(f"Failed to generate metadata path using patterns '{output_directory_pattern}' / '{output_filename_pattern}' and data {token_data}: {path_gen_err}", exc_info=True) - raise ProcessingEngineError(f"Failed to generate metadata path for asset '{asset_name}'") from path_gen_err - - output_path_temp_abs = self.temp_dir / full_relative_path_str # Save to engine's temp dir, preserving structure - log.debug(f"Writing metadata for asset '{asset_name}' to temporary file: {output_path_temp_abs}") - - # Ensure parent directory exists in temp (using the full path) - try: - output_path_temp_abs.parent.mkdir(parents=True, exist_ok=True) - except Exception as mkdir_err: - log.error(f"Failed to create temporary directory {output_path_temp_abs.parent} for metadata: {mkdir_err}", exc_info=True) - raise ProcessingEngineError(f"Failed to create temporary directory for metadata for asset '{asset_name}'") from mkdir_err - - try: - with open(output_path_temp_abs, 'w', encoding='utf-8') as f: - json.dump(final_metadata, f, indent=4, ensure_ascii=False, sort_keys=True) - log.info(f"Metadata file '{filename_str}' generated successfully for asset '{asset_name}' at relative temp path '{full_relative_path_str}'.") - # Return the RELATIVE directory Path object and the filename string - return relative_dir_path, filename_str - except Exception as e: - raise ProcessingEngineError(f"Failed to write metadata file {output_path_temp_abs} for asset '{asset_name}': {e}") from e - - - def _organize_output_files(self, asset_rule: AssetRule, workspace_path: Path, supplier_identifier: str, output_base_path: Path, processed_maps_details_asset: Dict[str, Dict[str, Dict]], merged_maps_details_asset: Dict[str, Dict[str, Dict]], temp_metadata_info: Tuple[Path, str]): - """ - Moves/copies processed files for a specific asset from the engine's temp dir - and copies EXTRA files from the original workspace to the final output structure, - using the relative paths generated by the token pattern. - - Args: - asset_rule: The AssetRule object for this asset. - workspace_path: Path to the original workspace containing source files. - supplier_identifier: The supplier identifier from the SourceRule. - output_base_path: The final base output directory. - processed_maps_details_asset: Details of processed maps for this asset. - merged_maps_details_asset: Details of merged maps for this asset. - temp_metadata_info: Tuple containing the relative directory Path and filename string for the metadata file within temp_dir. - """ - if not self.temp_dir or not self.temp_dir.exists(): raise ProcessingEngineError("Engine temp workspace missing.") - asset_name = asset_rule.asset_name - if not asset_name: raise ProcessingEngineError("Asset name missing for organization.") - - if not asset_name: raise ProcessingEngineError("Asset name missing for organization.") - asset_name_sanitized = _sanitize_filename(asset_name) # Still useful for logging - - # Get structure names from static config - extra_subdir_name = self.config_obj.extra_files_subdir - - log.info(f"Organizing output files for asset '{asset_name_sanitized}' using generated paths relative to: {output_base_path}") - - # --- Helper for moving files from engine's temp dir to final output --- - def _safe_move_to_final(src_rel_path_str: str | None, file_desc: str): - """Moves a file from temp to its final location based on its relative path string.""" - if not src_rel_path_str: - log.warning(f"Asset '{asset_name_sanitized}': Missing src relative path string for {file_desc}. Cannot move.") - return - - source_abs = self.temp_dir / src_rel_path_str # Absolute path in temp - dest_abs = output_base_path / src_rel_path_str # Final absolute path - - try: - if source_abs.exists(): - # Ensure final destination directory exists - dest_abs.parent.mkdir(parents=True, exist_ok=True) - log.debug(f"Asset '{asset_name_sanitized}': Moving {file_desc}: {src_rel_path_str} -> {dest_abs.relative_to(output_base_path)}") - shutil.move(str(source_abs), str(dest_abs)) - else: - log.warning(f"Asset '{asset_name_sanitized}': Source file missing in engine temp for {file_desc}: {source_abs}") - except Exception as e: - log.error(f"Asset '{asset_name_sanitized}': Failed moving {file_desc} '{src_rel_path_str}': {e}", exc_info=True) - - # --- Move Processed/Merged Maps --- - moved_map_count = 0 - for details_dict in [processed_maps_details_asset, merged_maps_details_asset]: - for map_type, res_dict in details_dict.items(): - # Skip if the whole map type failed (e.g., merge rule source missing) - if isinstance(res_dict, dict) and 'error' in res_dict and len(res_dict) == 1: - log.warning(f"Skipping move for map type '{map_type}' due to processing error: {res_dict['error']}") - continue - for res_key, details in res_dict.items(): - # Skip specific resolution errors - if isinstance(details, str) and details.startswith("error_"): - log.warning(f"Skipping move for {map_type} ({res_key}) due to error: {details}") - continue - if isinstance(details, dict) and 'path' in details: - # details['path'] is the relative path string within temp_dir - relative_path_str = details['path'] - _safe_move_to_final(relative_path_str, f"{map_type} ({res_key})") - moved_map_count += 1 - log.debug(f"Asset '{asset_name_sanitized}': Moved {moved_map_count} map files.") - - # --- Move Metadata File --- - if temp_metadata_info: - relative_dir_path, filename = temp_metadata_info - metadata_rel_path_str = str(relative_dir_path / filename) - _safe_move_to_final(metadata_rel_path_str, "metadata file") - else: - log.warning(f"Asset '{asset_name_sanitized}': Temporary metadata info missing. Cannot move metadata file.") - - # --- Handle "EXTRA" Files (copy from original workspace to final asset dir) --- - # Determine the final asset directory based on the metadata's relative directory path - final_asset_relative_dir = relative_dir_path if temp_metadata_info else None - if final_asset_relative_dir is not None: # Check explicitly for None - final_extra_dir_abs = output_base_path / final_asset_relative_dir / extra_subdir_name - log.debug(f"Asset '{asset_name_sanitized}': Determined final EXTRA directory: {final_extra_dir_abs}") - copied_extra_files = [] - for file_rule in asset_rule.files: - # Copy files explicitly marked as EXTRA or those with no item_type_override (unmatched) - if file_rule.item_type_override == "EXTRA" or file_rule.item_type_override is None: - try: - source_rel_path = Path(file_rule.file_path) - source_abs = workspace_path / source_rel_path - # Place in Extra subdir within the final asset dir, keep original name - dest_abs = final_extra_dir_abs / source_rel_path.name - - if source_abs.is_file(): - log.debug(f"Asset '{asset_name_sanitized}': Copying EXTRA/unmatched file: {source_rel_path} -> {final_extra_dir_abs.relative_to(output_base_path)}/") - final_extra_dir_abs.mkdir(parents=True, exist_ok=True) - shutil.copy2(str(source_abs), str(dest_abs)) # copy2 preserves metadata - copied_extra_files.append(source_rel_path.name) - elif source_abs.is_dir(): - log.debug(f"Asset '{asset_name_sanitized}': Skipping EXTRA/unmatched directory: {source_rel_path}") - else: - log.warning(f"Asset '{asset_name_sanitized}': Source file marked as EXTRA/unmatched not found in workspace: {source_abs}") - except Exception as copy_err: - log.error(f"Asset '{asset_name_sanitized}': Failed copying EXTRA/unmatched file '{file_rule.file_path}': {copy_err}", exc_info=True) - - if copied_extra_files: - log.info(f"Asset '{asset_name_sanitized}': Copied {len(copied_extra_files)} EXTRA/unmatched file(s) to '{final_extra_dir_abs.relative_to(output_base_path)}' subdirectory.") - else: - log.warning(f"Asset '{asset_name_sanitized}': Could not determine final asset directory from metadata info '{temp_metadata_info}'. Skipping EXTRA file copying.") - - - log.info(f"Finished organizing output for asset '{asset_name_sanitized}'.") diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..2e70fad --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +# This file makes the 'tests' directory a Python package. \ No newline at end of file diff --git a/tests/processing/pipeline/__init__.py b/tests/processing/pipeline/__init__.py new file mode 100644 index 0000000..f178d82 --- /dev/null +++ b/tests/processing/pipeline/__init__.py @@ -0,0 +1 @@ +# This file makes Python treat the directory as a package. \ No newline at end of file diff --git a/tests/processing/pipeline/stages/__init__.py b/tests/processing/pipeline/stages/__init__.py new file mode 100644 index 0000000..f178d82 --- /dev/null +++ b/tests/processing/pipeline/stages/__init__.py @@ -0,0 +1 @@ +# This file makes Python treat the directory as a package. \ No newline at end of file diff --git a/tests/processing/pipeline/stages/test_alpha_extraction_to_mask.py b/tests/processing/pipeline/stages/test_alpha_extraction_to_mask.py new file mode 100644 index 0000000..0589585 --- /dev/null +++ b/tests/processing/pipeline/stages/test_alpha_extraction_to_mask.py @@ -0,0 +1,273 @@ +import pytest +from unittest import mock +from pathlib import Path +import uuid +import numpy as np + +from processing.pipeline.stages.alpha_extraction_to_mask import AlphaExtractionToMaskStage +from processing.pipeline.asset_context import AssetProcessingContext +from rule_structure import AssetRule, SourceRule, FileRule, TransformSettings +from configuration import Configuration, GeneralSettings +import processing.utils.image_processing_utils as ipu # Ensure ipu is available for mocking + +# Helper Functions +def create_mock_file_rule_for_alpha_test( + id_val: uuid.UUID = None, + map_type: str = "ALBEDO", + filename_pattern: str = "albedo.png", + item_type: str = "MAP_COL", + active: bool = True +) -> mock.MagicMock: + mock_fr = mock.MagicMock(spec=FileRule) + mock_fr.id = id_val if id_val else uuid.uuid4() + mock_fr.map_type = map_type + mock_fr.filename_pattern = filename_pattern + mock_fr.item_type = item_type + mock_fr.active = active + mock_fr.transform_settings = mock.MagicMock(spec=TransformSettings) + return mock_fr + +def create_alpha_extraction_mock_context( + initial_file_rules: list = None, + initial_processed_details: dict = None, + skip_asset_flag: bool = False, + asset_name: str = "AlphaAsset", + # extract_alpha_globally: bool = True # If stage checks this +) -> AssetProcessingContext: + mock_asset_rule = mock.MagicMock(spec=AssetRule) + mock_asset_rule.name = asset_name + + mock_source_rule = mock.MagicMock(spec=SourceRule) + + mock_gs = mock.MagicMock(spec=GeneralSettings) + # if your stage uses a global flag: + # mock_gs.extract_alpha_to_mask_globally = extract_alpha_globally + + mock_config = mock.MagicMock(spec=Configuration) + mock_config.general_settings = mock_gs + + context = AssetProcessingContext( + source_rule=mock_source_rule, + asset_rule=mock_asset_rule, + workspace_path=Path("/fake/workspace"), + engine_temp_dir=Path("/fake/temp_engine_dir"), + output_base_path=Path("/fake/output"), + effective_supplier="ValidSupplier", + asset_metadata={'asset_name': asset_name}, + processed_maps_details=initial_processed_details if initial_processed_details is not None else {}, + merged_maps_details={}, + files_to_process=list(initial_file_rules) if initial_file_rules else [], + loaded_data_cache={}, + config_obj=mock_config, + status_flags={'skip_asset': skip_asset_flag}, + incrementing_value=None, + sha5_value=None + ) + return context + +# Unit Tests +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.save_image') +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.load_image') +@mock.patch('logging.info') # Mock logging to avoid console output during tests +def test_asset_skipped(mock_log_info, mock_load_image, mock_save_image): + stage = AlphaExtractionToMaskStage() + context = create_alpha_extraction_mock_context(skip_asset_flag=True) + + updated_context = stage.execute(context) + + assert updated_context == context # Context should be unchanged + mock_load_image.assert_not_called() + mock_save_image.assert_not_called() + assert len(updated_context.files_to_process) == 0 + assert not updated_context.processed_maps_details + +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.save_image') +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.load_image') +@mock.patch('logging.info') +def test_existing_mask_map(mock_log_info, mock_load_image, mock_save_image): + stage = AlphaExtractionToMaskStage() + + existing_mask_rule = create_mock_file_rule_for_alpha_test(map_type="MASK", filename_pattern="mask.png") + context = create_alpha_extraction_mock_context(initial_file_rules=[existing_mask_rule]) + + updated_context = stage.execute(context) + + assert updated_context == context + mock_load_image.assert_not_called() + mock_save_image.assert_not_called() + assert len(updated_context.files_to_process) == 1 + assert updated_context.files_to_process[0].map_type == "MASK" + +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.save_image') +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.load_image') +@mock.patch('logging.info') +def test_alpha_extraction_success(mock_log_info, mock_load_image, mock_save_image): + stage = AlphaExtractionToMaskStage() + + albedo_rule_id = uuid.uuid4() + albedo_fr = create_mock_file_rule_for_alpha_test(id_val=albedo_rule_id, map_type="ALBEDO") + + initial_processed_details = { + albedo_fr.id.hex: {'temp_processed_file': '/fake/temp_engine_dir/processed_albedo.png', 'status': 'Processed', 'map_type': 'ALBEDO', 'source_file_path': Path('/fake/source/albedo.png')} + } + context = create_alpha_extraction_mock_context( + initial_file_rules=[albedo_fr], + initial_processed_details=initial_processed_details + ) + + mock_rgba_data = np.zeros((10, 10, 4), dtype=np.uint8) + mock_rgba_data[:, :, 3] = 128 # Example alpha data + mock_load_image.side_effect = [mock_rgba_data, mock_rgba_data] + + mock_save_image.return_value = True + + updated_context = stage.execute(context) + + assert mock_load_image.call_count == 2 + # First call to check for alpha, second to get data for saving + mock_load_image.assert_any_call(Path('/fake/temp_engine_dir/processed_albedo.png')) + + mock_save_image.assert_called_once() + saved_path_arg = mock_save_image.call_args[0][0] + saved_data_arg = mock_save_image.call_args[0][1] + + assert isinstance(saved_path_arg, Path) + assert "mask_from_alpha_" in saved_path_arg.name + assert np.array_equal(saved_data_arg, mock_rgba_data[:, :, 3]) + + assert len(updated_context.files_to_process) == 2 + new_mask_rule = None + for fr in updated_context.files_to_process: + if fr.map_type == "MASK": + new_mask_rule = fr + break + assert new_mask_rule is not None + assert new_mask_rule.item_type == "MAP_DER" # Derived map + + assert new_mask_rule.id.hex in updated_context.processed_maps_details + new_mask_detail = updated_context.processed_maps_details[new_mask_rule.id.hex] + assert new_mask_detail['map_type'] == "MASK" + assert "mask_from_alpha_" in new_mask_detail['temp_processed_file'] + assert "Generated from alpha of ALBEDO" in new_mask_detail['notes'] # Check for specific note + assert new_mask_detail['status'] == 'Processed' + +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.save_image') +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.load_image') +@mock.patch('logging.info') +def test_no_alpha_channel_in_source(mock_log_info, mock_load_image, mock_save_image): + stage = AlphaExtractionToMaskStage() + + albedo_rule_id = uuid.uuid4() + albedo_fr = create_mock_file_rule_for_alpha_test(id_val=albedo_rule_id, map_type="ALBEDO") + initial_processed_details = { + albedo_fr.id.hex: {'temp_processed_file': '/fake/temp_engine_dir/processed_rgb_albedo.png', 'status': 'Processed', 'map_type': 'ALBEDO', 'source_file_path': Path('/fake/source/albedo_rgb.png')} + } + context = create_alpha_extraction_mock_context( + initial_file_rules=[albedo_fr], + initial_processed_details=initial_processed_details + ) + + mock_rgb_data = np.zeros((10, 10, 3), dtype=np.uint8) # RGB, no alpha + mock_load_image.return_value = mock_rgb_data # Only called once for check + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(Path('/fake/temp_engine_dir/processed_rgb_albedo.png')) + mock_save_image.assert_not_called() + assert len(updated_context.files_to_process) == 1 # No new MASK rule + assert albedo_fr.id.hex in updated_context.processed_maps_details + assert len(updated_context.processed_maps_details) == 1 + + +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.save_image') +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.load_image') +@mock.patch('logging.info') +def test_no_suitable_source_map_type(mock_log_info, mock_load_image, mock_save_image): + stage = AlphaExtractionToMaskStage() + + normal_rule_id = uuid.uuid4() + normal_fr = create_mock_file_rule_for_alpha_test(id_val=normal_rule_id, map_type="NORMAL") + initial_processed_details = { + normal_fr.id.hex: {'temp_processed_file': '/fake/temp_engine_dir/processed_normal.png', 'status': 'Processed', 'map_type': 'NORMAL'} + } + context = create_alpha_extraction_mock_context( + initial_file_rules=[normal_fr], + initial_processed_details=initial_processed_details + ) + + updated_context = stage.execute(context) + + mock_load_image.assert_not_called() + mock_save_image.assert_not_called() + assert len(updated_context.files_to_process) == 1 + assert normal_fr.id.hex in updated_context.processed_maps_details + assert len(updated_context.processed_maps_details) == 1 + +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.save_image') +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.load_image') +@mock.patch('logging.warning') # Expect a warning log +def test_load_image_fails(mock_log_warning, mock_load_image, mock_save_image): + stage = AlphaExtractionToMaskStage() + + albedo_rule_id = uuid.uuid4() + albedo_fr = create_mock_file_rule_for_alpha_test(id_val=albedo_rule_id, map_type="ALBEDO") + initial_processed_details = { + albedo_fr.id.hex: {'temp_processed_file': '/fake/temp_engine_dir/processed_albedo_load_fail.png', 'status': 'Processed', 'map_type': 'ALBEDO', 'source_file_path': Path('/fake/source/albedo_load_fail.png')} + } + context = create_alpha_extraction_mock_context( + initial_file_rules=[albedo_fr], + initial_processed_details=initial_processed_details + ) + + mock_load_image.return_value = None # Simulate load failure + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(Path('/fake/temp_engine_dir/processed_albedo_load_fail.png')) + mock_save_image.assert_not_called() + assert len(updated_context.files_to_process) == 1 + assert albedo_fr.id.hex in updated_context.processed_maps_details + assert len(updated_context.processed_maps_details) == 1 + mock_log_warning.assert_called_once() # Check that a warning was logged + +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.save_image') +@mock.patch('processing.pipeline.stages.alpha_extraction_to_mask.ipu.load_image') +@mock.patch('logging.error') # Expect an error log +def test_save_image_fails(mock_log_error, mock_load_image, mock_save_image): + stage = AlphaExtractionToMaskStage() + + albedo_rule_id = uuid.uuid4() + albedo_fr = create_mock_file_rule_for_alpha_test(id_val=albedo_rule_id, map_type="ALBEDO") + initial_processed_details = { + albedo_fr.id.hex: {'temp_processed_file': '/fake/temp_engine_dir/processed_albedo_save_fail.png', 'status': 'Processed', 'map_type': 'ALBEDO', 'source_file_path': Path('/fake/source/albedo_save_fail.png')} + } + context = create_alpha_extraction_mock_context( + initial_file_rules=[albedo_fr], + initial_processed_details=initial_processed_details + ) + + mock_rgba_data = np.zeros((10, 10, 4), dtype=np.uint8) + mock_rgba_data[:, :, 3] = 128 + mock_load_image.side_effect = [mock_rgba_data, mock_rgba_data] # Load succeeds + + mock_save_image.return_value = False # Simulate save failure + + updated_context = stage.execute(context) + + assert mock_load_image.call_count == 2 + mock_save_image.assert_called_once() # Save was attempted + + assert len(updated_context.files_to_process) == 1 # No new MASK rule should be successfully added and detailed + + # Check that no new MASK details were added, or if they were, they reflect failure. + # The current stage logic returns context early, so no new rule or details should be present. + mask_rule_found = any(fr.map_type == "MASK" for fr in updated_context.files_to_process) + assert not mask_rule_found + + mask_details_found = any( + details['map_type'] == "MASK" + for fr_id, details in updated_context.processed_maps_details.items() + if fr_id != albedo_fr.id.hex # Exclude the original albedo + ) + assert not mask_details_found + mock_log_error.assert_called_once() # Check that an error was logged \ No newline at end of file diff --git a/tests/processing/pipeline/stages/test_asset_skip_logic.py b/tests/processing/pipeline/stages/test_asset_skip_logic.py new file mode 100644 index 0000000..388cc8c --- /dev/null +++ b/tests/processing/pipeline/stages/test_asset_skip_logic.py @@ -0,0 +1,213 @@ +import pytest +from unittest import mock +from pathlib import Path +from typing import Dict, Optional, Any + +from processing.pipeline.stages.asset_skip_logic import AssetSkipLogicStage +from processing.pipeline.asset_context import AssetProcessingContext +from rule_structure import AssetRule, SourceRule +from configuration import Configuration, GeneralSettings + +# Helper function to create a mock AssetProcessingContext +def create_skip_logic_mock_context( + effective_supplier: Optional[str] = "ValidSupplier", + asset_process_status: str = "PENDING", + overwrite_existing: bool = False, + asset_name: str = "TestAssetSkip" +) -> AssetProcessingContext: + mock_asset_rule = mock.MagicMock(spec=AssetRule) + mock_asset_rule.name = asset_name + mock_asset_rule.process_status = asset_process_status + mock_asset_rule.source_path = "fake/source" # Added for completeness + mock_asset_rule.output_path = "fake/output" # Added for completeness + mock_asset_rule.maps = [] # Added for completeness + mock_asset_rule.metadata = {} # Added for completeness + mock_asset_rule.material_name = None # Added for completeness + mock_asset_rule.notes = None # Added for completeness + mock_asset_rule.tags = [] # Added for completeness + mock_asset_rule.enabled = True # Added for completeness + + + mock_source_rule = mock.MagicMock(spec=SourceRule) + mock_source_rule.name = "TestSourceRule" # Added for completeness + mock_source_rule.path = "fake/source_rule_path" # Added for completeness + mock_source_rule.default_supplier = None # Added for completeness + mock_source_rule.assets = [mock_asset_rule] # Added for completeness + mock_source_rule.enabled = True # Added for completeness + + mock_general_settings = mock.MagicMock(spec=GeneralSettings) + mock_general_settings.overwrite_existing = overwrite_existing + + mock_config = mock.MagicMock(spec=Configuration) + mock_config.general_settings = mock_general_settings + mock_config.suppliers = {"ValidSupplier": mock.MagicMock()} + + context = AssetProcessingContext( + source_rule=mock_source_rule, + asset_rule=mock_asset_rule, + workspace_path=Path("/fake/workspace"), + engine_temp_dir=Path("/fake/temp"), + output_base_path=Path("/fake/output"), + effective_supplier=effective_supplier, + asset_metadata={}, + processed_maps_details={}, + merged_maps_details={}, + files_to_process=[], + loaded_data_cache={}, + config_obj=mock_config, + status_flags={}, + incrementing_value=None, + sha5_value=None # Corrected from sha5_value to sha256_value if that's the actual field + ) + # Ensure status_flags is initialized if AssetSkipLogicStage expects it + # context.status_flags = {} # Already done in constructor + return context +@mock.patch('logging.info') +def test_skip_due_to_missing_supplier(mock_log_info): + """ + Test that the asset is skipped if effective_supplier is None. + """ + stage = AssetSkipLogicStage() + context = create_skip_logic_mock_context(effective_supplier=None, asset_name="MissingSupplierAsset") + + updated_context = stage.execute(context) + + assert updated_context.status_flags.get('skip_asset') is True + assert updated_context.status_flags.get('skip_reason') == "Invalid or missing supplier" + mock_log_info.assert_any_call(f"Asset 'MissingSupplierAsset': Skipping due to missing or invalid supplier.") + +@mock.patch('logging.info') +def test_skip_due_to_process_status_skip(mock_log_info): + """ + Test that the asset is skipped if asset_rule.process_status is "SKIP". + """ + stage = AssetSkipLogicStage() + context = create_skip_logic_mock_context(asset_process_status="SKIP", asset_name="SkipStatusAsset") + + updated_context = stage.execute(context) + + assert updated_context.status_flags.get('skip_asset') is True + assert updated_context.status_flags.get('skip_reason') == "Process status set to SKIP" + mock_log_info.assert_any_call(f"Asset 'SkipStatusAsset': Skipping because process_status is 'SKIP'.") + +@mock.patch('logging.info') +def test_skip_due_to_processed_and_overwrite_disabled(mock_log_info): + """ + Test that the asset is skipped if asset_rule.process_status is "PROCESSED" + and overwrite_existing is False. + """ + stage = AssetSkipLogicStage() + context = create_skip_logic_mock_context( + asset_process_status="PROCESSED", + overwrite_existing=False, + asset_name="ProcessedNoOverwriteAsset" + ) + + updated_context = stage.execute(context) + + assert updated_context.status_flags.get('skip_asset') is True + assert updated_context.status_flags.get('skip_reason') == "Already processed, overwrite disabled" + mock_log_info.assert_any_call(f"Asset 'ProcessedNoOverwriteAsset': Skipping because already processed and overwrite is disabled.") + +@mock.patch('logging.info') +def test_no_skip_when_processed_and_overwrite_enabled(mock_log_info): + """ + Test that the asset is NOT skipped if asset_rule.process_status is "PROCESSED" + but overwrite_existing is True. + """ + stage = AssetSkipLogicStage() + context = create_skip_logic_mock_context( + asset_process_status="PROCESSED", + overwrite_existing=True, + effective_supplier="ValidSupplier", # Ensure supplier is valid + asset_name="ProcessedOverwriteAsset" + ) + + updated_context = stage.execute(context) + + assert updated_context.status_flags.get('skip_asset', False) is False # Default to False if key not present + # No specific skip_reason to check if not skipped + # Check that no skip log message was called for this specific reason + for call_args in mock_log_info.call_args_list: + assert "Skipping because already processed and overwrite is disabled" not in call_args[0][0] + assert "Skipping due to missing or invalid supplier" not in call_args[0][0] + assert "Skipping because process_status is 'SKIP'" not in call_args[0][0] + + +@mock.patch('logging.info') +def test_no_skip_when_process_status_pending(mock_log_info): + """ + Test that the asset is NOT skipped if asset_rule.process_status is "PENDING". + """ + stage = AssetSkipLogicStage() + context = create_skip_logic_mock_context( + asset_process_status="PENDING", + effective_supplier="ValidSupplier", # Ensure supplier is valid + asset_name="PendingAsset" + ) + + updated_context = stage.execute(context) + + assert updated_context.status_flags.get('skip_asset', False) is False + # Check that no skip log message was called + for call_args in mock_log_info.call_args_list: + assert "Skipping" not in call_args[0][0] + + +@mock.patch('logging.info') +def test_no_skip_when_process_status_failed_previously(mock_log_info): + """ + Test that the asset is NOT skipped if asset_rule.process_status is "FAILED_PREVIOUSLY". + """ + stage = AssetSkipLogicStage() + context = create_skip_logic_mock_context( + asset_process_status="FAILED_PREVIOUSLY", + effective_supplier="ValidSupplier", # Ensure supplier is valid + asset_name="FailedPreviouslyAsset" + ) + + updated_context = stage.execute(context) + + assert updated_context.status_flags.get('skip_asset', False) is False + # Check that no skip log message was called + for call_args in mock_log_info.call_args_list: + assert "Skipping" not in call_args[0][0] + +@mock.patch('logging.info') +def test_no_skip_when_process_status_other_valid_status(mock_log_info): + """ + Test that the asset is NOT skipped for other valid, non-skip process statuses. + """ + stage = AssetSkipLogicStage() + context = create_skip_logic_mock_context( + asset_process_status="READY_FOR_PROCESSING", # Example of another non-skip status + effective_supplier="ValidSupplier", + asset_name="ReadyAsset" + ) + updated_context = stage.execute(context) + assert updated_context.status_flags.get('skip_asset', False) is False + for call_args in mock_log_info.call_args_list: + assert "Skipping" not in call_args[0][0] + +@mock.patch('logging.info') +def test_skip_asset_flag_initialized_if_not_present(mock_log_info): + """ + Test that 'skip_asset' is initialized to False in status_flags if not skipped and not present. + """ + stage = AssetSkipLogicStage() + context = create_skip_logic_mock_context( + asset_process_status="PENDING", + effective_supplier="ValidSupplier", + asset_name="InitFlagAsset" + ) + # Ensure status_flags is empty before execute + context.status_flags = {} + + updated_context = stage.execute(context) + + # If not skipped, 'skip_asset' should be explicitly False. + assert updated_context.status_flags.get('skip_asset') is False + # No skip reason should be set + assert 'skip_reason' not in updated_context.status_flags + for call_args in mock_log_info.call_args_list: + assert "Skipping" not in call_args[0][0] \ No newline at end of file diff --git a/tests/processing/pipeline/stages/test_file_rule_filter.py b/tests/processing/pipeline/stages/test_file_rule_filter.py new file mode 100644 index 0000000..4a79308 --- /dev/null +++ b/tests/processing/pipeline/stages/test_file_rule_filter.py @@ -0,0 +1,330 @@ +import pytest +from unittest import mock +from pathlib import Path +import uuid +from typing import Optional # Added Optional for type hinting + +from processing.pipeline.stages.file_rule_filter import FileRuleFilterStage +from processing.pipeline.asset_context import AssetProcessingContext +from rule_structure import AssetRule, SourceRule, FileRule # FileRule is key here +from configuration import Configuration # Minimal config needed + +def create_mock_file_rule( + id_val: Optional[uuid.UUID] = None, + map_type: str = "Diffuse", + filename_pattern: str = "*.tif", + item_type: str = "MAP_COL", # e.g., MAP_COL, FILE_IGNORE + active: bool = True +) -> mock.MagicMock: # Return MagicMock to easily set other attributes if needed + mock_fr = mock.MagicMock(spec=FileRule) + mock_fr.id = id_val if id_val else uuid.uuid4() + mock_fr.map_type = map_type + mock_fr.filename_pattern = filename_pattern + mock_fr.item_type = item_type + mock_fr.active = active + return mock_fr + +def create_file_filter_mock_context( + file_rules_list: Optional[list] = None, # List of mock FileRule objects + skip_asset_flag: bool = False, + asset_name: str = "FileFilterAsset" +) -> AssetProcessingContext: + mock_asset_rule = mock.MagicMock(spec=AssetRule) + mock_asset_rule.name = asset_name + mock_asset_rule.file_rules = file_rules_list if file_rules_list is not None else [] + + mock_source_rule = mock.MagicMock(spec=SourceRule) + mock_config = mock.MagicMock(spec=Configuration) + + context = AssetProcessingContext( + source_rule=mock_source_rule, + asset_rule=mock_asset_rule, + workspace_path=Path("/fake/workspace"), + engine_temp_dir=Path("/fake/temp"), + output_base_path=Path("/fake/output"), + effective_supplier="ValidSupplier", # Assume valid for this stage + asset_metadata={'asset_name': asset_name}, # Assume metadata init happened + processed_maps_details={}, + merged_maps_details={}, + files_to_process=[], # Stage will populate this + loaded_data_cache={}, + config_obj=mock_config, + status_flags={'skip_asset': skip_asset_flag}, + incrementing_value=None, + sha5_value=None # Corrected from sha5_value to sha256_value based on AssetProcessingContext + ) + return context +# Test Cases for FileRuleFilterStage.execute() + +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_file_rule_filter_asset_skipped(mock_log_debug, mock_log_info): + """ + Test case: Asset Skipped - status_flags['skip_asset'] is True. + Assert context.files_to_process remains empty. + """ + stage = FileRuleFilterStage() + context = create_file_filter_mock_context(skip_asset_flag=True) + + updated_context = stage.execute(context) + + assert len(updated_context.files_to_process) == 0 + mock_log_info.assert_any_call(f"Asset '{context.asset_rule.name}': Skipping file rule filtering as 'skip_asset' is True.") + +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_file_rule_filter_no_file_rules(mock_log_debug, mock_log_info): + """ + Test case: No File Rules - asset_rule.file_rules is empty. + Assert context.files_to_process is empty. + """ + stage = FileRuleFilterStage() + context = create_file_filter_mock_context(file_rules_list=[]) + + updated_context = stage.execute(context) + + assert len(updated_context.files_to_process) == 0 + mock_log_info.assert_any_call(f"Asset '{context.asset_rule.name}': No file rules defined. Skipping file rule filtering.") + +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_file_rule_filter_only_active_processable_rules(mock_log_debug, mock_log_info): + """ + Test case: Only Active, Processable Rules - All FileRules are active=True and item_type="MAP_COL". + Assert all are added to context.files_to_process. + """ + stage = FileRuleFilterStage() + fr1 = create_mock_file_rule(filename_pattern="diffuse.png", item_type="MAP_COL", active=True) + fr2 = create_mock_file_rule(filename_pattern="normal.png", item_type="MAP_COL", active=True) + context = create_file_filter_mock_context(file_rules_list=[fr1, fr2]) + + updated_context = stage.execute(context) + + assert len(updated_context.files_to_process) == 2 + assert fr1 in updated_context.files_to_process + assert fr2 in updated_context.files_to_process + mock_log_info.assert_any_call(f"Asset '{context.asset_rule.name}': 2 file rules queued for processing after filtering.") + +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_file_rule_filter_inactive_rules(mock_log_debug, mock_log_info): + """ + Test case: Inactive Rules - Some FileRules have active=False. + Assert only active rules are added. + """ + stage = FileRuleFilterStage() + fr_active = create_mock_file_rule(filename_pattern="active.png", item_type="MAP_COL", active=True) + fr_inactive = create_mock_file_rule(filename_pattern="inactive.png", item_type="MAP_COL", active=False) + fr_another_active = create_mock_file_rule(filename_pattern="another_active.jpg", item_type="MAP_COL", active=True) + context = create_file_filter_mock_context(file_rules_list=[fr_active, fr_inactive, fr_another_active]) + + updated_context = stage.execute(context) + + assert len(updated_context.files_to_process) == 2 + assert fr_active in updated_context.files_to_process + assert fr_another_active in updated_context.files_to_process + assert fr_inactive not in updated_context.files_to_process + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Skipping inactive file rule: '{fr_inactive.filename_pattern}'") + mock_log_info.assert_any_call(f"Asset '{context.asset_rule.name}': 2 file rules queued for processing after filtering.") + +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_file_rule_filter_file_ignore_simple_match(mock_log_debug, mock_log_info): + """ + Test case: FILE_IGNORE Rule (Simple Match). + One FILE_IGNORE rule with filename_pattern="*_ignore.png". + One MAP_COL rule with filename_pattern="diffuse_ignore.png". + One MAP_COL rule with filename_pattern="normal_process.png". + Assert only "normal_process.png" rule is added. + """ + stage = FileRuleFilterStage() + fr_ignore = create_mock_file_rule(filename_pattern="*_ignore.png", item_type="FILE_IGNORE", active=True) + fr_ignored_map = create_mock_file_rule(filename_pattern="diffuse_ignore.png", item_type="MAP_COL", active=True) + fr_process_map = create_mock_file_rule(filename_pattern="normal_process.png", item_type="MAP_COL", active=True) + context = create_file_filter_mock_context(file_rules_list=[fr_ignore, fr_ignored_map, fr_process_map]) + + updated_context = stage.execute(context) + + assert len(updated_context.files_to_process) == 1 + assert fr_process_map in updated_context.files_to_process + assert fr_ignored_map not in updated_context.files_to_process + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Registering ignore pattern: '{fr_ignore.filename_pattern}'") + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Skipping file rule '{fr_ignored_map.filename_pattern}' due to matching ignore pattern.") + mock_log_info.assert_any_call(f"Asset '{context.asset_rule.name}': 1 file rules queued for processing after filtering.") + +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_file_rule_filter_file_ignore_glob_pattern(mock_log_debug, mock_log_info): + """ + Test case: FILE_IGNORE Rule (Glob Pattern). + One FILE_IGNORE rule with filename_pattern="*_ignore.*". + MAP_COL rules: "tex_ignore.tif", "tex_process.png". + Assert only "tex_process.png" rule is added. + """ + stage = FileRuleFilterStage() + fr_ignore_glob = create_mock_file_rule(filename_pattern="*_ignore.*", item_type="FILE_IGNORE", active=True) + fr_ignored_tif = create_mock_file_rule(filename_pattern="tex_ignore.tif", item_type="MAP_COL", active=True) + fr_process_png = create_mock_file_rule(filename_pattern="tex_process.png", item_type="MAP_COL", active=True) + context = create_file_filter_mock_context(file_rules_list=[fr_ignore_glob, fr_ignored_tif, fr_process_png]) + + updated_context = stage.execute(context) + + assert len(updated_context.files_to_process) == 1 + assert fr_process_png in updated_context.files_to_process + assert fr_ignored_tif not in updated_context.files_to_process + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Registering ignore pattern: '{fr_ignore_glob.filename_pattern}'") + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Skipping file rule '{fr_ignored_tif.filename_pattern}' due to matching ignore pattern.") + mock_log_info.assert_any_call(f"Asset '{context.asset_rule.name}': 1 file rules queued for processing after filtering.") + +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_file_rule_filter_multiple_file_ignore_rules(mock_log_debug, mock_log_info): + """ + Test case: Multiple FILE_IGNORE Rules. + Test with several ignore patterns and ensure they are all respected. + """ + stage = FileRuleFilterStage() + fr_ignore1 = create_mock_file_rule(filename_pattern="*.tmp", item_type="FILE_IGNORE", active=True) + fr_ignore2 = create_mock_file_rule(filename_pattern="backup_*", item_type="FILE_IGNORE", active=True) + fr_ignore3 = create_mock_file_rule(filename_pattern="*_old.png", item_type="FILE_IGNORE", active=True) + + fr_map_ignored1 = create_mock_file_rule(filename_pattern="data.tmp", item_type="MAP_COL", active=True) + fr_map_ignored2 = create_mock_file_rule(filename_pattern="backup_diffuse.jpg", item_type="MAP_COL", active=True) + fr_map_ignored3 = create_mock_file_rule(filename_pattern="normal_old.png", item_type="MAP_COL", active=True) + fr_map_process = create_mock_file_rule(filename_pattern="final_texture.tif", item_type="MAP_COL", active=True) + + context = create_file_filter_mock_context(file_rules_list=[ + fr_ignore1, fr_ignore2, fr_ignore3, + fr_map_ignored1, fr_map_ignored2, fr_map_ignored3, fr_map_process + ]) + + updated_context = stage.execute(context) + + assert len(updated_context.files_to_process) == 1 + assert fr_map_process in updated_context.files_to_process + assert fr_map_ignored1 not in updated_context.files_to_process + assert fr_map_ignored2 not in updated_context.files_to_process + assert fr_map_ignored3 not in updated_context.files_to_process + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Registering ignore pattern: '{fr_ignore1.filename_pattern}'") + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Registering ignore pattern: '{fr_ignore2.filename_pattern}'") + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Registering ignore pattern: '{fr_ignore3.filename_pattern}'") + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Skipping file rule '{fr_map_ignored1.filename_pattern}' due to matching ignore pattern.") + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Skipping file rule '{fr_map_ignored2.filename_pattern}' due to matching ignore pattern.") + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Skipping file rule '{fr_map_ignored3.filename_pattern}' due to matching ignore pattern.") + mock_log_info.assert_any_call(f"Asset '{context.asset_rule.name}': 1 file rules queued for processing after filtering.") + +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_file_rule_filter_file_ignore_rule_is_inactive(mock_log_debug, mock_log_info): + """ + Test case: FILE_IGNORE Rule is Inactive. + An ignore rule itself is active=False. Assert its pattern is NOT used for filtering. + """ + stage = FileRuleFilterStage() + fr_inactive_ignore = create_mock_file_rule(filename_pattern="*_ignore.tif", item_type="FILE_IGNORE", active=False) + fr_should_process1 = create_mock_file_rule(filename_pattern="diffuse_ignore.tif", item_type="MAP_COL", active=True) # Should be processed + fr_should_process2 = create_mock_file_rule(filename_pattern="normal_ok.png", item_type="MAP_COL", active=True) + context = create_file_filter_mock_context(file_rules_list=[fr_inactive_ignore, fr_should_process1, fr_should_process2]) + + updated_context = stage.execute(context) + + assert len(updated_context.files_to_process) == 2 + assert fr_should_process1 in updated_context.files_to_process + assert fr_should_process2 in updated_context.files_to_process + # Ensure the inactive ignore rule's pattern was not registered + # We check this by ensuring no debug log for registering *that specific* pattern was made. + # A more robust way would be to check mock_log_debug.call_args_list, but this is simpler for now. + for call in mock_log_debug.call_args_list: + args, kwargs = call + if "Registering ignore pattern" in args[0] and fr_inactive_ignore.filename_pattern in args[0]: + pytest.fail(f"Inactive ignore pattern '{fr_inactive_ignore.filename_pattern}' was incorrectly registered.") + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Skipping inactive file rule: '{fr_inactive_ignore.filename_pattern}' (type: FILE_IGNORE)") + mock_log_info.assert_any_call(f"Asset '{context.asset_rule.name}': 2 file rules queued for processing after filtering.") + + +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_file_rule_filter_no_file_ignore_rules(mock_log_debug, mock_log_info): + """ + Test case: No FILE_IGNORE Rules. + All rules are MAP_COL or other processable types. + Assert all active, processable rules are included. + """ + stage = FileRuleFilterStage() + fr1 = create_mock_file_rule(filename_pattern="diffuse.png", item_type="MAP_COL", active=True) + fr2 = create_mock_file_rule(filename_pattern="normal.png", item_type="MAP_COL", active=True) + fr_other_type = create_mock_file_rule(filename_pattern="spec.tif", item_type="MAP_SPEC", active=True) # Assuming MAP_SPEC is processable + fr_inactive = create_mock_file_rule(filename_pattern="ao.jpg", item_type="MAP_AO", active=False) + + context = create_file_filter_mock_context(file_rules_list=[fr1, fr2, fr_other_type, fr_inactive]) + + updated_context = stage.execute(context) + + assert len(updated_context.files_to_process) == 3 + assert fr1 in updated_context.files_to_process + assert fr2 in updated_context.files_to_process + assert fr_other_type in updated_context.files_to_process + assert fr_inactive not in updated_context.files_to_process + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Skipping inactive file rule: '{fr_inactive.filename_pattern}'") + mock_log_info.assert_any_call(f"Asset '{context.asset_rule.name}': 3 file rules queued for processing after filtering.") + +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_file_rule_filter_item_type_not_processable(mock_log_debug, mock_log_info): + """ + Test case: Item type is not processable (e.g., not MAP_COL, MAP_AO etc., but something else like 'METADATA_ONLY'). + Assert such rules are not added to files_to_process, unless they are FILE_IGNORE. + """ + stage = FileRuleFilterStage() + fr_processable = create_mock_file_rule(filename_pattern="diffuse.png", item_type="MAP_COL", active=True) + fr_not_processable = create_mock_file_rule(filename_pattern="info.txt", item_type="METADATA_ONLY", active=True) + fr_ignore = create_mock_file_rule(filename_pattern="*.bak", item_type="FILE_IGNORE", active=True) + fr_ignored_by_bak = create_mock_file_rule(filename_pattern="diffuse.bak", item_type="MAP_COL", active=True) + + context = create_file_filter_mock_context(file_rules_list=[fr_processable, fr_not_processable, fr_ignore, fr_ignored_by_bak]) + + updated_context = stage.execute(context) + + assert len(updated_context.files_to_process) == 1 + assert fr_processable in updated_context.files_to_process + assert fr_not_processable not in updated_context.files_to_process + assert fr_ignored_by_bak not in updated_context.files_to_process + + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Registering ignore pattern: '{fr_ignore.filename_pattern}'") + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Skipping file rule '{fr_not_processable.filename_pattern}' as its item_type '{fr_not_processable.item_type}' is not processable.") + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Skipping file rule '{fr_ignored_by_bak.filename_pattern}' due to matching ignore pattern.") + mock_log_info.assert_any_call(f"Asset '{context.asset_rule.name}': 1 file rules queued for processing after filtering.") + +# Example tests from instructions (can be adapted or used as a base) +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_file_rule_filter_basic_active_example(mock_log_debug, mock_log_info): # Renamed to avoid conflict + stage = FileRuleFilterStage() + fr1 = create_mock_file_rule(filename_pattern="diffuse.png", item_type="MAP_COL", active=True) + fr2 = create_mock_file_rule(filename_pattern="normal.png", item_type="MAP_COL", active=True) + context = create_file_filter_mock_context(file_rules_list=[fr1, fr2]) + + updated_context = stage.execute(context) + + assert len(updated_context.files_to_process) == 2 + assert fr1 in updated_context.files_to_process + assert fr2 in updated_context.files_to_process + mock_log_info.assert_any_call(f"Asset '{context.asset_rule.name}': 2 file rules queued for processing after filtering.") + +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_file_rule_filter_with_file_ignore_example(mock_log_debug, mock_log_info): # Renamed to avoid conflict + stage = FileRuleFilterStage() + fr_ignore = create_mock_file_rule(filename_pattern="*_ignore.tif", item_type="FILE_IGNORE", active=True) + fr_process = create_mock_file_rule(filename_pattern="diffuse_ok.tif", item_type="MAP_COL", active=True) + fr_skip = create_mock_file_rule(filename_pattern="normal_ignore.tif", item_type="MAP_COL", active=True) + context = create_file_filter_mock_context(file_rules_list=[fr_ignore, fr_process, fr_skip]) + + updated_context = stage.execute(context) + + assert len(updated_context.files_to_process) == 1 + assert fr_process in updated_context.files_to_process + assert fr_skip not in updated_context.files_to_process + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Registering ignore pattern: '{fr_ignore.filename_pattern}'") + mock_log_debug.assert_any_call(f"Asset '{context.asset_rule.name}': Skipping file rule '{fr_skip.filename_pattern}' due to matching ignore pattern.") + mock_log_info.assert_any_call(f"Asset '{context.asset_rule.name}': 1 file rules queued for processing after filtering.") \ No newline at end of file diff --git a/tests/processing/pipeline/stages/test_gloss_to_rough_conversion.py b/tests/processing/pipeline/stages/test_gloss_to_rough_conversion.py new file mode 100644 index 0000000..934ad2c --- /dev/null +++ b/tests/processing/pipeline/stages/test_gloss_to_rough_conversion.py @@ -0,0 +1,486 @@ +import pytest +from unittest import mock +from pathlib import Path +import uuid +import numpy as np +from typing import Optional, List, Dict + +from processing.pipeline.stages.gloss_to_rough_conversion import GlossToRoughConversionStage +from processing.pipeline.asset_context import AssetProcessingContext +from rule_structure import AssetRule, SourceRule, FileRule +from configuration import Configuration, GeneralSettings +# No direct ipu import needed in test if we mock its usage by the stage + +def create_mock_file_rule_for_gloss_test( + id_val: Optional[uuid.UUID] = None, + map_type: str = "GLOSS", # Test with GLOSS and other types + filename_pattern: str = "gloss.png" +) -> mock.MagicMock: + mock_fr = mock.MagicMock(spec=FileRule) + mock_fr.id = id_val if id_val else uuid.uuid4() + mock_fr.map_type = map_type + mock_fr.filename_pattern = filename_pattern + mock_fr.item_type = "MAP_COL" + mock_fr.active = True + return mock_fr + +def create_gloss_conversion_mock_context( + initial_file_rules: Optional[List[FileRule]] = None, # Type hint corrected + initial_processed_details: Optional[Dict] = None, # Type hint corrected + skip_asset_flag: bool = False, + asset_name: str = "GlossAsset", + # Add a mock for general_settings if your stage checks a global flag + # convert_gloss_globally: bool = True +) -> AssetProcessingContext: + mock_asset_rule = mock.MagicMock(spec=AssetRule) + mock_asset_rule.name = asset_name + mock_asset_rule.file_rules = initial_file_rules if initial_file_rules is not None else [] + + mock_source_rule = mock.MagicMock(spec=SourceRule) + + mock_gs = mock.MagicMock(spec=GeneralSettings) + # if your stage uses a global flag: + # mock_gs.convert_gloss_to_rough_globally = convert_gloss_globally + + mock_config = mock.MagicMock(spec=Configuration) + mock_config.general_settings = mock_gs + + + context = AssetProcessingContext( + source_rule=mock_source_rule, + asset_rule=mock_asset_rule, + workspace_path=Path("/fake/workspace"), + engine_temp_dir=Path("/fake/temp_engine_dir"), # Important for new temp file paths + output_base_path=Path("/fake/output"), + effective_supplier="ValidSupplier", + asset_metadata={'asset_name': asset_name}, + processed_maps_details=initial_processed_details if initial_processed_details is not None else {}, + merged_maps_details={}, + files_to_process=list(initial_file_rules) if initial_file_rules else [], # Stage modifies this list + loaded_data_cache={}, + config_obj=mock_config, + status_flags={'skip_asset': skip_asset_flag}, + incrementing_value=None, # Added as per AssetProcessingContext definition + sha5_value=None # Added as per AssetProcessingContext definition + ) + return context + +# Unit tests will be added below +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.save_image') +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.load_image') +def test_asset_skipped(mock_load_image, mock_save_image): + """ + Test that if 'skip_asset' is True, no processing occurs. + """ + stage = GlossToRoughConversionStage() + + gloss_rule_id = uuid.uuid4() + gloss_fr = create_mock_file_rule_for_gloss_test(id_val=gloss_rule_id, map_type="GLOSS") + + initial_details = { + gloss_fr.id.hex: {'temp_processed_file': '/fake/temp_engine_dir/processed_gloss_map.png', 'status': 'Processed', 'map_type': 'GLOSS'} + } + context = create_gloss_conversion_mock_context( + initial_file_rules=[gloss_fr], + initial_processed_details=initial_details, + skip_asset_flag=True # Asset is skipped + ) + + # Keep a copy of files_to_process and processed_maps_details to compare + original_files_to_process = list(context.files_to_process) + original_processed_maps_details = context.processed_maps_details.copy() + + updated_context = stage.execute(context) + + mock_load_image.assert_not_called() + mock_save_image.assert_not_called() + + assert updated_context.files_to_process == original_files_to_process, "files_to_process should not change if asset is skipped" + assert updated_context.processed_maps_details == original_processed_maps_details, "processed_maps_details should not change if asset is skipped" + assert updated_context.status_flags['skip_asset'] is True +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.save_image') +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.load_image') +def test_no_gloss_map_present(mock_load_image, mock_save_image): + """ + Test that if no GLOSS maps are in files_to_process, no conversion occurs. + """ + stage = GlossToRoughConversionStage() + + normal_rule_id = uuid.uuid4() + normal_fr = create_mock_file_rule_for_gloss_test(id_val=normal_rule_id, map_type="NORMAL", filename_pattern="normal.png") + albedo_fr = create_mock_file_rule_for_gloss_test(map_type="ALBEDO", filename_pattern="albedo.jpg") + + initial_details = { + normal_fr.id.hex: {'temp_processed_file': '/fake/temp_engine_dir/processed_normal_map.png', 'status': 'Processed', 'map_type': 'NORMAL'} + } + context = create_gloss_conversion_mock_context( + initial_file_rules=[normal_fr, albedo_fr], + initial_processed_details=initial_details + ) + + original_files_to_process = list(context.files_to_process) + original_processed_maps_details = context.processed_maps_details.copy() + + updated_context = stage.execute(context) + + mock_load_image.assert_not_called() + mock_save_image.assert_not_called() + + assert updated_context.files_to_process == original_files_to_process, "files_to_process should not change if no GLOSS maps are present" + assert updated_context.processed_maps_details == original_processed_maps_details, "processed_maps_details should not change if no GLOSS maps are present" + + # Ensure map types of existing rules are unchanged + for fr_in_list in updated_context.files_to_process: + if fr_in_list.id == normal_fr.id: + assert fr_in_list.map_type == "NORMAL" + elif fr_in_list.id == albedo_fr.id: + assert fr_in_list.map_type == "ALBEDO" +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.logging') # Mock logging +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.save_image') +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.load_image') +def test_gloss_conversion_uint8_success(mock_load_image, mock_save_image, mock_logging): + """ + Test successful conversion of a GLOSS map (uint8 data) to ROUGHNESS. + """ + stage = GlossToRoughConversionStage() + + gloss_rule_id = uuid.uuid4() + # Use a distinct filename for the gloss map to ensure correct path construction + gloss_fr = create_mock_file_rule_for_gloss_test(id_val=gloss_rule_id, map_type="GLOSS", filename_pattern="my_gloss_map.png") + other_fr_id = uuid.uuid4() + other_fr = create_mock_file_rule_for_gloss_test(id_val=other_fr_id, map_type="NORMAL", filename_pattern="normal_map.png") + + initial_gloss_temp_path = Path("/fake/temp_engine_dir/processed_gloss_map.png") + initial_other_temp_path = Path("/fake/temp_engine_dir/processed_normal_map.png") + + initial_details = { + gloss_fr.id.hex: {'temp_processed_file': str(initial_gloss_temp_path), 'status': 'Processed', 'map_type': 'GLOSS'}, + other_fr.id.hex: {'temp_processed_file': str(initial_other_temp_path), 'status': 'Processed', 'map_type': 'NORMAL'} + } + context = create_gloss_conversion_mock_context( + initial_file_rules=[gloss_fr, other_fr], + initial_processed_details=initial_details + ) + + mock_loaded_gloss_data = np.array([10, 50, 250], dtype=np.uint8) + mock_load_image.return_value = mock_loaded_gloss_data + mock_save_image.return_value = True # Simulate successful save + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(initial_gloss_temp_path) + + # Check that save_image was called with inverted data and correct path + expected_inverted_data = 255 - mock_loaded_gloss_data + + # call_args[0] is a tuple of positional args, call_args[1] is a dict of kwargs + saved_path_arg = mock_save_image.call_args[0][0] + saved_data_arg = mock_save_image.call_args[0][1] + + assert np.array_equal(saved_data_arg, expected_inverted_data), "Image data passed to save_image is not correctly inverted." + assert "rough_from_gloss_" in saved_path_arg.name, "Saved file name should indicate conversion from gloss." + assert saved_path_arg.parent == Path("/fake/temp_engine_dir"), "Saved file should be in the engine temp directory." + # Ensure the new filename is based on the original gloss map's ID for uniqueness + assert gloss_fr.id.hex in saved_path_arg.name + + # Check context.files_to_process + assert len(updated_context.files_to_process) == 2, "Number of file rules in context should remain the same." + converted_rule_found = False + other_rule_untouched = False + for fr_in_list in updated_context.files_to_process: + if fr_in_list.id == gloss_fr.id: # Should be the same rule object, modified + assert fr_in_list.map_type == "ROUGHNESS", "GLOSS map_type should be changed to ROUGHNESS." + # Check if filename_pattern was updated (optional, depends on stage logic) + # For now, assume it might not be, as the primary identifier is map_type and ID + converted_rule_found = True + elif fr_in_list.id == other_fr.id: + assert fr_in_list.map_type == "NORMAL", "Other map_type should remain unchanged." + other_rule_untouched = True + assert converted_rule_found, "The converted GLOSS rule was not found or not updated correctly in files_to_process." + assert other_rule_untouched, "The non-GLOSS rule was modified unexpectedly." + + # Check context.processed_maps_details + assert len(updated_context.processed_maps_details) == 2, "Number of entries in processed_maps_details should remain the same." + + gloss_detail = updated_context.processed_maps_details[gloss_fr.id.hex] + assert "rough_from_gloss_" in gloss_detail['temp_processed_file'], "temp_processed_file for gloss map not updated." + assert Path(gloss_detail['temp_processed_file']).name == saved_path_arg.name, "Path in details should match saved path." + assert gloss_detail['original_map_type_before_conversion'] == "GLOSS", "original_map_type_before_conversion not set correctly." + assert "Converted from GLOSS to ROUGHNESS" in gloss_detail['notes'], "Conversion notes not added or incorrect." + assert gloss_detail['map_type'] == "ROUGHNESS", "map_type in details not updated to ROUGHNESS." + + + other_detail = updated_context.processed_maps_details[other_fr.id.hex] + assert other_detail['temp_processed_file'] == str(initial_other_temp_path), "Other map's temp_processed_file should be unchanged." + assert other_detail['map_type'] == "NORMAL", "Other map's map_type should be unchanged." + assert 'original_map_type_before_conversion' not in other_detail, "Other map should not have conversion history." + assert 'notes' not in other_detail or "Converted from GLOSS" not in other_detail['notes'], "Other map should not have conversion notes." + + mock_logging.info.assert_any_call(f"Successfully converted GLOSS map {gloss_fr.id.hex} to ROUGHNESS.") +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.logging') # Mock logging +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.save_image') +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.load_image') +def test_gloss_conversion_float_success(mock_load_image, mock_save_image, mock_logging): + """ + Test successful conversion of a GLOSS map (float data) to ROUGHNESS. + """ + stage = GlossToRoughConversionStage() + + gloss_rule_id = uuid.uuid4() + gloss_fr = create_mock_file_rule_for_gloss_test(id_val=gloss_rule_id, map_type="GLOSS", filename_pattern="gloss_float.hdr") # Example float format + + initial_gloss_temp_path = Path("/fake/temp_engine_dir/processed_gloss_float.hdr") + initial_details = { + gloss_fr.id.hex: {'temp_processed_file': str(initial_gloss_temp_path), 'status': 'Processed', 'map_type': 'GLOSS'} + } + context = create_gloss_conversion_mock_context( + initial_file_rules=[gloss_fr], + initial_processed_details=initial_details + ) + + mock_loaded_gloss_data = np.array([0.1, 0.5, 0.9], dtype=np.float32) + mock_load_image.return_value = mock_loaded_gloss_data + mock_save_image.return_value = True # Simulate successful save + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(initial_gloss_temp_path) + + expected_inverted_data = 1.0 - mock_loaded_gloss_data + + saved_path_arg = mock_save_image.call_args[0][0] + saved_data_arg = mock_save_image.call_args[0][1] + + assert np.allclose(saved_data_arg, expected_inverted_data), "Image data (float) passed to save_image is not correctly inverted." + assert "rough_from_gloss_" in saved_path_arg.name, "Saved file name should indicate conversion from gloss." + assert saved_path_arg.parent == Path("/fake/temp_engine_dir"), "Saved file should be in the engine temp directory." + assert gloss_fr.id.hex in saved_path_arg.name + + assert len(updated_context.files_to_process) == 1 + converted_rule = updated_context.files_to_process[0] + assert converted_rule.id == gloss_fr.id + assert converted_rule.map_type == "ROUGHNESS" + + gloss_detail = updated_context.processed_maps_details[gloss_fr.id.hex] + assert "rough_from_gloss_" in gloss_detail['temp_processed_file'] + assert Path(gloss_detail['temp_processed_file']).name == saved_path_arg.name + assert gloss_detail['original_map_type_before_conversion'] == "GLOSS" + assert "Converted from GLOSS to ROUGHNESS" in gloss_detail['notes'] + assert gloss_detail['map_type'] == "ROUGHNESS" + + mock_logging.info.assert_any_call(f"Successfully converted GLOSS map {gloss_fr.id.hex} to ROUGHNESS.") +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.logging') +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.save_image') +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.load_image') +def test_load_image_fails(mock_load_image, mock_save_image, mock_logging): + """ + Test behavior when ipu.load_image fails (returns None). + The original FileRule should be kept, and an error logged. + """ + stage = GlossToRoughConversionStage() + + gloss_rule_id = uuid.uuid4() + gloss_fr = create_mock_file_rule_for_gloss_test(id_val=gloss_rule_id, map_type="GLOSS", filename_pattern="gloss_fails_load.png") + + initial_gloss_temp_path = Path("/fake/temp_engine_dir/processed_gloss_fails_load.png") + initial_details = { + gloss_fr.id.hex: {'temp_processed_file': str(initial_gloss_temp_path), 'status': 'Processed', 'map_type': 'GLOSS'} + } + context = create_gloss_conversion_mock_context( + initial_file_rules=[gloss_fr], + initial_processed_details=initial_details + ) + + # Keep a copy for comparison + original_file_rule_map_type = gloss_fr.map_type + original_details_entry = context.processed_maps_details[gloss_fr.id.hex].copy() + + mock_load_image.return_value = None # Simulate load failure + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(initial_gloss_temp_path) + mock_save_image.assert_not_called() # Save should not be attempted + + # Check context.files_to_process: rule should be unchanged + assert len(updated_context.files_to_process) == 1 + processed_rule = updated_context.files_to_process[0] + assert processed_rule.id == gloss_fr.id + assert processed_rule.map_type == original_file_rule_map_type, "FileRule map_type should not change if load fails." + assert processed_rule.map_type == "GLOSS" # Explicitly check it's still GLOSS + + # Check context.processed_maps_details: details should be unchanged + current_details_entry = updated_context.processed_maps_details[gloss_fr.id.hex] + assert current_details_entry['temp_processed_file'] == str(initial_gloss_temp_path) + assert current_details_entry['map_type'] == "GLOSS" + assert 'original_map_type_before_conversion' not in current_details_entry + assert 'notes' not in current_details_entry or "Converted from GLOSS" not in current_details_entry['notes'] + + mock_logging.error.assert_called_once_with( + f"Failed to load image data for GLOSS map {gloss_fr.id.hex} from {initial_gloss_temp_path}. Skipping conversion for this map." + ) +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.logging') +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.save_image') +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.load_image') +def test_save_image_fails(mock_load_image, mock_save_image, mock_logging): + """ + Test behavior when ipu.save_image fails (returns False). + The original FileRule should be kept, and an error logged. + """ + stage = GlossToRoughConversionStage() + + gloss_rule_id = uuid.uuid4() + gloss_fr = create_mock_file_rule_for_gloss_test(id_val=gloss_rule_id, map_type="GLOSS", filename_pattern="gloss_fails_save.png") + + initial_gloss_temp_path = Path("/fake/temp_engine_dir/processed_gloss_fails_save.png") + initial_details = { + gloss_fr.id.hex: {'temp_processed_file': str(initial_gloss_temp_path), 'status': 'Processed', 'map_type': 'GLOSS'} + } + context = create_gloss_conversion_mock_context( + initial_file_rules=[gloss_fr], + initial_processed_details=initial_details + ) + + original_file_rule_map_type = gloss_fr.map_type + original_details_entry = context.processed_maps_details[gloss_fr.id.hex].copy() + + mock_loaded_gloss_data = np.array([10, 50, 250], dtype=np.uint8) + mock_load_image.return_value = mock_loaded_gloss_data + mock_save_image.return_value = False # Simulate save failure + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(initial_gloss_temp_path) + + # Check that save_image was called with correct data and path + expected_inverted_data = 255 - mock_loaded_gloss_data + # call_args[0] is a tuple of positional args + saved_path_arg = mock_save_image.call_args[0][0] + saved_data_arg = mock_save_image.call_args[0][1] + + assert np.array_equal(saved_data_arg, expected_inverted_data), "Image data passed to save_image is not correctly inverted even on failure." + assert "rough_from_gloss_" in saved_path_arg.name, "Attempted save file name should indicate conversion from gloss." + assert saved_path_arg.parent == Path("/fake/temp_engine_dir"), "Attempted save file should be in the engine temp directory." + + # Check context.files_to_process: rule should be unchanged + assert len(updated_context.files_to_process) == 1 + processed_rule = updated_context.files_to_process[0] + assert processed_rule.id == gloss_fr.id + assert processed_rule.map_type == original_file_rule_map_type, "FileRule map_type should not change if save fails." + assert processed_rule.map_type == "GLOSS" + + # Check context.processed_maps_details: details should be unchanged + current_details_entry = updated_context.processed_maps_details[gloss_fr.id.hex] + assert current_details_entry['temp_processed_file'] == str(initial_gloss_temp_path) + assert current_details_entry['map_type'] == "GLOSS" + assert 'original_map_type_before_conversion' not in current_details_entry + assert 'notes' not in current_details_entry or "Converted from GLOSS" not in current_details_entry['notes'] + + mock_logging.error.assert_called_once_with( + f"Failed to save inverted GLOSS map {gloss_fr.id.hex} to {saved_path_arg}. Retaining original GLOSS map." + ) +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.logging') +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.save_image') +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.load_image') +def test_gloss_map_in_files_to_process_but_not_in_details(mock_load_image, mock_save_image, mock_logging): + """ + Test behavior when a GLOSS FileRule is in files_to_process but its details + are missing from processed_maps_details. + The stage should log an error and skip this FileRule. + """ + stage = GlossToRoughConversionStage() + + gloss_rule_id = uuid.uuid4() + # This FileRule is in files_to_process + gloss_fr_in_list = create_mock_file_rule_for_gloss_test(id_val=gloss_rule_id, map_type="GLOSS", filename_pattern="orphan_gloss.png") + + # processed_maps_details is empty or does not contain gloss_fr_in_list.id.hex + initial_details = {} + + context = create_gloss_conversion_mock_context( + initial_file_rules=[gloss_fr_in_list], + initial_processed_details=initial_details + ) + + original_files_to_process = list(context.files_to_process) + original_processed_maps_details = context.processed_maps_details.copy() + + updated_context = stage.execute(context) + + mock_load_image.assert_not_called() # Load should not be attempted if details are missing + mock_save_image.assert_not_called() # Save should not be attempted + + # Check context.files_to_process: rule should be unchanged + assert len(updated_context.files_to_process) == 1 + processed_rule = updated_context.files_to_process[0] + assert processed_rule.id == gloss_fr_in_list.id + assert processed_rule.map_type == "GLOSS", "FileRule map_type should not change if its details are missing." + + # Check context.processed_maps_details: should remain unchanged + assert updated_context.processed_maps_details == original_processed_maps_details, "processed_maps_details should not change." + + mock_logging.error.assert_called_once_with( + f"GLOSS map {gloss_fr_in_list.id.hex} found in files_to_process but missing from processed_maps_details. Skipping conversion." + ) + +# Test for Case 8.2 (GLOSS map ID in processed_maps_details but no corresponding FileRule in files_to_process) +# This case is implicitly handled because the stage iterates files_to_process. +# If a FileRule isn't in files_to_process, its corresponding entry in processed_maps_details (if any) won't be acted upon. +# We can add a simple test to ensure no errors occur and non-relevant details are untouched. + +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.logging') +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.save_image') +@mock.patch('processing.pipeline.stages.gloss_to_rough_conversion.ipu.load_image') +def test_gloss_detail_exists_but_not_in_files_to_process(mock_load_image, mock_save_image, mock_logging): + """ + Test that if a GLOSS map detail exists in processed_maps_details but + no corresponding FileRule is in files_to_process, it's simply ignored + without error, and other valid conversions proceed. + """ + stage = GlossToRoughConversionStage() + + # This rule will be processed + convert_rule_id = uuid.uuid4() + convert_fr = create_mock_file_rule_for_gloss_test(id_val=convert_rule_id, map_type="GLOSS", filename_pattern="convert_me.png") + convert_initial_temp_path = Path("/fake/temp_engine_dir/processed_convert_me.png") + + # This rule's details exist, but the rule itself is not in files_to_process + orphan_detail_id = uuid.uuid4() + + initial_details = { + convert_fr.id.hex: {'temp_processed_file': str(convert_initial_temp_path), 'status': 'Processed', 'map_type': 'GLOSS'}, + orphan_detail_id.hex: {'temp_processed_file': '/fake/temp_engine_dir/orphan.png', 'status': 'Processed', 'map_type': 'GLOSS', 'notes': 'This is an orphan'} + } + + context = create_gloss_conversion_mock_context( + initial_file_rules=[convert_fr], # Only convert_fr is in files_to_process + initial_processed_details=initial_details + ) + + mock_loaded_data = np.array([100], dtype=np.uint8) + mock_load_image.return_value = mock_loaded_data + mock_save_image.return_value = True + + updated_context = stage.execute(context) + + # Assert that load/save were called only for the rule in files_to_process + mock_load_image.assert_called_once_with(convert_initial_temp_path) + mock_save_image.assert_called_once() # Check it was called, details checked in other tests + + # Check that the orphan detail in processed_maps_details is untouched + assert orphan_detail_id.hex in updated_context.processed_maps_details + orphan_entry = updated_context.processed_maps_details[orphan_detail_id.hex] + assert orphan_entry['temp_processed_file'] == '/fake/temp_engine_dir/orphan.png' + assert orphan_entry['map_type'] == 'GLOSS' + assert orphan_entry['notes'] == 'This is an orphan' + assert 'original_map_type_before_conversion' not in orphan_entry + + # Check that the processed rule was indeed converted + assert convert_fr.id.hex in updated_context.processed_maps_details + converted_entry = updated_context.processed_maps_details[convert_fr.id.hex] + assert converted_entry['map_type'] == 'ROUGHNESS' + assert "rough_from_gloss_" in converted_entry['temp_processed_file'] + + # No errors should have been logged regarding the orphan detail + for call_args in mock_logging.error.call_args_list: + assert str(orphan_detail_id.hex) not in call_args[0][0], "Error logged for orphan detail" \ No newline at end of file diff --git a/tests/processing/pipeline/stages/test_individual_map_processing.py b/tests/processing/pipeline/stages/test_individual_map_processing.py new file mode 100644 index 0000000..2d78c00 --- /dev/null +++ b/tests/processing/pipeline/stages/test_individual_map_processing.py @@ -0,0 +1,555 @@ +import pytest +from unittest import mock +from pathlib import Path +import uuid +import numpy as np +from typing import Optional # Added for type hinting in helper functions + +from processing.pipeline.stages.individual_map_processing import IndividualMapProcessingStage +from processing.pipeline.asset_context import AssetProcessingContext +from rule_structure import AssetRule, SourceRule, FileRule, TransformSettings # Key models +from configuration import Configuration, GeneralSettings +# cv2 might be imported by the stage for interpolation constants, ensure it's mockable if so. +# For now, assume ipu handles interpolation details. + +def create_mock_transform_settings( + target_width=0, target_height=0, resize_mode="FIT", + ensure_pot=False, allow_upscale=True, target_color_profile="RGB" # Add other fields as needed +) -> mock.MagicMock: + ts = mock.MagicMock(spec=TransformSettings) + ts.target_width = target_width + ts.target_height = target_height + ts.resize_mode = resize_mode + ts.ensure_pot = ensure_pot + ts.allow_upscale = allow_upscale + ts.target_color_profile = target_color_profile + # ts.resize_filter = "AREA" # if your stage uses this + return ts + +def create_mock_file_rule_for_individual_processing( + id_val: Optional[uuid.UUID] = None, + map_type: str = "ALBEDO", + filename_pattern: str = "albedo_*.png", # Pattern for glob + item_type: str = "MAP_COL", + active: bool = True, + transform_settings: Optional[mock.MagicMock] = None +) -> mock.MagicMock: + mock_fr = mock.MagicMock(spec=FileRule) + mock_fr.id = id_val if id_val else uuid.uuid4() + mock_fr.map_type = map_type + mock_fr.filename_pattern = filename_pattern + mock_fr.item_type = item_type + mock_fr.active = active + mock_fr.transform_settings = transform_settings if transform_settings else create_mock_transform_settings() + return mock_fr + +def create_individual_map_proc_mock_context( + initial_file_rules: Optional[list] = None, + asset_source_path_str: str = "/fake/asset_source", + skip_asset_flag: bool = False, + asset_name: str = "IndividualMapAsset" +) -> AssetProcessingContext: + mock_asset_rule = mock.MagicMock(spec=AssetRule) + mock_asset_rule.name = asset_name + mock_asset_rule.source_path = Path(asset_source_path_str) + # file_rules on AssetRule not directly used by stage, context.files_to_process is + + mock_source_rule = mock.MagicMock(spec=SourceRule) + mock_config = mock.MagicMock(spec=Configuration) + # mock_config.general_settings = mock.MagicMock(spec=GeneralSettings) # If needed + + context = AssetProcessingContext( + source_rule=mock_source_rule, + asset_rule=mock_asset_rule, + workspace_path=Path("/fake/workspace"), + engine_temp_dir=Path("/fake/temp_engine_dir"), + output_base_path=Path("/fake/output"), + effective_supplier="ValidSupplier", + asset_metadata={'asset_name': asset_name}, + processed_maps_details={}, # Stage populates this + merged_maps_details={}, + files_to_process=list(initial_file_rules) if initial_file_rules else [], + loaded_data_cache={}, + config_obj=mock_config, + status_flags={'skip_asset': skip_asset_flag}, + incrementing_value=None, + sha5_value=None # Corrected from sha5_value to sha_value if that's the actual param + ) + return context + +# Placeholder for tests to be added next +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu') +@mock.patch('logging.info') +def test_asset_skipped_if_flag_is_true(mock_log_info, mock_ipu): + stage = IndividualMapProcessingStage() + context = create_individual_map_proc_mock_context(skip_asset_flag=True) + + # Add a dummy file rule to ensure it's not processed + file_rule = create_mock_file_rule_for_individual_processing() + context.files_to_process = [file_rule] + + updated_context = stage.execute(context) + + mock_ipu.load_image.assert_not_called() + mock_ipu.save_image.assert_not_called() + assert not updated_context.processed_maps_details # No details should be added + # Check for a log message indicating skip, if applicable (depends on stage's logging) + # mock_log_info.assert_any_call("Skipping asset IndividualMapAsset due to status_flags['skip_asset'] = True") # Example + + +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu') +@mock.patch('logging.info') +def test_no_processing_if_no_map_col_rules(mock_log_info, mock_ipu): + stage = IndividualMapProcessingStage() + + # Create a file rule that is NOT of item_type MAP_COL + non_map_col_rule = create_mock_file_rule_for_individual_processing(item_type="METADATA") + context = create_individual_map_proc_mock_context(initial_file_rules=[non_map_col_rule]) + + updated_context = stage.execute(context) + + mock_ipu.load_image.assert_not_called() + mock_ipu.save_image.assert_not_called() + assert not updated_context.processed_maps_details + # mock_log_info.assert_any_call("No FileRules of item_type 'MAP_COL' to process for asset IndividualMapAsset.") # Example + + +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.save_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.resize_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.calculate_target_dimensions') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.load_image') +@mock.patch('pathlib.Path.glob') # Mocking Path.glob used by the stage's _find_source_file +@mock.patch('logging.info') +@mock.patch('logging.error') +def test_individual_map_processing_success_no_resize( + mock_log_error, mock_log_info, mock_path_glob, mock_load_image, + mock_calc_dims, mock_resize_image, mock_save_image +): + stage = IndividualMapProcessingStage() + + source_file_name = "albedo_source.png" + # The glob is called on context.asset_rule.source_path, so mock that Path object's glob + mock_asset_source_path = Path("/fake/asset_source") + mock_found_source_path = mock_asset_source_path / source_file_name + + # We need to mock the glob method of the Path instance + # that represents the asset's source directory. + # The stage does something like: Path(context.asset_rule.source_path).glob(...) + # So, we need to ensure that when Path() is called with that specific string, + # the resulting object's glob method is our mock. + # A more robust way is to mock Path itself to return a mock object + # whose glob method is also a mock. + + # Simpler approach for now: assume Path.glob is used as a static/class method call + # or that the instance it's called on is correctly patched by @mock.patch('pathlib.Path.glob') + # if the stage does `from pathlib import Path` and then `Path(path_str).glob(...)`. + # The prompt example uses @mock.patch('pathlib.Path.glob'), implying the stage might do this: + # for f_pattern in patterns: + # for found_file in Path(base_dir).glob(f_pattern): ... + # Let's refine the mock_path_glob setup. + # The stage's _find_source_file likely does: + # search_path = Path(self.context.asset_rule.source_path) + # found_files = list(search_path.glob(filename_pattern)) + + # To correctly mock this, we need to mock the `glob` method of the specific Path instance. + # Or, if `_find_source_file` instantiates `Path` like `Path(str(context.asset_rule.source_path)).glob(...)`, + # then patching `pathlib.Path.glob` might work if it's treated as a method that gets bound. + # Let's stick to the example's @mock.patch('pathlib.Path.glob') and assume it covers the usage. + mock_path_glob.return_value = [mock_found_source_path] # Glob finds one file + + ts = create_mock_transform_settings(target_width=100, target_height=100) + file_rule = create_mock_file_rule_for_individual_processing( + map_type="ALBEDO", filename_pattern="albedo_*.png", transform_settings=ts + ) + context = create_individual_map_proc_mock_context( + initial_file_rules=[file_rule], + asset_source_path_str=str(mock_asset_source_path) # Ensure context uses this path + ) + + mock_img_data = np.zeros((100, 100, 3), dtype=np.uint8) # Original dimensions + mock_load_image.return_value = mock_img_data + mock_calc_dims.return_value = (100, 100) # No resize needed + mock_save_image.return_value = True + + updated_context = stage.execute(context) + + # Assert that Path(context.asset_rule.source_path).glob was called + # This requires a bit more intricate mocking if Path instances are created inside. + # For now, assert mock_path_glob was called with the pattern. + # The actual call in stage is `Path(context.asset_rule.source_path).glob(file_rule.filename_pattern)` + # So, `mock_path_glob` (if it patches `Path.glob` globally) should be called. + # We need to ensure the mock_path_glob is associated with the correct Path instance or that + # the global patch works as intended. + # A common pattern is: + # with mock.patch.object(Path, 'glob', return_value=[mock_found_source_path]) as specific_glob_mock: + # # execute code + # specific_glob_mock.assert_called_once_with(file_rule.filename_pattern) + # However, the decorator @mock.patch('pathlib.Path.glob') should work if the stage code is + # `from pathlib import Path; p = Path(...); p.glob(...)` + + # The stage's _find_source_file will instantiate a Path object from context.asset_rule.source_path + # and then call glob on it. + # So, @mock.patch('pathlib.Path.glob') is patching the method on the class. + # When an instance calls it, the mock is used. + mock_path_glob.assert_called_once_with(file_rule.filename_pattern) + + + mock_load_image.assert_called_once_with(mock_found_source_path) + # The actual call to calculate_target_dimensions is: + # ipu.calculate_target_dimensions(original_dims, ts.target_width, ts.target_height, ts.resize_mode, ts.ensure_pot, ts.allow_upscale) + mock_calc_dims.assert_called_once_with( + (100, 100), ts.target_width, ts.target_height, ts.resize_mode, ts.ensure_pot, ts.allow_upscale + ) + mock_resize_image.assert_not_called() # Crucial for this test case + mock_save_image.assert_called_once() + + # Check save path and data + saved_image_arg, saved_path_arg = mock_save_image.call_args[0] + assert np.array_equal(saved_image_arg, mock_img_data) # Ensure correct image data is passed to save + assert "processed_ALBEDO_" in saved_path_arg.name # Based on map_type + assert file_rule.id.hex in saved_path_arg.name # Ensure unique name with FileRule ID + assert saved_path_arg.parent == context.engine_temp_dir + + assert file_rule.id.hex in updated_context.processed_maps_details + details = updated_context.processed_maps_details[file_rule.id.hex] + assert details['status'] == 'Processed' + assert details['source_file'] == str(mock_found_source_path) + assert Path(details['temp_processed_file']) == saved_path_arg + assert details['original_dimensions'] == (100, 100) + assert details['processed_dimensions'] == (100, 100) + assert details['map_type'] == file_rule.map_type + mock_log_error.assert_not_called() + mock_log_info.assert_any_call(f"Successfully processed map {file_rule.map_type} (ID: {file_rule.id.hex}) for asset {context.asset_rule.name}. Output: {saved_path_arg}") + + +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.save_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.resize_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.calculate_target_dimensions') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.load_image') +@mock.patch('pathlib.Path.glob') +@mock.patch('logging.info') +@mock.patch('logging.error') +def test_source_file_not_found( + mock_log_error, mock_log_info, mock_path_glob, mock_load_image, + mock_calc_dims, mock_resize_image, mock_save_image +): + stage = IndividualMapProcessingStage() + mock_asset_source_path = Path("/fake/asset_source") + + mock_path_glob.return_value = [] # Glob finds no files + + file_rule = create_mock_file_rule_for_individual_processing(filename_pattern="nonexistent_*.png") + context = create_individual_map_proc_mock_context( + initial_file_rules=[file_rule], + asset_source_path_str=str(mock_asset_source_path) + ) + + updated_context = stage.execute(context) + + mock_path_glob.assert_called_once_with(file_rule.filename_pattern) + mock_load_image.assert_not_called() + mock_calc_dims.assert_not_called() + mock_resize_image.assert_not_called() + mock_save_image.assert_not_called() + + assert file_rule.id.hex in updated_context.processed_maps_details + details = updated_context.processed_maps_details[file_rule.id.hex] + assert details['status'] == 'Source Not Found' + assert details['source_file'] is None + assert details['temp_processed_file'] is None + assert details['error_message'] is not None # Check an error message is present + mock_log_error.assert_called_once() + # Example: mock_log_error.assert_called_with(f"Could not find source file for rule {file_rule.id} (pattern: {file_rule.filename_pattern}) in {context.asset_rule.source_path}") + + +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.save_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.resize_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.calculate_target_dimensions') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.load_image') +@mock.patch('pathlib.Path.glob') +@mock.patch('logging.info') +@mock.patch('logging.error') +def test_load_image_fails( + mock_log_error, mock_log_info, mock_path_glob, mock_load_image, + mock_calc_dims, mock_resize_image, mock_save_image +): + stage = IndividualMapProcessingStage() + source_file_name = "albedo_corrupt.png" + mock_asset_source_path = Path("/fake/asset_source") + mock_found_source_path = mock_asset_source_path / source_file_name + mock_path_glob.return_value = [mock_found_source_path] + + mock_load_image.return_value = None # Simulate load failure + + file_rule = create_mock_file_rule_for_individual_processing(filename_pattern="albedo_*.png") + context = create_individual_map_proc_mock_context( + initial_file_rules=[file_rule], + asset_source_path_str=str(mock_asset_source_path) + ) + + updated_context = stage.execute(context) + + mock_path_glob.assert_called_once_with(file_rule.filename_pattern) + mock_load_image.assert_called_once_with(mock_found_source_path) + mock_calc_dims.assert_not_called() + mock_resize_image.assert_not_called() + mock_save_image.assert_not_called() + + assert file_rule.id.hex in updated_context.processed_maps_details + details = updated_context.processed_maps_details[file_rule.id.hex] + assert details['status'] == 'Load Failed' + assert details['source_file'] == str(mock_found_source_path) + assert details['temp_processed_file'] is None + assert details['error_message'] is not None + mock_log_error.assert_called_once() + # Example: mock_log_error.assert_called_with(f"Failed to load image {mock_found_source_path} for rule {file_rule.id}") + + +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.save_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.resize_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.calculate_target_dimensions') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.load_image') +@mock.patch('pathlib.Path.glob') +@mock.patch('logging.info') +@mock.patch('logging.error') +def test_resize_occurs_when_dimensions_differ( + mock_log_error, mock_log_info, mock_path_glob, mock_load_image, + mock_calc_dims, mock_resize_image, mock_save_image +): + stage = IndividualMapProcessingStage() + source_file_name = "albedo_resize.png" + mock_asset_source_path = Path("/fake/asset_source") + mock_found_source_path = mock_asset_source_path / source_file_name + mock_path_glob.return_value = [mock_found_source_path] + + original_dims = (100, 100) + target_dims = (50, 50) # Different dimensions + mock_img_data = np.zeros((*original_dims, 3), dtype=np.uint8) + mock_resized_img_data = np.zeros((*target_dims, 3), dtype=np.uint8) + + mock_load_image.return_value = mock_img_data + ts = create_mock_transform_settings(target_width=target_dims[0], target_height=target_dims[1]) + file_rule = create_mock_file_rule_for_individual_processing(transform_settings=ts) + context = create_individual_map_proc_mock_context( + initial_file_rules=[file_rule], + asset_source_path_str=str(mock_asset_source_path) + ) + + mock_calc_dims.return_value = target_dims # Simulate calc_dims returning new dimensions + mock_resize_image.return_value = mock_resized_img_data # Simulate resize returning new image data + mock_save_image.return_value = True + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(mock_found_source_path) + mock_calc_dims.assert_called_once_with( + original_dims, ts.target_width, ts.target_height, ts.resize_mode, ts.ensure_pot, ts.allow_upscale + ) + # The actual call to resize_image is: + # ipu.resize_image(loaded_image, target_dims, ts.resize_filter) # Assuming resize_filter is used + # If resize_filter is not on TransformSettings or not used, adjust this. + # For now, let's assume it's ipu.resize_image(loaded_image, target_dims) or similar + # The stage code is: resized_image = ipu.resize_image(loaded_image, target_dims_calculated, file_rule.transform_settings.resize_filter) + # So we need to mock ts.resize_filter + ts.resize_filter = "LANCZOS4" # Example filter + mock_resize_image.assert_called_once_with(mock_img_data, target_dims, ts.resize_filter) + + saved_image_arg, saved_path_arg = mock_save_image.call_args[0] + assert np.array_equal(saved_image_arg, mock_resized_img_data) # Check resized data is saved + assert "processed_ALBEDO_" in saved_path_arg.name + assert saved_path_arg.parent == context.engine_temp_dir + + assert file_rule.id.hex in updated_context.processed_maps_details + details = updated_context.processed_maps_details[file_rule.id.hex] + assert details['status'] == 'Processed' + assert details['original_dimensions'] == original_dims + assert details['processed_dimensions'] == target_dims + mock_log_error.assert_not_called() + + +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.save_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.resize_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.calculate_target_dimensions') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.load_image') +@mock.patch('pathlib.Path.glob') +@mock.patch('logging.info') +@mock.patch('logging.error') +def test_save_image_fails( + mock_log_error, mock_log_info, mock_path_glob, mock_load_image, + mock_calc_dims, mock_resize_image, mock_save_image +): + stage = IndividualMapProcessingStage() + source_file_name = "albedo_save_fail.png" + mock_asset_source_path = Path("/fake/asset_source") + mock_found_source_path = mock_asset_source_path / source_file_name + mock_path_glob.return_value = [mock_found_source_path] + + mock_img_data = np.zeros((100, 100, 3), dtype=np.uint8) + mock_load_image.return_value = mock_img_data + mock_calc_dims.return_value = (100, 100) # No resize + mock_save_image.return_value = False # Simulate save failure + + ts = create_mock_transform_settings() + file_rule = create_mock_file_rule_for_individual_processing(transform_settings=ts) + context = create_individual_map_proc_mock_context( + initial_file_rules=[file_rule], + asset_source_path_str=str(mock_asset_source_path) + ) + + updated_context = stage.execute(context) + + mock_save_image.assert_called_once() # Attempt to save should still happen + + assert file_rule.id.hex in updated_context.processed_maps_details + details = updated_context.processed_maps_details[file_rule.id.hex] + assert details['status'] == 'Save Failed' + assert details['source_file'] == str(mock_found_source_path) + assert details['temp_processed_file'] is not None # Path was generated + assert details['error_message'] is not None + mock_log_error.assert_called_once() + # Example: mock_log_error.assert_called_with(f"Failed to save processed image for rule {file_rule.id} to {details['temp_processed_file']}") + + +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.save_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.resize_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.calculate_target_dimensions') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.load_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.convert_bgr_to_rgb') +@mock.patch('pathlib.Path.glob') +@mock.patch('logging.info') +@mock.patch('logging.error') +def test_color_conversion_bgr_to_rgb( + mock_log_error, mock_log_info, mock_path_glob, mock_convert_bgr, mock_load_image, + mock_calc_dims, mock_resize_image, mock_save_image +): + stage = IndividualMapProcessingStage() + source_file_name = "albedo_bgr.png" + mock_asset_source_path = Path("/fake/asset_source") + mock_found_source_path = mock_asset_source_path / source_file_name + mock_path_glob.return_value = [mock_found_source_path] + + mock_bgr_img_data = np.zeros((100, 100, 3), dtype=np.uint8) # Loaded as BGR + mock_rgb_img_data = np.zeros((100, 100, 3), dtype=np.uint8) # After conversion + + mock_load_image.return_value = mock_bgr_img_data # Image is loaded (assume BGR by default from cv2) + mock_convert_bgr.return_value = mock_rgb_img_data # Mock the conversion + mock_calc_dims.return_value = (100, 100) # No resize + mock_save_image.return_value = True + + # Transform settings request RGB, and stage assumes load might be BGR + ts = create_mock_transform_settings(target_color_profile="RGB") + file_rule = create_mock_file_rule_for_individual_processing(transform_settings=ts) + context = create_individual_map_proc_mock_context( + initial_file_rules=[file_rule], + asset_source_path_str=str(mock_asset_source_path) + ) + # The stage code is: + # if file_rule.transform_settings.target_color_profile == "RGB" and loaded_image.shape[2] == 3: + # logger.info(f"Attempting to convert image from BGR to RGB for {file_rule_id_hex}") + # processed_image_data = ipu.convert_bgr_to_rgb(processed_image_data) + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(mock_found_source_path) + mock_convert_bgr.assert_called_once_with(mock_bgr_img_data) + mock_resize_image.assert_not_called() + + saved_image_arg, _ = mock_save_image.call_args[0] + assert np.array_equal(saved_image_arg, mock_rgb_img_data) # Ensure RGB data is saved + mock_log_error.assert_not_called() + mock_log_info.assert_any_call(f"Attempting to convert image from BGR to RGB for {file_rule.id.hex}") + + +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.save_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.resize_image') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.calculate_target_dimensions') +@mock.patch('processing.pipeline.stages.individual_map_processing.ipu.load_image') +@mock.patch('pathlib.Path.glob') +@mock.patch('logging.info') +@mock.patch('logging.error') +def test_multiple_map_col_rules_processed( + mock_log_error, mock_log_info, mock_path_glob, mock_load_image, + mock_calc_dims, mock_resize_image, mock_save_image +): + stage = IndividualMapProcessingStage() + mock_asset_source_path = Path("/fake/asset_source") + + # Rule 1: Albedo + ts1 = create_mock_transform_settings(target_width=100, target_height=100) + file_rule1_id = uuid.uuid4() + file_rule1 = create_mock_file_rule_for_individual_processing( + id_val=file_rule1_id, map_type="ALBEDO", filename_pattern="albedo_*.png", transform_settings=ts1 + ) + source_file1 = mock_asset_source_path / "albedo_map.png" + img_data1 = np.zeros((100, 100, 3), dtype=np.uint8) + + # Rule 2: Roughness + ts2 = create_mock_transform_settings(target_width=50, target_height=50) # Resize + ts2.resize_filter = "AREA" + file_rule2_id = uuid.uuid4() + file_rule2 = create_mock_file_rule_for_individual_processing( + id_val=file_rule2_id, map_type="ROUGHNESS", filename_pattern="rough_*.png", transform_settings=ts2 + ) + source_file2 = mock_asset_source_path / "rough_map.png" + img_data2_orig = np.zeros((200, 200, 1), dtype=np.uint8) # Original, needs resize + img_data2_resized = np.zeros((50, 50, 1), dtype=np.uint8) # Resized + + context = create_individual_map_proc_mock_context( + initial_file_rules=[file_rule1, file_rule2], + asset_source_path_str=str(mock_asset_source_path) + ) + + # Mock behaviors for Path.glob, load_image, calc_dims, resize, save + # Path.glob will be called twice + mock_path_glob.side_effect = [ + [source_file1], # For albedo_*.png + [source_file2] # For rough_*.png + ] + mock_load_image.side_effect = [img_data1, img_data2_orig] + mock_calc_dims.side_effect = [ + (100, 100), # For rule1 (no change) + (50, 50) # For rule2 (change) + ] + mock_resize_image.return_value = img_data2_resized # Only called for rule2 + mock_save_image.return_value = True + + updated_context = stage.execute(context) + + # Assertions for Rule 1 (Albedo) + assert mock_path_glob.call_args_list[0][0][0] == file_rule1.filename_pattern + assert mock_load_image.call_args_list[0][0][0] == source_file1 + assert mock_calc_dims.call_args_list[0][0] == ((100,100), ts1.target_width, ts1.target_height, ts1.resize_mode, ts1.ensure_pot, ts1.allow_upscale) + + # Assertions for Rule 2 (Roughness) + assert mock_path_glob.call_args_list[1][0][0] == file_rule2.filename_pattern + assert mock_load_image.call_args_list[1][0][0] == source_file2 + assert mock_calc_dims.call_args_list[1][0] == ((200,200), ts2.target_width, ts2.target_height, ts2.resize_mode, ts2.ensure_pot, ts2.allow_upscale) + mock_resize_image.assert_called_once_with(img_data2_orig, (50,50), ts2.resize_filter) + + assert mock_save_image.call_count == 2 + # Check saved image for rule 1 + saved_img1_arg, saved_path1_arg = mock_save_image.call_args_list[0][0] + assert np.array_equal(saved_img1_arg, img_data1) + assert "processed_ALBEDO_" in saved_path1_arg.name + assert file_rule1_id.hex in saved_path1_arg.name + + # Check saved image for rule 2 + saved_img2_arg, saved_path2_arg = mock_save_image.call_args_list[1][0] + assert np.array_equal(saved_img2_arg, img_data2_resized) + assert "processed_ROUGHNESS_" in saved_path2_arg.name + assert file_rule2_id.hex in saved_path2_arg.name + + # Check context details + assert file_rule1_id.hex in updated_context.processed_maps_details + details1 = updated_context.processed_maps_details[file_rule1_id.hex] + assert details1['status'] == 'Processed' + assert details1['original_dimensions'] == (100, 100) + assert details1['processed_dimensions'] == (100, 100) + + assert file_rule2_id.hex in updated_context.processed_maps_details + details2 = updated_context.processed_maps_details[file_rule2_id.hex] + assert details2['status'] == 'Processed' + assert details2['original_dimensions'] == (200, 200) # Original dims of img_data2_orig + assert details2['processed_dimensions'] == (50, 50) + + mock_log_error.assert_not_called() \ No newline at end of file diff --git a/tests/processing/pipeline/stages/test_map_merging.py b/tests/processing/pipeline/stages/test_map_merging.py new file mode 100644 index 0000000..d222f62 --- /dev/null +++ b/tests/processing/pipeline/stages/test_map_merging.py @@ -0,0 +1,538 @@ +import pytest +from unittest import mock +from pathlib import Path +import uuid +import numpy as np +from typing import Optional # Added Optional for type hinting + +from processing.pipeline.stages.map_merging import MapMergingStage +from processing.pipeline.asset_context import AssetProcessingContext +from rule_structure import AssetRule, SourceRule, FileRule, MergeSettings, MergeInputChannel +from configuration import Configuration + +# Mock Helper Functions +def create_mock_merge_input_channel( + file_rule_id: uuid.UUID, source_channel: int = 0, target_channel: int = 0, invert: bool = False +) -> mock.MagicMock: + mic = mock.MagicMock(spec=MergeInputChannel) + mic.file_rule_id = file_rule_id + mic.source_channel = source_channel + mic.target_channel = target_channel + mic.invert_source_channel = invert + mic.default_value_if_missing = 0 # Or some other default + return mic + +def create_mock_merge_settings( + input_maps: Optional[list] = None, # List of mock MergeInputChannel + output_channels: int = 3 +) -> mock.MagicMock: + ms = mock.MagicMock(spec=MergeSettings) + ms.input_maps = input_maps if input_maps is not None else [] + ms.output_channels = output_channels + return ms + +def create_mock_file_rule_for_merging( + id_val: Optional[uuid.UUID] = None, + map_type: str = "ORM", # Output map type + item_type: str = "MAP_MERGE", + merge_settings: Optional[mock.MagicMock] = None +) -> mock.MagicMock: + mock_fr = mock.MagicMock(spec=FileRule) + mock_fr.id = id_val if id_val else uuid.uuid4() + mock_fr.map_type = map_type + mock_fr.filename_pattern = f"{map_type.lower()}_merged.png" # Placeholder + mock_fr.item_type = item_type + mock_fr.active = True + mock_fr.merge_settings = merge_settings if merge_settings else create_mock_merge_settings() + return mock_fr + +def create_map_merging_mock_context( + initial_file_rules: Optional[list] = None, # Will contain the MAP_MERGE rule + initial_processed_details: Optional[dict] = None, # Pre-processed inputs for merge + skip_asset_flag: bool = False, + asset_name: str = "MergeAsset" +) -> AssetProcessingContext: + mock_asset_rule = mock.MagicMock(spec=AssetRule) + mock_asset_rule.name = asset_name + mock_source_rule = mock.MagicMock(spec=SourceRule) + mock_config = mock.MagicMock(spec=Configuration) + + context = AssetProcessingContext( + source_rule=mock_source_rule, + asset_rule=mock_asset_rule, + workspace_path=Path("/fake/workspace"), + engine_temp_dir=Path("/fake/temp_engine_dir"), + output_base_path=Path("/fake/output"), + effective_supplier="ValidSupplier", + asset_metadata={'asset_name': asset_name}, + processed_maps_details=initial_processed_details if initial_processed_details is not None else {}, + merged_maps_details={}, # Stage populates this + files_to_process=list(initial_file_rules) if initial_file_rules else [], + loaded_data_cache={}, + config_obj=mock_config, + status_flags={'skip_asset': skip_asset_flag}, + incrementing_value=None, + sha5_value=None # Corrected from sha5_value to sha_value based on AssetProcessingContext + ) + return context +def test_asset_skipped(): + stage = MapMergingStage() + context = create_map_merging_mock_context(skip_asset_flag=True) + + updated_context = stage.execute(context) + + assert updated_context == context # No changes expected + assert not updated_context.merged_maps_details # No maps should be merged + +def test_no_map_merge_rules(): + stage = MapMergingStage() + # Context with a non-MAP_MERGE rule + non_merge_rule = create_mock_file_rule_for_merging(item_type="TEXTURE_MAP", map_type="Diffuse") + context = create_map_merging_mock_context(initial_file_rules=[non_merge_rule]) + + updated_context = stage.execute(context) + + assert updated_context == context # No changes expected + assert not updated_context.merged_maps_details # No maps should be merged + +@mock.patch('processing.pipeline.stages.map_merging.ipu.save_image') +@mock.patch('processing.pipeline.stages.map_merging.ipu.resize_image') # If testing resize +@mock.patch('processing.pipeline.stages.map_merging.ipu.load_image') +@mock.patch('logging.info') +@mock.patch('logging.error') +def test_map_merging_rgb_success(mock_log_error, mock_log_info, mock_load_image, mock_resize_image, mock_save_image): + stage = MapMergingStage() + + # Input FileRules (mocked as already processed) + r_id, g_id, b_id = uuid.uuid4(), uuid.uuid4(), uuid.uuid4() + processed_details = { + r_id.hex: {'temp_processed_file': '/fake/red.png', 'status': 'Processed', 'map_type': 'RED_SRC'}, + g_id.hex: {'temp_processed_file': '/fake/green.png', 'status': 'Processed', 'map_type': 'GREEN_SRC'}, + b_id.hex: {'temp_processed_file': '/fake/blue.png', 'status': 'Processed', 'map_type': 'BLUE_SRC'} + } + # Mock loaded image data (grayscale for inputs) + mock_r_data = np.full((10, 10), 200, dtype=np.uint8) + mock_g_data = np.full((10, 10), 100, dtype=np.uint8) + mock_b_data = np.full((10, 10), 50, dtype=np.uint8) + mock_load_image.side_effect = [mock_r_data, mock_g_data, mock_b_data] + + # Merge Rule setup + merge_inputs = [ + create_mock_merge_input_channel(file_rule_id=r_id, source_channel=0, target_channel=0), # R to R + create_mock_merge_input_channel(file_rule_id=g_id, source_channel=0, target_channel=1), # G to G + create_mock_merge_input_channel(file_rule_id=b_id, source_channel=0, target_channel=2) # B to B + ] + merge_settings = create_mock_merge_settings(input_maps=merge_inputs, output_channels=3) + merge_rule_id = uuid.uuid4() + merge_rule = create_mock_file_rule_for_merging(id_val=merge_rule_id, map_type="RGB_Combined", merge_settings=merge_settings) + + context = create_map_merging_mock_context( + initial_file_rules=[merge_rule], + initial_processed_details=processed_details + ) + mock_save_image.return_value = True + + updated_context = stage.execute(context) + + assert mock_load_image.call_count == 3 + mock_resize_image.assert_not_called() # Assuming all inputs are same size for this test + mock_save_image.assert_called_once() + + # Check that the correct filename was passed to save_image + # The filename is constructed as: f"{context.asset_rule.name}_merged_{merge_rule.map_type}{Path(first_input_path).suffix}" + # In this case, first_input_path is '/fake/red.png', so suffix is '.png' + # Asset name is "MergeAsset" + expected_filename_part = f"{context.asset_rule.name}_merged_{merge_rule.map_type}.png" + saved_path_arg = mock_save_image.call_args[0][0] + assert expected_filename_part in str(saved_path_arg) + + + saved_data = mock_save_image.call_args[0][1] + assert saved_data.shape == (10, 10, 3) + assert np.all(saved_data[:,:,0] == 200) # Red channel + assert np.all(saved_data[:,:,1] == 100) # Green channel + assert np.all(saved_data[:,:,2] == 50) # Blue channel + + assert merge_rule.id.hex in updated_context.merged_maps_details + details = updated_context.merged_maps_details[merge_rule.id.hex] + assert details['status'] == 'Processed' + # The temp_merged_file path will be under engine_temp_dir / asset_name / filename + assert f"{context.engine_temp_dir / context.asset_rule.name / expected_filename_part}" == details['temp_merged_file'] + mock_log_error.assert_not_called() + mock_log_info.assert_any_call(f"Successfully merged map '{merge_rule.map_type}' for asset '{context.asset_rule.name}'.") + +# Unit tests will be added below this line +@mock.patch('processing.pipeline.stages.map_merging.ipu.save_image') +@mock.patch('processing.pipeline.stages.map_merging.ipu.resize_image') +@mock.patch('processing.pipeline.stages.map_merging.ipu.load_image') +@mock.patch('logging.info') +@mock.patch('logging.error') +def test_map_merging_channel_inversion(mock_log_error, mock_log_info, mock_load_image, mock_resize_image, mock_save_image): + stage = MapMergingStage() + + # Input FileRule + input_id = uuid.uuid4() + processed_details = { + input_id.hex: {'temp_processed_file': '/fake/source.png', 'status': 'Processed', 'map_type': 'SOURCE_MAP'} + } + # Mock loaded image data (single channel for simplicity, to be inverted) + mock_source_data = np.array([[0, 100], [155, 255]], dtype=np.uint8) + mock_load_image.return_value = mock_source_data + + # Merge Rule setup: one input, inverted, to one output channel + merge_inputs = [ + create_mock_merge_input_channel(file_rule_id=input_id, source_channel=0, target_channel=0, invert=True) + ] + merge_settings = create_mock_merge_settings(input_maps=merge_inputs, output_channels=1) + merge_rule_id = uuid.uuid4() + merge_rule = create_mock_file_rule_for_merging(id_val=merge_rule_id, map_type="Inverted_Gray", merge_settings=merge_settings) + + context = create_map_merging_mock_context( + initial_file_rules=[merge_rule], + initial_processed_details=processed_details + ) + mock_save_image.return_value = True + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(Path('/fake/source.png')) + mock_resize_image.assert_not_called() + mock_save_image.assert_called_once() + + saved_data = mock_save_image.call_args[0][1] + assert saved_data.shape == (2, 2) # Grayscale output + + # Expected inverted data: 255-original + expected_inverted_data = np.array([[255, 155], [100, 0]], dtype=np.uint8) + assert np.all(saved_data == expected_inverted_data) + + assert merge_rule.id.hex in updated_context.merged_maps_details + details = updated_context.merged_maps_details[merge_rule.id.hex] + assert details['status'] == 'Processed' + assert "merged_Inverted_Gray" in details['temp_merged_file'] + mock_log_error.assert_not_called() + mock_log_info.assert_any_call(f"Successfully merged map '{merge_rule.map_type}' for asset '{context.asset_rule.name}'.") +@mock.patch('processing.pipeline.stages.map_merging.ipu.save_image') +@mock.patch('processing.pipeline.stages.map_merging.ipu.load_image') +@mock.patch('logging.error') +def test_map_merging_input_map_missing(mock_log_error, mock_load_image, mock_save_image): + stage = MapMergingStage() + + # Input FileRule ID that will be missing from processed_details + missing_input_id = uuid.uuid4() + + # Merge Rule setup + merge_inputs = [ + create_mock_merge_input_channel(file_rule_id=missing_input_id, source_channel=0, target_channel=0) + ] + merge_settings = create_mock_merge_settings(input_maps=merge_inputs, output_channels=1) + merge_rule_id = uuid.uuid4() + merge_rule = create_mock_file_rule_for_merging(id_val=merge_rule_id, map_type="TestMissing", merge_settings=merge_settings) + + # processed_details is empty, so missing_input_id will not be found + context = create_map_merging_mock_context( + initial_file_rules=[merge_rule], + initial_processed_details={} + ) + + updated_context = stage.execute(context) + + mock_load_image.assert_not_called() + mock_save_image.assert_not_called() + + assert merge_rule.id.hex in updated_context.merged_maps_details + details = updated_context.merged_maps_details[merge_rule.id.hex] + assert details['status'] == 'Failed' + assert 'error_message' in details + assert f"Input map FileRule ID {missing_input_id.hex} not found in processed_maps_details or not successfully processed" in details['error_message'] + + mock_log_error.assert_called_once() + assert f"Failed to merge map '{merge_rule.map_type}' for asset '{context.asset_rule.name}'" in mock_log_error.call_args[0][0] + assert f"Input map FileRule ID {missing_input_id.hex} not found in processed_maps_details or not successfully processed" in mock_log_error.call_args[0][0] + +@mock.patch('processing.pipeline.stages.map_merging.ipu.save_image') +@mock.patch('processing.pipeline.stages.map_merging.ipu.load_image') +@mock.patch('logging.error') +def test_map_merging_input_map_status_not_processed(mock_log_error, mock_load_image, mock_save_image): + stage = MapMergingStage() + + input_id = uuid.uuid4() + processed_details = { + # Status is 'Failed', not 'Processed' + input_id.hex: {'temp_processed_file': '/fake/source.png', 'status': 'Failed', 'map_type': 'SOURCE_MAP'} + } + + merge_inputs = [ + create_mock_merge_input_channel(file_rule_id=input_id, source_channel=0, target_channel=0) + ] + merge_settings = create_mock_merge_settings(input_maps=merge_inputs, output_channels=1) + merge_rule_id = uuid.uuid4() + merge_rule = create_mock_file_rule_for_merging(id_val=merge_rule_id, map_type="TestNotProcessed", merge_settings=merge_settings) + + context = create_map_merging_mock_context( + initial_file_rules=[merge_rule], + initial_processed_details=processed_details + ) + + updated_context = stage.execute(context) + + mock_load_image.assert_not_called() + mock_save_image.assert_not_called() + + assert merge_rule.id.hex in updated_context.merged_maps_details + details = updated_context.merged_maps_details[merge_rule.id.hex] + assert details['status'] == 'Failed' + assert 'error_message' in details + assert f"Input map FileRule ID {input_id.hex} not found in processed_maps_details or not successfully processed" in details['error_message'] + + mock_log_error.assert_called_once() + assert f"Failed to merge map '{merge_rule.map_type}' for asset '{context.asset_rule.name}'" in mock_log_error.call_args[0][0] + assert f"Input map FileRule ID {input_id.hex} not found in processed_maps_details or not successfully processed" in mock_log_error.call_args[0][0] +@mock.patch('processing.pipeline.stages.map_merging.ipu.save_image') +@mock.patch('processing.pipeline.stages.map_merging.ipu.load_image') +@mock.patch('logging.error') +def test_map_merging_load_image_fails(mock_log_error, mock_load_image, mock_save_image): + stage = MapMergingStage() + + input_id = uuid.uuid4() + processed_details = { + input_id.hex: {'temp_processed_file': '/fake/source.png', 'status': 'Processed', 'map_type': 'SOURCE_MAP'} + } + + # Configure mock_load_image to raise an exception + mock_load_image.side_effect = Exception("Failed to load image") + + merge_inputs = [ + create_mock_merge_input_channel(file_rule_id=input_id, source_channel=0, target_channel=0) + ] + merge_settings = create_mock_merge_settings(input_maps=merge_inputs, output_channels=1) + merge_rule_id = uuid.uuid4() + merge_rule = create_mock_file_rule_for_merging(id_val=merge_rule_id, map_type="TestLoadFail", merge_settings=merge_settings) + + context = create_map_merging_mock_context( + initial_file_rules=[merge_rule], + initial_processed_details=processed_details + ) + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(Path('/fake/source.png')) + mock_save_image.assert_not_called() + + assert merge_rule.id.hex in updated_context.merged_maps_details + details = updated_context.merged_maps_details[merge_rule.id.hex] + assert details['status'] == 'Failed' + assert 'error_message' in details + assert "Failed to load image for merge input" in details['error_message'] + assert str(Path('/fake/source.png')) in details['error_message'] + + mock_log_error.assert_called_once() + assert f"Failed to merge map '{merge_rule.map_type}' for asset '{context.asset_rule.name}'" in mock_log_error.call_args[0][0] + assert "Failed to load image for merge input" in mock_log_error.call_args[0][0] +@mock.patch('processing.pipeline.stages.map_merging.ipu.save_image') +@mock.patch('processing.pipeline.stages.map_merging.ipu.load_image') +@mock.patch('logging.error') +def test_map_merging_save_image_fails(mock_log_error, mock_load_image, mock_save_image): + stage = MapMergingStage() + + input_id = uuid.uuid4() + processed_details = { + input_id.hex: {'temp_processed_file': '/fake/source.png', 'status': 'Processed', 'map_type': 'SOURCE_MAP'} + } + mock_source_data = np.full((10, 10), 128, dtype=np.uint8) + mock_load_image.return_value = mock_source_data + + # Configure mock_save_image to return False (indicating failure) + mock_save_image.return_value = False + + merge_inputs = [ + create_mock_merge_input_channel(file_rule_id=input_id, source_channel=0, target_channel=0) + ] + merge_settings = create_mock_merge_settings(input_maps=merge_inputs, output_channels=1) + merge_rule_id = uuid.uuid4() + merge_rule = create_mock_file_rule_for_merging(id_val=merge_rule_id, map_type="TestSaveFail", merge_settings=merge_settings) + + context = create_map_merging_mock_context( + initial_file_rules=[merge_rule], + initial_processed_details=processed_details + ) + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(Path('/fake/source.png')) + mock_save_image.assert_called_once() # save_image is called, but returns False + + assert merge_rule.id.hex in updated_context.merged_maps_details + details = updated_context.merged_maps_details[merge_rule.id.hex] + assert details['status'] == 'Failed' + assert 'error_message' in details + assert "Failed to save merged map" in details['error_message'] + + mock_log_error.assert_called_once() + assert f"Failed to merge map '{merge_rule.map_type}' for asset '{context.asset_rule.name}'" in mock_log_error.call_args[0][0] + assert "Failed to save merged map" in mock_log_error.call_args[0][0] +@mock.patch('processing.pipeline.stages.map_merging.ipu.save_image') +@mock.patch('processing.pipeline.stages.map_merging.ipu.resize_image') +@mock.patch('processing.pipeline.stages.map_merging.ipu.load_image') +@mock.patch('logging.info') +@mock.patch('logging.error') +def test_map_merging_dimension_mismatch_handling(mock_log_error, mock_log_info, mock_load_image, mock_resize_image, mock_save_image): + stage = MapMergingStage() + + # Input FileRules + id1, id2 = uuid.uuid4(), uuid.uuid4() + processed_details = { + id1.hex: {'temp_processed_file': '/fake/img1.png', 'status': 'Processed', 'map_type': 'IMG1_SRC'}, + id2.hex: {'temp_processed_file': '/fake/img2.png', 'status': 'Processed', 'map_type': 'IMG2_SRC'} + } + + # Mock loaded image data with different dimensions + mock_img1_data = np.full((10, 10), 100, dtype=np.uint8) # 10x10 + mock_img2_data_original = np.full((5, 5), 200, dtype=np.uint8) # 5x5, will be resized + + mock_load_image.side_effect = [mock_img1_data, mock_img2_data_original] + + # Mock resize_image to return an image of the target dimensions + # For simplicity, it just creates a new array of the target size filled with a value. + mock_img2_data_resized = np.full((10, 10), 210, dtype=np.uint8) # Resized to 10x10 + mock_resize_image.return_value = mock_img2_data_resized + + # Merge Rule setup: two inputs, one output channel (e.g., averaging them) + # Target channel 0 for both, the stage should handle combining them if they map to the same target. + # However, the current stage logic for multiple inputs to the same target channel is to take the last one. + # Let's make them target different channels for a clearer test of resize. + merge_inputs = [ + create_mock_merge_input_channel(file_rule_id=id1, source_channel=0, target_channel=0), + create_mock_merge_input_channel(file_rule_id=id2, source_channel=0, target_channel=1) + ] + merge_settings = create_mock_merge_settings(input_maps=merge_inputs, output_channels=2) # Outputting 2 channels + merge_rule_id = uuid.uuid4() + merge_rule = create_mock_file_rule_for_merging(id_val=merge_rule_id, map_type="ResizedMerge", merge_settings=merge_settings) + + context = create_map_merging_mock_context( + initial_file_rules=[merge_rule], + initial_processed_details=processed_details + ) + mock_save_image.return_value = True + + updated_context = stage.execute(context) + + assert mock_load_image.call_count == 2 + mock_load_image.assert_any_call(Path('/fake/img1.png')) + mock_load_image.assert_any_call(Path('/fake/img2.png')) + + # Assert resize_image was called for the second image to match the first's dimensions + mock_resize_image.assert_called_once() + # The first argument to resize_image is the image data, second is target_shape tuple (height, width) + # np.array_equal is needed for comparing numpy arrays in mock calls + assert np.array_equal(mock_resize_image.call_args[0][0], mock_img2_data_original) + assert mock_resize_image.call_args[0][1] == (10, 10) + + mock_save_image.assert_called_once() + + saved_data = mock_save_image.call_args[0][1] + assert saved_data.shape == (10, 10, 2) # 2 output channels + assert np.all(saved_data[:,:,0] == mock_img1_data) # First channel from img1 + assert np.all(saved_data[:,:,1] == mock_img2_data_resized) # Second channel from resized img2 + + assert merge_rule.id.hex in updated_context.merged_maps_details + details = updated_context.merged_maps_details[merge_rule.id.hex] + assert details['status'] == 'Processed' + assert "merged_ResizedMerge" in details['temp_merged_file'] + mock_log_error.assert_not_called() + mock_log_info.assert_any_call(f"Resized input map from {Path('/fake/img2.png')} from {mock_img2_data_original.shape} to {(10,10)} to match first loaded map.") + mock_log_info.assert_any_call(f"Successfully merged map '{merge_rule.map_type}' for asset '{context.asset_rule.name}'.") +@mock.patch('processing.pipeline.stages.map_merging.ipu.save_image') +@mock.patch('processing.pipeline.stages.map_merging.ipu.resize_image') +@mock.patch('processing.pipeline.stages.map_merging.ipu.load_image') +@mock.patch('logging.info') +@mock.patch('logging.error') +def test_map_merging_to_grayscale_output(mock_log_error, mock_log_info, mock_load_image, mock_resize_image, mock_save_image): + stage = MapMergingStage() + + # Input FileRule (e.g., an RGB image) + input_id = uuid.uuid4() + processed_details = { + input_id.hex: {'temp_processed_file': '/fake/rgb_source.png', 'status': 'Processed', 'map_type': 'RGB_SRC'} + } + # Mock loaded image data (3 channels) + mock_rgb_data = np.full((10, 10, 3), [50, 100, 150], dtype=np.uint8) + mock_load_image.return_value = mock_rgb_data + + # Merge Rule setup: take the Green channel (source_channel=1) from input and map it to the single output channel (target_channel=0) + merge_inputs = [ + create_mock_merge_input_channel(file_rule_id=input_id, source_channel=1, target_channel=0) # G to Grayscale + ] + # output_channels = 1 for grayscale + merge_settings = create_mock_merge_settings(input_maps=merge_inputs, output_channels=1) + merge_rule_id = uuid.uuid4() + merge_rule = create_mock_file_rule_for_merging(id_val=merge_rule_id, map_type="GrayscaleFromGreen", merge_settings=merge_settings) + + context = create_map_merging_mock_context( + initial_file_rules=[merge_rule], + initial_processed_details=processed_details + ) + mock_save_image.return_value = True + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(Path('/fake/rgb_source.png')) + mock_resize_image.assert_not_called() + mock_save_image.assert_called_once() + + saved_data = mock_save_image.call_args[0][1] + assert saved_data.shape == (10, 10) # Grayscale output (2D) + assert np.all(saved_data == 100) # Green channel's value + + assert merge_rule.id.hex in updated_context.merged_maps_details + details = updated_context.merged_maps_details[merge_rule.id.hex] + assert details['status'] == 'Processed' + assert "merged_GrayscaleFromGreen" in details['temp_merged_file'] + mock_log_error.assert_not_called() + mock_log_info.assert_any_call(f"Successfully merged map '{merge_rule.map_type}' for asset '{context.asset_rule.name}'.") + +@mock.patch('processing.pipeline.stages.map_merging.ipu.save_image') +@mock.patch('processing.pipeline.stages.map_merging.ipu.load_image') +@mock.patch('logging.error') +def test_map_merging_default_value_if_missing_channel(mock_log_error, mock_load_image, mock_save_image): + stage = MapMergingStage() + + input_id = uuid.uuid4() + processed_details = { + # Input is a grayscale image (1 channel) + input_id.hex: {'temp_processed_file': '/fake/gray_source.png', 'status': 'Processed', 'map_type': 'GRAY_SRC'} + } + mock_gray_data = np.full((10, 10), 50, dtype=np.uint8) + mock_load_image.return_value = mock_gray_data + + # Merge Rule: try to read source_channel 1 (which doesn't exist in grayscale) + # and use default_value_if_missing for target_channel 0. + # Also, read source_channel 0 (which exists) for target_channel 1. + mic1 = create_mock_merge_input_channel(file_rule_id=input_id, source_channel=1, target_channel=0) + mic1.default_value_if_missing = 128 # Set a specific default value + mic2 = create_mock_merge_input_channel(file_rule_id=input_id, source_channel=0, target_channel=1) + + merge_settings = create_mock_merge_settings(input_maps=[mic1, mic2], output_channels=2) + merge_rule_id = uuid.uuid4() + merge_rule = create_mock_file_rule_for_merging(id_val=merge_rule_id, map_type="DefaultValueTest", merge_settings=merge_settings) + + context = create_map_merging_mock_context( + initial_file_rules=[merge_rule], + initial_processed_details=processed_details + ) + mock_save_image.return_value = True + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(Path('/fake/gray_source.png')) + mock_save_image.assert_called_once() + + saved_data = mock_save_image.call_args[0][1] + assert saved_data.shape == (10, 10, 2) + assert np.all(saved_data[:,:,0] == 128) # Default value for missing source channel 1 + assert np.all(saved_data[:,:,1] == 50) # Value from existing source channel 0 + + assert merge_rule.id.hex in updated_context.merged_maps_details + details = updated_context.merged_maps_details[merge_rule.id.hex] + assert details['status'] == 'Processed' + mock_log_error.assert_not_called() \ No newline at end of file diff --git a/tests/processing/pipeline/stages/test_metadata_finalization_save.py b/tests/processing/pipeline/stages/test_metadata_finalization_save.py new file mode 100644 index 0000000..68741ce --- /dev/null +++ b/tests/processing/pipeline/stages/test_metadata_finalization_save.py @@ -0,0 +1,359 @@ +import pytest +from unittest import mock +from pathlib import Path +import datetime +import json # For comparing dumped content +import uuid +from typing import Optional, Dict, Any + +from processing.pipeline.stages.metadata_finalization_save import MetadataFinalizationAndSaveStage +from processing.pipeline.asset_context import AssetProcessingContext +from rule_structure import AssetRule, SourceRule +from configuration import Configuration, GeneralSettings # Added GeneralSettings as it's in the helper + + +def create_metadata_save_mock_context( + status_flags: Optional[Dict[str, Any]] = None, + initial_asset_metadata: Optional[Dict[str, Any]] = None, + processed_details: Optional[Dict[str, Any]] = None, + merged_details: Optional[Dict[str, Any]] = None, + asset_name: str = "MetaSaveAsset", + output_path_pattern_val: str = "{asset_name}/metadata/{filename}", + # ... other common context fields ... +) -> AssetProcessingContext: + mock_asset_rule = mock.MagicMock(spec=AssetRule) + mock_asset_rule.name = asset_name + mock_asset_rule.output_path_pattern = output_path_pattern_val + mock_asset_rule.id = uuid.uuid4() # Needed for generate_path_from_pattern if it uses it + + mock_source_rule = mock.MagicMock(spec=SourceRule) + mock_source_rule.name = "MetaSaveSource" + + mock_config = mock.MagicMock(spec=Configuration) + # mock_config.general_settings = mock.MagicMock(spec=GeneralSettings) # If needed + + context = AssetProcessingContext( + source_rule=mock_source_rule, + asset_rule=mock_asset_rule, + workspace_path=Path("/fake/workspace"), + engine_temp_dir=Path("/fake/temp_engine_dir"), + output_base_path=Path("/fake/output_base"), # For generate_path + effective_supplier="ValidSupplier", + asset_metadata=initial_asset_metadata if initial_asset_metadata is not None else {}, + processed_maps_details=processed_details if processed_details is not None else {}, + merged_maps_details=merged_details if merged_details is not None else {}, + files_to_process=[], + loaded_data_cache={}, + config_obj=mock_config, + status_flags=status_flags if status_flags is not None else {}, + incrementing_value="001", # Example for path generation + sha5_value="abc" # Example for path generation + ) + return context +@mock.patch('processing.pipeline.stages.metadata_finalization_save.json.dump') +@mock.patch('builtins.open', new_callable=mock.mock_open) +@mock.patch('pathlib.Path.mkdir') +@mock.patch('processing.pipeline.stages.metadata_finalization_save.generate_path_from_pattern') +@mock.patch('datetime.datetime') +def test_asset_skipped_before_metadata_init( + mock_dt, mock_gen_path, mock_mkdir, mock_file_open, mock_json_dump +): + """ + Tests that if an asset is marked for skipping and has no initial metadata, + the stage returns early without attempting to save metadata. + """ + stage = MetadataFinalizationAndSaveStage() + context = create_metadata_save_mock_context( + status_flags={'skip_asset': True}, + initial_asset_metadata={} # Explicitly empty + ) + + updated_context = stage.execute(context) + + # Assert that no processing or saving attempts were made + mock_dt.now.assert_not_called() # Should not even try to set end time if no metadata + mock_gen_path.assert_not_called() + mock_mkdir.assert_not_called() + mock_file_open.assert_not_called() + mock_json_dump.assert_not_called() + + assert updated_context.asset_metadata == {} # Metadata remains empty + assert 'metadata_file_path' not in updated_context.asset_metadata + assert updated_context.status_flags.get('metadata_save_error') is None +@mock.patch('processing.pipeline.stages.metadata_finalization_save.json.dump') +@mock.patch('builtins.open', new_callable=mock.mock_open) +@mock.patch('pathlib.Path.mkdir') +@mock.patch('processing.pipeline.stages.metadata_finalization_save.generate_path_from_pattern') +@mock.patch('datetime.datetime') +def test_asset_skipped_after_metadata_init( + mock_dt, mock_gen_path, mock_mkdir, mock_file_open, mock_json_dump +): + """ + Tests that if an asset is marked for skipping but has initial metadata, + the status is updated to 'Skipped' and metadata is saved. + """ + stage = MetadataFinalizationAndSaveStage() + + fixed_now = datetime.datetime(2023, 1, 1, 12, 0, 0) + mock_dt.now.return_value = fixed_now + + fake_metadata_path_str = "/fake/output_base/SkippedAsset/metadata/SkippedAsset_metadata.json" + mock_gen_path.return_value = fake_metadata_path_str + + initial_meta = {'asset_name': "SkippedAsset", 'status': "Pending"} + + context = create_metadata_save_mock_context( + asset_name="SkippedAsset", + status_flags={'skip_asset': True}, + initial_asset_metadata=initial_meta + ) + + updated_context = stage.execute(context) + + mock_dt.now.assert_called_once() + mock_gen_path.assert_called_once_with( + context.asset_rule.output_path_pattern, + context.asset_rule, + context.source_rule, + context.output_base_path, + context.asset_metadata, # Original metadata passed for path gen + context.incrementing_value, + context.sha5_value, + filename_override=f"{context.asset_rule.name}_metadata.json" + ) + mock_mkdir.assert_called_once_with(parents=True, exist_ok=True) + mock_file_open.assert_called_once_with(Path(fake_metadata_path_str), 'w') + mock_json_dump.assert_called_once() + + dumped_data = mock_json_dump.call_args[0][0] + assert dumped_data['status'] == "Skipped" + assert dumped_data['processing_end_time'] == fixed_now.isoformat() + assert 'processed_map_details' not in dumped_data # Should not be present if skipped early + assert 'merged_map_details' not in dumped_data # Should not be present if skipped early + + assert updated_context.asset_metadata['status'] == "Skipped" + assert updated_context.asset_metadata['processing_end_time'] == fixed_now.isoformat() + assert updated_context.asset_metadata['metadata_file_path'] == fake_metadata_path_str + assert updated_context.status_flags.get('metadata_save_error') is None +@mock.patch('processing.pipeline.stages.metadata_finalization_save.json.dump') +@mock.patch('builtins.open', new_callable=mock.mock_open) # Mocks open() +@mock.patch('pathlib.Path.mkdir') +@mock.patch('processing.pipeline.stages.metadata_finalization_save.generate_path_from_pattern') +@mock.patch('datetime.datetime') +def test_metadata_save_success(mock_dt, mock_gen_path, mock_mkdir, mock_file_open, mock_json_dump): + """ + Tests successful metadata finalization and saving, including serialization of Path objects. + """ + stage = MetadataFinalizationAndSaveStage() + + fixed_now = datetime.datetime(2023, 1, 1, 12, 30, 0) + mock_dt.now.return_value = fixed_now + + fake_metadata_path_str = "/fake/output_base/MetaSaveAsset/metadata/MetaSaveAsset_metadata.json" + mock_gen_path.return_value = fake_metadata_path_str + + initial_meta = {'asset_name': "MetaSaveAsset", 'status': "Pending", 'processing_start_time': "2023-01-01T12:00:00"} + # Example of a Path object that needs serialization + proc_details = {'map1': {'temp_processed_file': Path('/fake/temp_engine_dir/map1.png'), 'final_file_path': Path('/fake/output_base/MetaSaveAsset/map1.png')}} + merged_details = {'merged_map_A': {'output_path': Path('/fake/output_base/MetaSaveAsset/merged_A.png')}} + + context = create_metadata_save_mock_context( + initial_asset_metadata=initial_meta, + processed_details=proc_details, + merged_details=merged_details, + status_flags={} # No errors, no skip + ) + + updated_context = stage.execute(context) + + mock_dt.now.assert_called_once() + mock_gen_path.assert_called_once_with( + context.asset_rule.output_path_pattern, + context.asset_rule, + context.source_rule, + context.output_base_path, + context.asset_metadata, # The metadata *before* adding end_time, status etc. + context.incrementing_value, + context.sha5_value, + filename_override=f"{context.asset_rule.name}_metadata.json" + ) + mock_mkdir.assert_called_once_with(parents=True, exist_ok=True) # Checks parent dir of fake_metadata_path_str + mock_file_open.assert_called_once_with(Path(fake_metadata_path_str), 'w') + mock_json_dump.assert_called_once() + + # Check what was passed to json.dump + dumped_data = mock_json_dump.call_args[0][0] + assert dumped_data['status'] == "Processed" + assert dumped_data['processing_end_time'] == fixed_now.isoformat() + assert 'processing_start_time' in dumped_data # Ensure existing fields are preserved + + # Verify processed_map_details and Path serialization + assert 'processed_map_details' in dumped_data + assert dumped_data['processed_map_details']['map1']['temp_processed_file'] == '/fake/temp_engine_dir/map1.png' + assert dumped_data['processed_map_details']['map1']['final_file_path'] == '/fake/output_base/MetaSaveAsset/map1.png' + + # Verify merged_map_details and Path serialization + assert 'merged_map_details' in dumped_data + assert dumped_data['merged_map_details']['merged_map_A']['output_path'] == '/fake/output_base/MetaSaveAsset/merged_A.png' + + assert updated_context.asset_metadata['metadata_file_path'] == fake_metadata_path_str + assert updated_context.asset_metadata['status'] == "Processed" + assert updated_context.status_flags.get('metadata_save_error') is None +@mock.patch('processing.pipeline.stages.metadata_finalization_save.json.dump') +@mock.patch('builtins.open', new_callable=mock.mock_open) +@mock.patch('pathlib.Path.mkdir') +@mock.patch('processing.pipeline.stages.metadata_finalization_save.generate_path_from_pattern') +@mock.patch('datetime.datetime') +def test_processing_failed_due_to_previous_error( + mock_dt, mock_gen_path, mock_mkdir, mock_file_open, mock_json_dump +): + """ + Tests that if a previous stage set an error flag, the status is 'Failed' + and metadata (including any existing details) is saved. + """ + stage = MetadataFinalizationAndSaveStage() + + fixed_now = datetime.datetime(2023, 1, 1, 12, 45, 0) + mock_dt.now.return_value = fixed_now + + fake_metadata_path_str = "/fake/output_base/FailedAsset/metadata/FailedAsset_metadata.json" + mock_gen_path.return_value = fake_metadata_path_str + + initial_meta = {'asset_name': "FailedAsset", 'status': "Processing"} + # Simulate some details might exist even if a later stage failed + proc_details = {'map1_partial': {'temp_processed_file': Path('/fake/temp_engine_dir/map1_partial.png')}} + + context = create_metadata_save_mock_context( + asset_name="FailedAsset", + initial_asset_metadata=initial_meta, + processed_details=proc_details, + merged_details={}, # No merged details if processing failed before that + status_flags={'file_processing_error': True, 'error_message': "Something went wrong"} + ) + + updated_context = stage.execute(context) + + mock_dt.now.assert_called_once() + mock_gen_path.assert_called_once() # Path generation should still occur + mock_mkdir.assert_called_once_with(parents=True, exist_ok=True) + mock_file_open.assert_called_once_with(Path(fake_metadata_path_str), 'w') + mock_json_dump.assert_called_once() + + dumped_data = mock_json_dump.call_args[0][0] + assert dumped_data['status'] == "Failed" + assert dumped_data['processing_end_time'] == fixed_now.isoformat() + assert 'error_message' in dumped_data # Assuming error messages from status_flags are copied + assert dumped_data['error_message'] == "Something went wrong" + + # Check that existing details are included + assert 'processed_map_details' in dumped_data + assert dumped_data['processed_map_details']['map1_partial']['temp_processed_file'] == '/fake/temp_engine_dir/map1_partial.png' + assert 'merged_map_details' in dumped_data # Should be present, even if empty + assert dumped_data['merged_map_details'] == {} + + assert updated_context.asset_metadata['status'] == "Failed" + assert updated_context.asset_metadata['metadata_file_path'] == fake_metadata_path_str + assert updated_context.status_flags.get('metadata_save_error') is None + # Ensure the original error flag is preserved + assert updated_context.status_flags['file_processing_error'] is True +@mock.patch('processing.pipeline.stages.metadata_finalization_save.json.dump') +@mock.patch('builtins.open', new_callable=mock.mock_open) +@mock.patch('pathlib.Path.mkdir') +@mock.patch('processing.pipeline.stages.metadata_finalization_save.generate_path_from_pattern') +@mock.patch('datetime.datetime') +@mock.patch('logging.error') # To check if error is logged +def test_generate_path_fails( + mock_log_error, mock_dt, mock_gen_path, mock_mkdir, mock_file_open, mock_json_dump +): + """ + Tests behavior when generate_path_from_pattern raises an exception. + Ensures status is updated, error flag is set, and no save is attempted. + """ + stage = MetadataFinalizationAndSaveStage() + + fixed_now = datetime.datetime(2023, 1, 1, 12, 50, 0) + mock_dt.now.return_value = fixed_now + + mock_gen_path.side_effect = Exception("Simulated path generation error") + + initial_meta = {'asset_name': "PathFailAsset", 'status': "Processing"} + context = create_metadata_save_mock_context( + asset_name="PathFailAsset", + initial_asset_metadata=initial_meta, + status_flags={} + ) + + updated_context = stage.execute(context) + + mock_dt.now.assert_called_once() # Time is set before path generation + mock_gen_path.assert_called_once() # generate_path_from_pattern is called + + # File operations should NOT be called if path generation fails + mock_mkdir.assert_not_called() + mock_file_open.assert_not_called() + mock_json_dump.assert_not_called() + + mock_log_error.assert_called_once() # Check that an error was logged + # Example: check if the log message contains relevant info, if needed + # assert "Failed to generate metadata path" in mock_log_error.call_args[0][0] + + assert updated_context.asset_metadata['status'] == "Failed" # Or a more specific error status + assert 'processing_end_time' in updated_context.asset_metadata # End time should still be set + assert updated_context.asset_metadata['processing_end_time'] == fixed_now.isoformat() + assert 'metadata_file_path' not in updated_context.asset_metadata # Path should not be set + + assert updated_context.status_flags.get('metadata_save_error') is True + assert 'error_message' in updated_context.asset_metadata # Check if error message is populated + assert "Simulated path generation error" in updated_context.asset_metadata['error_message'] +@mock.patch('processing.pipeline.stages.metadata_finalization_save.json.dump') +@mock.patch('builtins.open', new_callable=mock.mock_open) +@mock.patch('pathlib.Path.mkdir') +@mock.patch('processing.pipeline.stages.metadata_finalization_save.generate_path_from_pattern') +@mock.patch('datetime.datetime') +@mock.patch('logging.error') # To check if error is logged +def test_json_dump_fails( + mock_log_error, mock_dt, mock_gen_path, mock_mkdir, mock_file_open, mock_json_dump +): + """ + Tests behavior when json.dump raises an exception during saving. + Ensures status is updated, error flag is set, and error is logged. + """ + stage = MetadataFinalizationAndSaveStage() + + fixed_now = datetime.datetime(2023, 1, 1, 12, 55, 0) + mock_dt.now.return_value = fixed_now + + fake_metadata_path_str = "/fake/output_base/JsonDumpFailAsset/metadata/JsonDumpFailAsset_metadata.json" + mock_gen_path.return_value = fake_metadata_path_str + + mock_json_dump.side_effect = IOError("Simulated JSON dump error") # Or TypeError for non-serializable + + initial_meta = {'asset_name': "JsonDumpFailAsset", 'status': "Processing"} + context = create_metadata_save_mock_context( + asset_name="JsonDumpFailAsset", + initial_asset_metadata=initial_meta, + status_flags={} + ) + + updated_context = stage.execute(context) + + mock_dt.now.assert_called_once() + mock_gen_path.assert_called_once() + mock_mkdir.assert_called_once_with(parents=True, exist_ok=True) + mock_file_open.assert_called_once_with(Path(fake_metadata_path_str), 'w') + mock_json_dump.assert_called_once() # json.dump was attempted + + mock_log_error.assert_called_once() + # assert "Failed to save metadata JSON" in mock_log_error.call_args[0][0] + + assert updated_context.asset_metadata['status'] == "Failed" # Or specific "Metadata Save Failed" + assert 'processing_end_time' in updated_context.asset_metadata + assert updated_context.asset_metadata['processing_end_time'] == fixed_now.isoformat() + # metadata_file_path might be set if path generation succeeded, even if dump failed. + # Depending on desired behavior, this could be asserted or not. + # For now, let's assume it's set if path generation was successful. + assert updated_context.asset_metadata['metadata_file_path'] == fake_metadata_path_str + + assert updated_context.status_flags.get('metadata_save_error') is True + assert 'error_message' in updated_context.asset_metadata + assert "Simulated JSON dump error" in updated_context.asset_metadata['error_message'] \ No newline at end of file diff --git a/tests/processing/pipeline/stages/test_metadata_initialization.py b/tests/processing/pipeline/stages/test_metadata_initialization.py new file mode 100644 index 0000000..5b358fd --- /dev/null +++ b/tests/processing/pipeline/stages/test_metadata_initialization.py @@ -0,0 +1,169 @@ +import pytest +from unittest import mock +from pathlib import Path +import datetime +import uuid +from typing import Optional + +from processing.pipeline.stages.metadata_initialization import MetadataInitializationStage +from processing.pipeline.asset_context import AssetProcessingContext +from rule_structure import AssetRule, SourceRule +from configuration import Configuration, GeneralSettings + +# Helper function to create a mock AssetProcessingContext +def create_metadata_init_mock_context( + skip_asset_flag: bool = False, + asset_name: str = "MetaAsset", + asset_id: uuid.UUID = None, # Allow None to default to uuid.uuid4() + source_path_str: str = "source/meta_asset", + output_pattern: str = "{asset_name}/{map_type}", + tags: list = None, + custom_fields: dict = None, + source_rule_name: str = "MetaSource", + source_rule_id: uuid.UUID = None, # Allow None to default to uuid.uuid4() + eff_supplier: Optional[str] = "SupplierMeta", + app_version_str: str = "1.0.0-test", + inc_val: Optional[str] = None, + sha_val: Optional[str] = None +) -> AssetProcessingContext: + mock_asset_rule = mock.MagicMock(spec=AssetRule) + mock_asset_rule.name = asset_name + mock_asset_rule.id = asset_id if asset_id is not None else uuid.uuid4() + mock_asset_rule.source_path = Path(source_path_str) + mock_asset_rule.output_path_pattern = output_pattern + mock_asset_rule.tags = tags if tags is not None else ["tag1", "test_tag"] + mock_asset_rule.custom_fields = custom_fields if custom_fields is not None else {"custom_key": "custom_value"} + + mock_source_rule = mock.MagicMock(spec=SourceRule) + mock_source_rule.name = source_rule_name + mock_source_rule.id = source_rule_id if source_rule_id is not None else uuid.uuid4() + + mock_general_settings = mock.MagicMock(spec=GeneralSettings) + mock_general_settings.app_version = app_version_str + + mock_config = mock.MagicMock(spec=Configuration) + mock_config.general_settings = mock_general_settings + + context = AssetProcessingContext( + source_rule=mock_source_rule, + asset_rule=mock_asset_rule, + workspace_path=Path("/fake/workspace"), + engine_temp_dir=Path("/fake/temp"), + output_base_path=Path("/fake/output"), + effective_supplier=eff_supplier, + asset_metadata={}, + processed_maps_details={}, + merged_maps_details={}, + files_to_process=[], + loaded_data_cache={}, + config_obj=mock_config, + status_flags={'skip_asset': skip_asset_flag}, + incrementing_value=inc_val, + sha5_value=sha_val + ) + return context + +@mock.patch('processing.pipeline.stages.metadata_initialization.datetime') +def test_metadata_initialization_not_skipped(mock_datetime_module): + stage = MetadataInitializationStage() + + fixed_now = datetime.datetime(2023, 10, 26, 12, 0, 0, tzinfo=datetime.timezone.utc) + mock_datetime_module.datetime.now.return_value = fixed_now + + asset_id_val = uuid.uuid4() + source_id_val = uuid.uuid4() + + context = create_metadata_init_mock_context( + skip_asset_flag=False, + asset_id=asset_id_val, + source_rule_id=source_id_val, + inc_val="001", + sha_val="abcde" + ) + + updated_context = stage.execute(context) + + assert isinstance(updated_context.asset_metadata, dict) + assert isinstance(updated_context.processed_maps_details, dict) + assert isinstance(updated_context.merged_maps_details, dict) + + md = updated_context.asset_metadata + assert md['asset_name'] == "MetaAsset" + assert md['asset_id'] == str(asset_id_val) + assert md['source_rule_name'] == "MetaSource" + assert md['source_rule_id'] == str(source_id_val) + assert md['source_path'] == "source/meta_asset" + assert md['effective_supplier'] == "SupplierMeta" + assert md['output_path_pattern'] == "{asset_name}/{map_type}" + assert md['processing_start_time'] == fixed_now.isoformat() + assert md['status'] == "Pending" + assert md['version'] == "1.0.0-test" + assert md['tags'] == ["tag1", "test_tag"] + assert md['custom_fields'] == {"custom_key": "custom_value"} + assert md['incrementing_value'] == "001" + assert md['sha5_value'] == "abcde" + +@mock.patch('processing.pipeline.stages.metadata_initialization.datetime') +def test_metadata_initialization_not_skipped_none_inc_sha(mock_datetime_module): + stage = MetadataInitializationStage() + + fixed_now = datetime.datetime(2023, 10, 26, 12, 0, 0, tzinfo=datetime.timezone.utc) + mock_datetime_module.datetime.now.return_value = fixed_now + + context = create_metadata_init_mock_context( + skip_asset_flag=False, + inc_val=None, + sha_val=None + ) + + updated_context = stage.execute(context) + + md = updated_context.asset_metadata + assert 'incrementing_value' not in md # Or assert md['incrementing_value'] is None, depending on desired behavior + assert 'sha5_value' not in md # Or assert md['sha5_value'] is None + +def test_metadata_initialization_skipped(): + stage = MetadataInitializationStage() + context = create_metadata_init_mock_context(skip_asset_flag=True) + + # Make copies of initial state to ensure they are not modified + initial_asset_metadata = dict(context.asset_metadata) + initial_processed_maps = dict(context.processed_maps_details) + initial_merged_maps = dict(context.merged_maps_details) + + updated_context = stage.execute(context) + + assert updated_context.asset_metadata == initial_asset_metadata + assert updated_context.processed_maps_details == initial_processed_maps + assert updated_context.merged_maps_details == initial_merged_maps + assert not updated_context.asset_metadata # Explicitly check it's empty as per initial setup + assert not updated_context.processed_maps_details + assert not updated_context.merged_maps_details + +@mock.patch('processing.pipeline.stages.metadata_initialization.datetime') +def test_tags_and_custom_fields_are_copies(mock_datetime_module): + stage = MetadataInitializationStage() + fixed_now = datetime.datetime(2023, 10, 26, 12, 0, 0, tzinfo=datetime.timezone.utc) + mock_datetime_module.datetime.now.return_value = fixed_now + + original_tags = ["original_tag"] + original_custom_fields = {"original_key": "original_value"} + + context = create_metadata_init_mock_context( + skip_asset_flag=False, + tags=original_tags, + custom_fields=original_custom_fields + ) + + # Modify originals after context creation but before stage execution + original_tags.append("modified_after_creation") + original_custom_fields["new_key_after_creation"] = "new_value" + + updated_context = stage.execute(context) + + md = updated_context.asset_metadata + assert md['tags'] == ["original_tag"] # Should not have "modified_after_creation" + assert md['tags'] is not original_tags # Ensure it's a different object + + assert md['custom_fields'] == {"original_key": "original_value"} # Should not have "new_key_after_creation" + assert md['custom_fields'] is not original_custom_fields # Ensure it's a different object \ No newline at end of file diff --git a/tests/processing/pipeline/stages/test_normal_map_green_channel.py b/tests/processing/pipeline/stages/test_normal_map_green_channel.py new file mode 100644 index 0000000..3120655 --- /dev/null +++ b/tests/processing/pipeline/stages/test_normal_map_green_channel.py @@ -0,0 +1,323 @@ +import pytest +from unittest import mock +from pathlib import Path +import uuid +import numpy as np +import logging # Added for mocking logger + +from processing.pipeline.stages.normal_map_green_channel import NormalMapGreenChannelStage +from processing.pipeline.asset_context import AssetProcessingContext +from rule_structure import AssetRule, SourceRule, FileRule +from configuration import Configuration, GeneralSettings + +# Helper functions +def create_mock_file_rule_for_normal_test( + id_val: uuid.UUID = None, # Corrected type hint from Optional[uuid.UUID] + map_type: str = "NORMAL", + filename_pattern: str = "normal.png" +) -> mock.MagicMock: + mock_fr = mock.MagicMock(spec=FileRule) + mock_fr.id = id_val if id_val else uuid.uuid4() + mock_fr.map_type = map_type + mock_fr.filename_pattern = filename_pattern + mock_fr.item_type = "MAP_COL" # As per example, though not directly used by stage + mock_fr.active = True # As per example + return mock_fr + +def create_normal_map_mock_context( + initial_file_rules: list = None, # Corrected type hint + initial_processed_details: dict = None, # Corrected type hint + invert_green_globally: bool = False, + skip_asset_flag: bool = False, + asset_name: str = "NormalMapAsset" +) -> AssetProcessingContext: + mock_asset_rule = mock.MagicMock(spec=AssetRule) + mock_asset_rule.name = asset_name + + mock_source_rule = mock.MagicMock(spec=SourceRule) + + mock_gs = mock.MagicMock(spec=GeneralSettings) + mock_gs.invert_normal_map_green_channel_globally = invert_green_globally + + mock_config = mock.MagicMock(spec=Configuration) + mock_config.general_settings = mock_gs + + context = AssetProcessingContext( + source_rule=mock_source_rule, + asset_rule=mock_asset_rule, + workspace_path=Path("/fake/workspace"), + engine_temp_dir=Path("/fake/temp_engine_dir"), + output_base_path=Path("/fake/output"), + effective_supplier="ValidSupplier", + asset_metadata={'asset_name': asset_name}, + processed_maps_details=initial_processed_details if initial_processed_details is not None else {}, + merged_maps_details={}, + files_to_process=list(initial_file_rules) if initial_file_rules else [], + loaded_data_cache={}, + config_obj=mock_config, + status_flags={'skip_asset': skip_asset_flag}, + incrementing_value=None, # Added as per AssetProcessingContext constructor + sha5_value=None # Added as per AssetProcessingContext constructor + ) + return context + +# Unit tests will be added below +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.save_image') +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.load_image') +def test_asset_skipped(mock_load_image, mock_save_image): + stage = NormalMapGreenChannelStage() + normal_fr = create_mock_file_rule_for_normal_test(map_type="NORMAL") + initial_details = { + normal_fr.id.hex: {'temp_processed_file': '/fake/temp_engine_dir/processed_normal.png', 'status': 'Processed', 'map_type': 'NORMAL', 'notes': ''} + } + context = create_normal_map_mock_context( + initial_file_rules=[normal_fr], + initial_processed_details=initial_details, + invert_green_globally=True, + skip_asset_flag=True # Asset is skipped + ) + original_details = context.processed_maps_details.copy() + + updated_context = stage.execute(context) + + mock_load_image.assert_not_called() + mock_save_image.assert_not_called() + assert updated_context.processed_maps_details == original_details + assert normal_fr in updated_context.files_to_process # Ensure rule is still there + +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.save_image') +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.load_image') +def test_no_normal_map_present(mock_load_image, mock_save_image): + stage = NormalMapGreenChannelStage() + # Create a non-normal map rule + diffuse_fr = create_mock_file_rule_for_normal_test(map_type="DIFFUSE", filename_pattern="diffuse.png") + initial_details = { + diffuse_fr.id.hex: {'temp_processed_file': '/fake/temp_engine_dir/processed_diffuse.png', 'status': 'Processed', 'map_type': 'DIFFUSE', 'notes': ''} + } + context = create_normal_map_mock_context( + initial_file_rules=[diffuse_fr], + initial_processed_details=initial_details, + invert_green_globally=True # Inversion enabled, but no normal map + ) + original_details = context.processed_maps_details.copy() + + updated_context = stage.execute(context) + + mock_load_image.assert_not_called() + mock_save_image.assert_not_called() + assert updated_context.processed_maps_details == original_details + assert diffuse_fr in updated_context.files_to_process + +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.save_image') +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.load_image') +def test_normal_map_present_inversion_disabled(mock_load_image, mock_save_image): + stage = NormalMapGreenChannelStage() + normal_rule_id = uuid.uuid4() + normal_fr = create_mock_file_rule_for_normal_test(id_val=normal_rule_id, map_type="NORMAL") + initial_details = { + normal_fr.id.hex: {'temp_processed_file': '/fake/temp_engine_dir/processed_normal.png', 'status': 'Processed', 'map_type': 'NORMAL', 'notes': 'Initial note'} + } + context = create_normal_map_mock_context( + initial_file_rules=[normal_fr], + initial_processed_details=initial_details, + invert_green_globally=False # Inversion disabled + ) + original_details_entry = context.processed_maps_details[normal_fr.id.hex].copy() + + updated_context = stage.execute(context) + + mock_load_image.assert_not_called() + mock_save_image.assert_not_called() + assert updated_context.processed_maps_details[normal_fr.id.hex] == original_details_entry + assert normal_fr in updated_context.files_to_process + +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.save_image') +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.load_image') +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_normal_map_inversion_uint8_success(mock_log_debug, mock_log_info, mock_load_image, mock_save_image): + stage = NormalMapGreenChannelStage() + + normal_rule_id = uuid.uuid4() + normal_fr = create_mock_file_rule_for_normal_test(id_val=normal_rule_id, map_type="NORMAL") + + initial_temp_path = Path('/fake/temp_engine_dir/processed_normal.png') + initial_details = { + normal_fr.id.hex: {'temp_processed_file': str(initial_temp_path), 'status': 'Processed', 'map_type': 'NORMAL', 'notes': 'Initial note'} + } + context = create_normal_map_mock_context( + initial_file_rules=[normal_fr], + initial_processed_details=initial_details, + invert_green_globally=True # Enable inversion + ) + + # R=10, G=50, B=100 + mock_loaded_normal_data = np.array([[[10, 50, 100]]], dtype=np.uint8) + mock_load_image.return_value = mock_loaded_normal_data + mock_save_image.return_value = True # Simulate successful save + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(initial_temp_path) + + # Check that save_image was called with green channel inverted + assert mock_save_image.call_count == 1 + saved_path_arg, saved_data_arg = mock_save_image.call_args[0] + + assert saved_data_arg[0,0,0] == 10 # R unchanged + assert saved_data_arg[0,0,1] == 255 - 50 # G inverted + assert saved_data_arg[0,0,2] == 100 # B unchanged + + assert isinstance(saved_path_arg, Path) + assert "normal_g_inv_" in saved_path_arg.name + assert saved_path_arg.parent == initial_temp_path.parent # Should be in same temp dir + + normal_detail = updated_context.processed_maps_details[normal_fr.id.hex] + assert "normal_g_inv_" in normal_detail['temp_processed_file'] + assert Path(normal_detail['temp_processed_file']).name == saved_path_arg.name + assert "Green channel inverted" in normal_detail['notes'] + assert "Initial note" in normal_detail['notes'] # Check existing notes preserved + + assert normal_fr in updated_context.files_to_process + +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.save_image') +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.load_image') +@mock.patch('logging.info') +@mock.patch('logging.debug') +def test_normal_map_inversion_float_success(mock_log_debug, mock_log_info, mock_load_image, mock_save_image): + stage = NormalMapGreenChannelStage() + normal_rule_id = uuid.uuid4() + normal_fr = create_mock_file_rule_for_normal_test(id_val=normal_rule_id, map_type="NORMAL") + initial_temp_path = Path('/fake/temp_engine_dir/processed_normal_float.png') + initial_details = { + normal_fr.id.hex: {'temp_processed_file': str(initial_temp_path), 'status': 'Processed', 'map_type': 'NORMAL', 'notes': 'Float image'} + } + context = create_normal_map_mock_context( + initial_file_rules=[normal_fr], + initial_processed_details=initial_details, + invert_green_globally=True + ) + + # R=0.1, G=0.25, B=0.75 + mock_loaded_normal_data = np.array([[[0.1, 0.25, 0.75]]], dtype=np.float32) + mock_load_image.return_value = mock_loaded_normal_data + mock_save_image.return_value = True + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(initial_temp_path) + + assert mock_save_image.call_count == 1 + saved_path_arg, saved_data_arg = mock_save_image.call_args[0] + + assert np.isclose(saved_data_arg[0,0,0], 0.1) # R unchanged + assert np.isclose(saved_data_arg[0,0,1], 1.0 - 0.25) # G inverted + assert np.isclose(saved_data_arg[0,0,2], 0.75) # B unchanged + + assert "normal_g_inv_" in saved_path_arg.name + normal_detail = updated_context.processed_maps_details[normal_fr.id.hex] + assert "normal_g_inv_" in normal_detail['temp_processed_file'] + assert "Green channel inverted" in normal_detail['notes'] + assert "Float image" in normal_detail['notes'] + assert normal_fr in updated_context.files_to_process + +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.save_image') +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.load_image') +@mock.patch('logging.error') +def test_load_image_fails(mock_log_error, mock_load_image, mock_save_image): + stage = NormalMapGreenChannelStage() + normal_rule_id = uuid.uuid4() + normal_fr = create_mock_file_rule_for_normal_test(id_val=normal_rule_id, map_type="NORMAL") + initial_temp_path_str = '/fake/temp_engine_dir/processed_normal_load_fail.png' + initial_details = { + normal_fr.id.hex: {'temp_processed_file': initial_temp_path_str, 'status': 'Processed', 'map_type': 'NORMAL', 'notes': 'Load fail test'} + } + context = create_normal_map_mock_context( + initial_file_rules=[normal_fr], + initial_processed_details=initial_details, + invert_green_globally=True + ) + original_details_entry = context.processed_maps_details[normal_fr.id.hex].copy() + + mock_load_image.return_value = None # Simulate load failure + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(Path(initial_temp_path_str)) + mock_save_image.assert_not_called() + mock_log_error.assert_called_once() + assert f"Failed to load image {Path(initial_temp_path_str)} for green channel inversion." in mock_log_error.call_args[0][0] + + # Details should be unchanged + assert updated_context.processed_maps_details[normal_fr.id.hex] == original_details_entry + assert normal_fr in updated_context.files_to_process + +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.save_image') +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.load_image') +@mock.patch('logging.error') +def test_save_image_fails(mock_log_error, mock_load_image, mock_save_image): + stage = NormalMapGreenChannelStage() + normal_rule_id = uuid.uuid4() + normal_fr = create_mock_file_rule_for_normal_test(id_val=normal_rule_id, map_type="NORMAL") + initial_temp_path = Path('/fake/temp_engine_dir/processed_normal_save_fail.png') + initial_details = { + normal_fr.id.hex: {'temp_processed_file': str(initial_temp_path), 'status': 'Processed', 'map_type': 'NORMAL', 'notes': 'Save fail test'} + } + context = create_normal_map_mock_context( + initial_file_rules=[normal_fr], + initial_processed_details=initial_details, + invert_green_globally=True + ) + original_details_entry = context.processed_maps_details[normal_fr.id.hex].copy() + + mock_loaded_normal_data = np.array([[[10, 50, 100]]], dtype=np.uint8) + mock_load_image.return_value = mock_loaded_normal_data + mock_save_image.return_value = False # Simulate save failure + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(initial_temp_path) + mock_save_image.assert_called_once() # Save is attempted + + saved_path_arg = mock_save_image.call_args[0][0] # Get the path it tried to save to + mock_log_error.assert_called_once() + assert f"Failed to save green channel inverted image to {saved_path_arg}." in mock_log_error.call_args[0][0] + + # Details should be unchanged + assert updated_context.processed_maps_details[normal_fr.id.hex] == original_details_entry + assert normal_fr in updated_context.files_to_process + +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.save_image') +@mock.patch('processing.pipeline.stages.normal_map_green_channel.ipu.load_image') +@mock.patch('logging.error') +@pytest.mark.parametrize("unsuitable_data, description", [ + (np.array([[1, 2], [3, 4]], dtype=np.uint8), "2D array"), # 2D array + (np.array([[[1, 2]]], dtype=np.uint8), "2-channel image") # Image with less than 3 channels +]) +def test_image_not_suitable_for_inversion(mock_log_error, mock_load_image, mock_save_image, unsuitable_data, description): + stage = NormalMapGreenChannelStage() + normal_rule_id = uuid.uuid4() + normal_fr = create_mock_file_rule_for_normal_test(id_val=normal_rule_id, map_type="NORMAL") + initial_temp_path_str = f'/fake/temp_engine_dir/unsuitable_{description.replace(" ", "_")}.png' + initial_details = { + normal_fr.id.hex: {'temp_processed_file': initial_temp_path_str, 'status': 'Processed', 'map_type': 'NORMAL', 'notes': f'Unsuitable: {description}'} + } + context = create_normal_map_mock_context( + initial_file_rules=[normal_fr], + initial_processed_details=initial_details, + invert_green_globally=True + ) + original_details_entry = context.processed_maps_details[normal_fr.id.hex].copy() + + mock_load_image.return_value = unsuitable_data + + updated_context = stage.execute(context) + + mock_load_image.assert_called_once_with(Path(initial_temp_path_str)) + mock_save_image.assert_not_called() # Save should not be attempted + mock_log_error.assert_called_once() + assert f"Image at {Path(initial_temp_path_str)} is not suitable for green channel inversion (e.g., not RGB/RGBA)." in mock_log_error.call_args[0][0] + + # Details should be unchanged + assert updated_context.processed_maps_details[normal_fr.id.hex] == original_details_entry + assert normal_fr in updated_context.files_to_process \ No newline at end of file diff --git a/tests/processing/pipeline/stages/test_output_organization.py b/tests/processing/pipeline/stages/test_output_organization.py new file mode 100644 index 0000000..ccf6c08 --- /dev/null +++ b/tests/processing/pipeline/stages/test_output_organization.py @@ -0,0 +1,417 @@ +import pytest +from unittest import mock +from pathlib import Path +import shutil # To check if shutil.copy2 is called +import uuid +from typing import Optional # Added for type hinting in helper + +from processing.pipeline.stages.output_organization import OutputOrganizationStage +from processing.pipeline.asset_context import AssetProcessingContext +from rule_structure import AssetRule, SourceRule, FileRule # For context setup +from configuration import Configuration, GeneralSettings + +def create_output_org_mock_context( + status_flags: Optional[dict] = None, + asset_metadata_status: str = "Processed", # Default to processed for testing copy + processed_map_details: Optional[dict] = None, + merged_map_details: Optional[dict] = None, + overwrite_setting: bool = False, + asset_name: str = "OutputOrgAsset", + output_path_pattern_val: str = "{asset_name}/{map_type}/{filename}" +) -> AssetProcessingContext: + mock_asset_rule = mock.MagicMock(spec=AssetRule) + mock_asset_rule.name = asset_name + mock_asset_rule.output_path_pattern = output_path_pattern_val + # Need FileRules on AssetRule if stage tries to look up output_filename_pattern from them + # For simplicity, assume stage constructs output_filename for now if not found on FileRule + mock_asset_rule.file_rules = [] # Or mock FileRules if stage uses them for output_filename_pattern + + mock_source_rule = mock.MagicMock(spec=SourceRule) + mock_source_rule.name = "OutputOrgSource" + + mock_gs = mock.MagicMock(spec=GeneralSettings) + mock_gs.overwrite_existing = overwrite_setting + + mock_config = mock.MagicMock(spec=Configuration) + mock_config.general_settings = mock_gs + + # Ensure asset_metadata has a status + initial_asset_metadata = {'asset_name': asset_name, 'status': asset_metadata_status} + + context = AssetProcessingContext( + source_rule=mock_source_rule, + asset_rule=mock_asset_rule, + workspace_path=Path("/fake/workspace"), + engine_temp_dir=Path("/fake/temp_engine_dir"), + output_base_path=Path("/fake/output_final"), + effective_supplier="ValidSupplier", + asset_metadata=initial_asset_metadata, + processed_maps_details=processed_map_details if processed_map_details is not None else {}, + merged_maps_details=merged_map_details if merged_map_details is not None else {}, + files_to_process=[], # Not directly used by this stage, but good to have + loaded_data_cache={}, + config_obj=mock_config, + status_flags=status_flags if status_flags is not None else {}, + incrementing_value="001", + sha5_value="xyz" # Corrected from sha5_value to sha256_value if that's the actual param, or ensure it's a valid param. Assuming sha5_value is a typo and should be something like 'unique_id' or similar if not sha256. For now, keeping as sha5_value as per instructions. + ) + return context +@mock.patch('shutil.copy2') +@mock.patch('logging.info') # To check for log messages +def test_output_organization_asset_skipped_by_status_flag(mock_log_info, mock_shutil_copy): + stage = OutputOrganizationStage() + context = create_output_org_mock_context(status_flags={'skip_asset': True}) + + updated_context = stage.execute(context) + + mock_shutil_copy.assert_not_called() + # Check if a log message indicates skipping, if applicable + # e.g., mock_log_info.assert_any_call("Skipping output organization for asset OutputOrgAsset due to skip_asset flag.") + assert 'final_output_files' not in updated_context.asset_metadata # Or assert it's empty + assert updated_context.asset_metadata['status'] == "Processed" # Status should not change if skipped due to flag before stage logic + # Add specific log check if the stage logs this event + # For now, assume no copy is the primary check + +@mock.patch('shutil.copy2') +@mock.patch('logging.warning') # Or info, depending on how failure is logged +def test_output_organization_asset_failed_by_metadata_status(mock_log_warning, mock_shutil_copy): + stage = OutputOrganizationStage() + context = create_output_org_mock_context(asset_metadata_status="Failed") + + updated_context = stage.execute(context) + + mock_shutil_copy.assert_not_called() + # Check for a log message indicating skipping due to failure status + # e.g., mock_log_warning.assert_any_call("Skipping output organization for asset OutputOrgAsset as its status is Failed.") + assert 'final_output_files' not in updated_context.asset_metadata # Or assert it's empty + assert updated_context.asset_metadata['status'] == "Failed" # Status remains Failed + +@mock.patch('shutil.copy2') +@mock.patch('pathlib.Path.mkdir') +@mock.patch('pathlib.Path.exists') +@mock.patch('processing.pipeline.stages.output_organization.generate_path_from_pattern') +@mock.patch('logging.info') +@mock.patch('logging.error') +def test_output_organization_success_no_overwrite( + mock_log_error, mock_log_info, mock_gen_path, mock_path_exists, mock_mkdir, mock_shutil_copy +): + stage = OutputOrganizationStage() + + proc_id_1 = uuid.uuid4().hex + merged_id_1 = uuid.uuid4().hex + + processed_details = { + proc_id_1: {'status': 'Processed', 'temp_processed_file': '/fake/temp_engine_dir/proc1.png', 'map_type': 'Diffuse', 'output_filename': 'OutputOrgAsset_Diffuse.png'} + } + merged_details = { + merged_id_1: {'status': 'Processed', 'temp_merged_file': '/fake/temp_engine_dir/merged1.png', 'map_type': 'ORM', 'output_filename': 'OutputOrgAsset_ORM.png'} + } + + context = create_output_org_mock_context( + processed_map_details=processed_details, + merged_map_details=merged_details, + overwrite_setting=False + ) + + # Mock generate_path_from_pattern to return different paths for each call + final_path_proc1 = Path("/fake/output_final/OutputOrgAsset/Diffuse/OutputOrgAsset_Diffuse.png") + final_path_merged1 = Path("/fake/output_final/OutputOrgAsset/ORM/OutputOrgAsset_ORM.png") + # Ensure generate_path_from_pattern is called with the correct context and details + # The actual call in the stage is: generate_path_from_pattern(context, map_detail, map_type_key, temp_file_key) + # We need to ensure our side_effect matches these calls. + + def gen_path_side_effect(ctx, detail, map_type_key, temp_file_key, output_filename_key): + if detail['temp_processed_file'] == '/fake/temp_engine_dir/proc1.png': + return final_path_proc1 + elif detail['temp_merged_file'] == '/fake/temp_engine_dir/merged1.png': + return final_path_merged1 + raise ValueError("Unexpected call to generate_path_from_pattern") + + mock_gen_path.side_effect = gen_path_side_effect + + mock_path_exists.return_value = False # Files do not exist at destination + + updated_context = stage.execute(context) + + assert mock_shutil_copy.call_count == 2 + mock_shutil_copy.assert_any_call(Path(processed_details[proc_id_1]['temp_processed_file']), final_path_proc1) + mock_shutil_copy.assert_any_call(Path(merged_details[merged_id_1]['temp_merged_file']), final_path_merged1) + + # Check mkdir calls + # It should be called for each unique parent directory + expected_mkdir_calls = [ + mock.call(Path("/fake/output_final/OutputOrgAsset/Diffuse"), parents=True, exist_ok=True), + mock.call(Path("/fake/output_final/OutputOrgAsset/ORM"), parents=True, exist_ok=True) + ] + mock_mkdir.assert_has_calls(expected_mkdir_calls, any_order=True) + # Ensure mkdir was called for the parent of each file + assert mock_mkdir.call_count >= 1 # Could be 1 or 2 if paths share a base that's created once + + assert len(updated_context.asset_metadata['final_output_files']) == 2 + assert str(final_path_proc1) in updated_context.asset_metadata['final_output_files'] + assert str(final_path_merged1) in updated_context.asset_metadata['final_output_files'] + + assert updated_context.processed_maps_details[proc_id_1]['final_output_path'] == str(final_path_proc1) + assert updated_context.merged_maps_details[merged_id_1]['final_output_path'] == str(final_path_merged1) + mock_log_error.assert_not_called() + # Check for specific info logs if necessary + # mock_log_info.assert_any_call(f"Copying {processed_details[proc_id_1]['temp_processed_file']} to {final_path_proc1}") + # mock_log_info.assert_any_call(f"Copying {merged_details[merged_id_1]['temp_merged_file']} to {final_path_merged1}") +@mock.patch('shutil.copy2') +@mock.patch('pathlib.Path.mkdir') # Still might be called if other files are processed +@mock.patch('pathlib.Path.exists') +@mock.patch('processing.pipeline.stages.output_organization.generate_path_from_pattern') +@mock.patch('logging.info') +def test_output_organization_overwrite_disabled_file_exists( + mock_log_info, mock_gen_path, mock_path_exists, mock_mkdir, mock_shutil_copy +): + stage = OutputOrganizationStage() + proc_id_1 = uuid.uuid4().hex + processed_details = { + proc_id_1: {'status': 'Processed', 'temp_processed_file': '/fake/temp_engine_dir/proc_exists.png', 'map_type': 'Diffuse', 'output_filename': 'OutputOrgAsset_Diffuse_Exists.png'} + } + context = create_output_org_mock_context( + processed_map_details=processed_details, + overwrite_setting=False + ) + + final_path_proc1 = Path("/fake/output_final/OutputOrgAsset/Diffuse/OutputOrgAsset_Diffuse_Exists.png") + mock_gen_path.return_value = final_path_proc1 # Only one file + mock_path_exists.return_value = True # File exists at destination + + updated_context = stage.execute(context) + + mock_shutil_copy.assert_not_called() + mock_log_info.assert_any_call( + f"Skipping copy for {final_path_proc1} as it already exists and overwrite is disabled." + ) + # final_output_files should still be populated if the file exists and is considered "organized" + assert str(final_path_proc1) in updated_context.asset_metadata['final_output_files'] + assert updated_context.processed_maps_details[proc_id_1]['final_output_path'] == str(final_path_proc1) + + +@mock.patch('shutil.copy2') +@mock.patch('pathlib.Path.mkdir') +@mock.patch('pathlib.Path.exists') +@mock.patch('processing.pipeline.stages.output_organization.generate_path_from_pattern') +@mock.patch('logging.info') +@mock.patch('logging.error') +def test_output_organization_overwrite_enabled_file_exists( + mock_log_error, mock_log_info, mock_gen_path, mock_path_exists, mock_mkdir, mock_shutil_copy +): + stage = OutputOrganizationStage() + proc_id_1 = uuid.uuid4().hex + processed_details = { + proc_id_1: {'status': 'Processed', 'temp_processed_file': '/fake/temp_engine_dir/proc_overwrite.png', 'map_type': 'Diffuse', 'output_filename': 'OutputOrgAsset_Diffuse_Overwrite.png'} + } + context = create_output_org_mock_context( + processed_map_details=processed_details, + overwrite_setting=True # Overwrite is enabled + ) + + final_path_proc1 = Path("/fake/output_final/OutputOrgAsset/Diffuse/OutputOrgAsset_Diffuse_Overwrite.png") + mock_gen_path.return_value = final_path_proc1 + mock_path_exists.return_value = True # File exists, but we should overwrite + + updated_context = stage.execute(context) + + mock_shutil_copy.assert_called_once_with(Path(processed_details[proc_id_1]['temp_processed_file']), final_path_proc1) + mock_mkdir.assert_called_once_with(final_path_proc1.parent, parents=True, exist_ok=True) + assert str(final_path_proc1) in updated_context.asset_metadata['final_output_files'] + assert updated_context.processed_maps_details[proc_id_1]['final_output_path'] == str(final_path_proc1) + mock_log_error.assert_not_called() + # Optionally check for a log message indicating overwrite, if implemented + # mock_log_info.assert_any_call(f"Overwriting existing file {final_path_proc1}...") + + +@mock.patch('shutil.copy2') +@mock.patch('pathlib.Path.mkdir') +@mock.patch('pathlib.Path.exists') +@mock.patch('processing.pipeline.stages.output_organization.generate_path_from_pattern') +@mock.patch('logging.error') +def test_output_organization_only_processed_maps( + mock_log_error, mock_gen_path, mock_path_exists, mock_mkdir, mock_shutil_copy +): + stage = OutputOrganizationStage() + proc_id_1 = uuid.uuid4().hex + processed_details = { + proc_id_1: {'status': 'Processed', 'temp_processed_file': '/fake/temp_engine_dir/proc_only.png', 'map_type': 'Albedo', 'output_filename': 'OutputOrgAsset_Albedo.png'} + } + context = create_output_org_mock_context( + processed_map_details=processed_details, + merged_map_details={}, # No merged maps + overwrite_setting=False + ) + + final_path_proc1 = Path("/fake/output_final/OutputOrgAsset/Albedo/OutputOrgAsset_Albedo.png") + mock_gen_path.return_value = final_path_proc1 + mock_path_exists.return_value = False + + updated_context = stage.execute(context) + + mock_shutil_copy.assert_called_once_with(Path(processed_details[proc_id_1]['temp_processed_file']), final_path_proc1) + mock_mkdir.assert_called_once_with(final_path_proc1.parent, parents=True, exist_ok=True) + assert len(updated_context.asset_metadata['final_output_files']) == 1 + assert str(final_path_proc1) in updated_context.asset_metadata['final_output_files'] + assert updated_context.processed_maps_details[proc_id_1]['final_output_path'] == str(final_path_proc1) + assert not updated_context.merged_maps_details # Should remain empty + mock_log_error.assert_not_called() + +@mock.patch('shutil.copy2') +@mock.patch('pathlib.Path.mkdir') +@mock.patch('pathlib.Path.exists') +@mock.patch('processing.pipeline.stages.output_organization.generate_path_from_pattern') +@mock.patch('logging.error') +def test_output_organization_only_merged_maps( + mock_log_error, mock_gen_path, mock_path_exists, mock_mkdir, mock_shutil_copy +): + stage = OutputOrganizationStage() + merged_id_1 = uuid.uuid4().hex + merged_details = { + merged_id_1: {'status': 'Processed', 'temp_merged_file': '/fake/temp_engine_dir/merged_only.png', 'map_type': 'Metallic', 'output_filename': 'OutputOrgAsset_Metallic.png'} + } + context = create_output_org_mock_context( + processed_map_details={}, # No processed maps + merged_map_details=merged_details, + overwrite_setting=False + ) + + final_path_merged1 = Path("/fake/output_final/OutputOrgAsset/Metallic/OutputOrgAsset_Metallic.png") + mock_gen_path.return_value = final_path_merged1 + mock_path_exists.return_value = False + + updated_context = stage.execute(context) + + mock_shutil_copy.assert_called_once_with(Path(merged_details[merged_id_1]['temp_merged_file']), final_path_merged1) + mock_mkdir.assert_called_once_with(final_path_merged1.parent, parents=True, exist_ok=True) + assert len(updated_context.asset_metadata['final_output_files']) == 1 + assert str(final_path_merged1) in updated_context.asset_metadata['final_output_files'] + assert updated_context.merged_maps_details[merged_id_1]['final_output_path'] == str(final_path_merged1) + assert not updated_context.processed_maps_details # Should remain empty + mock_log_error.assert_not_called() + +@mock.patch('shutil.copy2') +@mock.patch('pathlib.Path.mkdir') +@mock.patch('pathlib.Path.exists') +@mock.patch('processing.pipeline.stages.output_organization.generate_path_from_pattern') +@mock.patch('logging.warning') # Expect a warning for skipped map +@mock.patch('logging.error') +def test_output_organization_map_status_not_processed( + mock_log_error, mock_log_warning, mock_gen_path, mock_path_exists, mock_mkdir, mock_shutil_copy +): + stage = OutputOrganizationStage() + + proc_id_1_failed = uuid.uuid4().hex + proc_id_2_ok = uuid.uuid4().hex + + processed_details = { + proc_id_1_failed: {'status': 'Failed', 'temp_processed_file': '/fake/temp_engine_dir/proc_failed.png', 'map_type': 'Diffuse', 'output_filename': 'OutputOrgAsset_Diffuse_Failed.png'}, + proc_id_2_ok: {'status': 'Processed', 'temp_processed_file': '/fake/temp_engine_dir/proc_ok.png', 'map_type': 'Normal', 'output_filename': 'OutputOrgAsset_Normal_OK.png'} + } + context = create_output_org_mock_context( + processed_map_details=processed_details, + overwrite_setting=False + ) + + final_path_proc_ok = Path("/fake/output_final/OutputOrgAsset/Normal/OutputOrgAsset_Normal_OK.png") + # generate_path_from_pattern should only be called for the 'Processed' map + mock_gen_path.return_value = final_path_proc_ok + mock_path_exists.return_value = False + + updated_context = stage.execute(context) + + # Assert copy was only called for the 'Processed' map + mock_shutil_copy.assert_called_once_with(Path(processed_details[proc_id_2_ok]['temp_processed_file']), final_path_proc_ok) + mock_mkdir.assert_called_once_with(final_path_proc_ok.parent, parents=True, exist_ok=True) + + # Assert final_output_files only contains the successfully processed map + assert len(updated_context.asset_metadata['final_output_files']) == 1 + assert str(final_path_proc_ok) in updated_context.asset_metadata['final_output_files'] + + # Assert final_output_path is set for the processed map + assert updated_context.processed_maps_details[proc_id_2_ok]['final_output_path'] == str(final_path_proc_ok) + # Assert final_output_path is NOT set for the failed map + assert 'final_output_path' not in updated_context.processed_maps_details[proc_id_1_failed] + + mock_log_warning.assert_any_call( + f"Skipping output organization for map with ID {proc_id_1_failed} (type: Diffuse) as its status is 'Failed'." + ) + mock_log_error.assert_not_called() +@mock.patch('shutil.copy2') +@mock.patch('pathlib.Path.mkdir') +@mock.patch('pathlib.Path.exists') +@mock.patch('processing.pipeline.stages.output_organization.generate_path_from_pattern') +@mock.patch('logging.error') +def test_output_organization_generate_path_fails( + mock_log_error, mock_gen_path, mock_path_exists, mock_mkdir, mock_shutil_copy +): + stage = OutputOrganizationStage() + proc_id_1 = uuid.uuid4().hex + processed_details = { + proc_id_1: {'status': 'Processed', 'temp_processed_file': '/fake/temp_engine_dir/proc_path_fail.png', 'map_type': 'Roughness', 'output_filename': 'OutputOrgAsset_Roughness_PathFail.png'} + } + context = create_output_org_mock_context( + processed_map_details=processed_details, + overwrite_setting=False + ) + + mock_gen_path.side_effect = Exception("Simulated path generation error") + mock_path_exists.return_value = False # Should not matter if path gen fails + + updated_context = stage.execute(context) + + mock_shutil_copy.assert_not_called() # No copy if path generation fails + mock_mkdir.assert_not_called() # No mkdir if path generation fails + + assert not updated_context.asset_metadata.get('final_output_files') # No files should be listed + assert 'final_output_path' not in updated_context.processed_maps_details[proc_id_1] + + assert updated_context.status_flags.get('output_organization_error') is True + assert updated_context.asset_metadata['status'] == "Error" # Or "Failed" depending on desired behavior + + mock_log_error.assert_any_call( + f"Error generating output path for map ID {proc_id_1} (type: Roughness): Simulated path generation error" + ) + +@mock.patch('shutil.copy2') +@mock.patch('pathlib.Path.mkdir') +@mock.patch('pathlib.Path.exists') +@mock.patch('processing.pipeline.stages.output_organization.generate_path_from_pattern') +@mock.patch('logging.error') +def test_output_organization_shutil_copy_fails( + mock_log_error, mock_gen_path, mock_path_exists, mock_mkdir, mock_shutil_copy +): + stage = OutputOrganizationStage() + proc_id_1 = uuid.uuid4().hex + processed_details = { + proc_id_1: {'status': 'Processed', 'temp_processed_file': '/fake/temp_engine_dir/proc_copy_fail.png', 'map_type': 'AO', 'output_filename': 'OutputOrgAsset_AO_CopyFail.png'} + } + context = create_output_org_mock_context( + processed_map_details=processed_details, + overwrite_setting=False + ) + + final_path_proc1 = Path("/fake/output_final/OutputOrgAsset/AO/OutputOrgAsset_AO_CopyFail.png") + mock_gen_path.return_value = final_path_proc1 + mock_path_exists.return_value = False + mock_shutil_copy.side_effect = shutil.Error("Simulated copy error") # Can also be IOError, OSError + + updated_context = stage.execute(context) + + mock_mkdir.assert_called_once_with(final_path_proc1.parent, parents=True, exist_ok=True) # mkdir would be called before copy + mock_shutil_copy.assert_called_once_with(Path(processed_details[proc_id_1]['temp_processed_file']), final_path_proc1) + + # Even if copy fails, the path might be added to final_output_files before the error is caught, + # or the design might be to not add it. Let's assume it's not added on error. + # Check the stage's actual behavior for this. + # If the intention is to record the *attempted* path, this assertion might change. + # For now, assume failure means it's not a "final" output. + assert not updated_context.asset_metadata.get('final_output_files') + assert 'final_output_path' not in updated_context.processed_maps_details[proc_id_1] # Or it might contain the path but status is error + + assert updated_context.status_flags.get('output_organization_error') is True + assert updated_context.asset_metadata['status'] == "Error" # Or "Failed" + + mock_log_error.assert_any_call( + f"Error copying file {processed_details[proc_id_1]['temp_processed_file']} to {final_path_proc1}: Simulated copy error" + ) \ No newline at end of file diff --git a/tests/processing/pipeline/stages/test_supplier_determination.py b/tests/processing/pipeline/stages/test_supplier_determination.py new file mode 100644 index 0000000..a1613b1 --- /dev/null +++ b/tests/processing/pipeline/stages/test_supplier_determination.py @@ -0,0 +1,213 @@ +import pytest +from unittest import mock +from pathlib import Path +from typing import Dict, List, Optional, Any + +# Assuming pytest is run from project root, adjust if necessary +from processing.pipeline.stages.supplier_determination import SupplierDeterminationStage +from processing.pipeline.asset_context import AssetProcessingContext +from rule_structure import AssetRule, SourceRule, FileRule # For constructing mock context +from configuration import Configuration, GeneralSettings, Supplier # For mock config + +# Example helper (can be a pytest fixture too) +def create_mock_context( + asset_rule_supplier_override: Optional[str] = None, + source_rule_supplier: Optional[str] = None, + config_suppliers: Optional[Dict[str, Any]] = None, # Mocked Supplier objects or dicts + asset_name: str = "TestAsset" +) -> AssetProcessingContext: + mock_asset_rule = mock.MagicMock(spec=AssetRule) + mock_asset_rule.name = asset_name + mock_asset_rule.supplier_override = asset_rule_supplier_override + # ... other AssetRule fields if needed by the stage ... + + mock_source_rule = mock.MagicMock(spec=SourceRule) + mock_source_rule.supplier = source_rule_supplier + # ... other SourceRule fields ... + + mock_config = mock.MagicMock(spec=Configuration) + mock_config.suppliers = config_suppliers if config_suppliers is not None else {} + + # Basic AssetProcessingContext fields + context = AssetProcessingContext( + source_rule=mock_source_rule, + asset_rule=mock_asset_rule, + workspace_path=Path("/fake/workspace"), + engine_temp_dir=Path("/fake/temp"), + output_base_path=Path("/fake/output"), + effective_supplier=None, + asset_metadata={}, + processed_maps_details={}, + merged_maps_details={}, + files_to_process=[], + loaded_data_cache={}, + config_obj=mock_config, + status_flags={}, + incrementing_value=None, + sha5_value=None # Corrected from sha5_value to sha256_value if that's the actual field name + ) + return context + +@pytest.fixture +def supplier_stage(): + return SupplierDeterminationStage() + +@mock.patch('logging.error') +@mock.patch('logging.info') +def test_supplier_from_asset_rule_override_valid(mock_log_info, mock_log_error, supplier_stage): + mock_suppliers_config = {"SupplierA": mock.MagicMock(spec=Supplier)} + context = create_mock_context( + asset_rule_supplier_override="SupplierA", + config_suppliers=mock_suppliers_config + ) + + updated_context = supplier_stage.execute(context) + + assert updated_context.effective_supplier == "SupplierA" + assert not updated_context.status_flags.get('supplier_error') + mock_log_info.assert_any_call("Effective supplier for asset 'TestAsset' set to 'SupplierA' from asset rule override.") + mock_log_error.assert_not_called() + +@mock.patch('logging.error') +@mock.patch('logging.info') +def test_supplier_from_source_rule_fallback_valid(mock_log_info, mock_log_error, supplier_stage): + mock_suppliers_config = {"SupplierB": mock.MagicMock(spec=Supplier)} + context = create_mock_context( + asset_rule_supplier_override=None, + source_rule_supplier="SupplierB", + config_suppliers=mock_suppliers_config + ) + + updated_context = supplier_stage.execute(context) + + assert updated_context.effective_supplier == "SupplierB" + assert not updated_context.status_flags.get('supplier_error') + mock_log_info.assert_any_call("Effective supplier for asset 'TestAsset' set to 'SupplierB' from source rule.") + mock_log_error.assert_not_called() + +@mock.patch('logging.error') +@mock.patch('logging.warning') # supplier_determination uses logging.warning for invalid suppliers +def test_asset_rule_override_invalid_supplier(mock_log_warning, mock_log_error, supplier_stage): + context = create_mock_context( + asset_rule_supplier_override="InvalidSupplier", + config_suppliers={"SupplierA": mock.MagicMock(spec=Supplier)} # "InvalidSupplier" not in config + ) + + updated_context = supplier_stage.execute(context) + + assert updated_context.effective_supplier is None + assert updated_context.status_flags.get('supplier_error') is True + mock_log_warning.assert_any_call( + "Asset 'TestAsset' has supplier_override 'InvalidSupplier' which is not defined in global suppliers. No supplier set." + ) + mock_log_error.assert_not_called() + + +@mock.patch('logging.error') +@mock.patch('logging.warning') +def test_source_rule_fallback_invalid_supplier(mock_log_warning, mock_log_error, supplier_stage): + context = create_mock_context( + asset_rule_supplier_override=None, + source_rule_supplier="InvalidSupplierB", + config_suppliers={"SupplierA": mock.MagicMock(spec=Supplier)} # "InvalidSupplierB" not in config + ) + + updated_context = supplier_stage.execute(context) + + assert updated_context.effective_supplier is None + assert updated_context.status_flags.get('supplier_error') is True + mock_log_warning.assert_any_call( + "Asset 'TestAsset' has source rule supplier 'InvalidSupplierB' which is not defined in global suppliers. No supplier set." + ) + mock_log_error.assert_not_called() + +@mock.patch('logging.error') +@mock.patch('logging.warning') +def test_no_supplier_defined(mock_log_warning, mock_log_error, supplier_stage): + context = create_mock_context( + asset_rule_supplier_override=None, + source_rule_supplier=None, + config_suppliers={"SupplierA": mock.MagicMock(spec=Supplier)} + ) + + updated_context = supplier_stage.execute(context) + + assert updated_context.effective_supplier is None + assert updated_context.status_flags.get('supplier_error') is True + mock_log_warning.assert_any_call( + "No supplier could be determined for asset 'TestAsset'. " + "AssetRule override is None and SourceRule supplier is None or empty." + ) + mock_log_error.assert_not_called() + +@mock.patch('logging.error') +@mock.patch('logging.warning') +def test_empty_config_suppliers_with_asset_override(mock_log_warning, mock_log_error, supplier_stage): + context = create_mock_context( + asset_rule_supplier_override="SupplierX", + config_suppliers={} # Empty global supplier config + ) + + updated_context = supplier_stage.execute(context) + + assert updated_context.effective_supplier is None + assert updated_context.status_flags.get('supplier_error') is True + mock_log_warning.assert_any_call( + "Asset 'TestAsset' has supplier_override 'SupplierX' which is not defined in global suppliers. No supplier set." + ) + mock_log_error.assert_not_called() + +@mock.patch('logging.error') +@mock.patch('logging.warning') +def test_empty_config_suppliers_with_source_rule(mock_log_warning, mock_log_error, supplier_stage): + context = create_mock_context( + source_rule_supplier="SupplierY", + config_suppliers={} # Empty global supplier config + ) + + updated_context = supplier_stage.execute(context) + + assert updated_context.effective_supplier is None + assert updated_context.status_flags.get('supplier_error') is True + mock_log_warning.assert_any_call( + "Asset 'TestAsset' has source rule supplier 'SupplierY' which is not defined in global suppliers. No supplier set." + ) + mock_log_error.assert_not_called() + +@mock.patch('logging.error') +@mock.patch('logging.info') +def test_asset_rule_override_empty_string(mock_log_info, mock_log_error, supplier_stage): + # This scenario should fall back to source_rule.supplier if asset_rule.supplier_override is "" + mock_suppliers_config = {"SupplierB": mock.MagicMock(spec=Supplier)} + context = create_mock_context( + asset_rule_supplier_override="", # Empty string override + source_rule_supplier="SupplierB", + config_suppliers=mock_suppliers_config + ) + + updated_context = supplier_stage.execute(context) + + assert updated_context.effective_supplier == "SupplierB" # Falls back to SourceRule + assert not updated_context.status_flags.get('supplier_error') + mock_log_info.assert_any_call("Effective supplier for asset 'TestAsset' set to 'SupplierB' from source rule.") + mock_log_error.assert_not_called() + +@mock.patch('logging.error') +@mock.patch('logging.warning') +def test_source_rule_supplier_empty_string(mock_log_warning, mock_log_error, supplier_stage): + # This scenario should result in an error if asset_rule.supplier_override is None and source_rule.supplier is "" + context = create_mock_context( + asset_rule_supplier_override=None, + source_rule_supplier="", # Empty string source supplier + config_suppliers={"SupplierA": mock.MagicMock(spec=Supplier)} + ) + + updated_context = supplier_stage.execute(context) + + assert updated_context.effective_supplier is None + assert updated_context.status_flags.get('supplier_error') is True + mock_log_warning.assert_any_call( + "No supplier could be determined for asset 'TestAsset'. " + "AssetRule override is None and SourceRule supplier is None or empty." + ) + mock_log_error.assert_not_called() \ No newline at end of file diff --git a/tests/processing/pipeline/test_orchestrator.py b/tests/processing/pipeline/test_orchestrator.py new file mode 100644 index 0000000..3f52908 --- /dev/null +++ b/tests/processing/pipeline/test_orchestrator.py @@ -0,0 +1,383 @@ +import pytest +from unittest import mock +from pathlib import Path +import uuid +import shutil # For checking rmtree +import tempfile # For mocking mkdtemp + +from processing.pipeline.orchestrator import PipelineOrchestrator +from processing.pipeline.asset_context import AssetProcessingContext +from processing.pipeline.stages.base_stage import ProcessingStage # For mocking stages +from rule_structure import SourceRule, AssetRule, FileRule +from configuration import Configuration, GeneralSettings + +# Mock Stage that modifies context +class MockPassThroughStage(ProcessingStage): + def __init__(self, stage_name="mock_stage"): + self.stage_name = stage_name + self.execute_call_count = 0 + self.contexts_called_with = [] # To store contexts for verification + + def execute(self, context: AssetProcessingContext) -> AssetProcessingContext: + self.execute_call_count += 1 + self.contexts_called_with.append(context) + # Optionally, modify context for testing + context.asset_metadata[f'{self.stage_name}_executed'] = True + if self.stage_name == "skipper_stage": # Example conditional logic + context.status_flags['skip_asset'] = True + context.status_flags['skip_reason'] = "Skipped by skipper_stage" + elif self.stage_name == "error_stage": # Example error-raising stage + raise ValueError("Simulated error in error_stage") + + # Simulate status update based on stage execution + if not context.status_flags.get('skip_asset') and not context.status_flags.get('asset_failed'): + context.asset_metadata['status'] = "Processed" # Default to processed if not skipped/failed + return context + +def create_orchestrator_test_config() -> mock.MagicMock: + mock_config = mock.MagicMock(spec=Configuration) + mock_config.general_settings = mock.MagicMock(spec=GeneralSettings) + mock_config.general_settings.temp_dir_override = None # Default, can be overridden in tests + # Add other config details if orchestrator or stages depend on them directly + return mock_config + +def create_orchestrator_test_asset_rule(name: str, num_file_rules: int = 1) -> mock.MagicMock: + asset_rule = mock.MagicMock(spec=AssetRule) + asset_rule.name = name + asset_rule.id = uuid.uuid4() + asset_rule.source_path = Path(f"/fake/source/{name}") # Using Path object + asset_rule.file_rules = [mock.MagicMock(spec=FileRule) for _ in range(num_file_rules)] + asset_rule.enabled = True + asset_rule.map_types = {} # Initialize as dict + asset_rule.material_name_scheme = "{asset_name}" + asset_rule.texture_name_scheme = "{asset_name}_{map_type}" + asset_rule.output_path_scheme = "{source_name}/{asset_name}" + # ... other necessary AssetRule fields ... + return asset_rule + +def create_orchestrator_test_source_rule(name: str, num_assets: int = 1, asset_names: list = None) -> mock.MagicMock: + source_rule = mock.MagicMock(spec=SourceRule) + source_rule.name = name + source_rule.id = uuid.uuid4() + if asset_names: + source_rule.assets = [create_orchestrator_test_asset_rule(an) for an in asset_names] + else: + source_rule.assets = [create_orchestrator_test_asset_rule(f"Asset_{i+1}_in_{name}") for i in range(num_assets)] + source_rule.enabled = True + source_rule.source_path = Path(f"/fake/source_root/{name}") # Using Path object + # ... other necessary SourceRule fields ... + return source_rule + +# --- Test Cases for PipelineOrchestrator.process_source_rule() --- + +@mock.patch('shutil.rmtree') +@mock.patch('tempfile.mkdtemp') +def test_orchestrator_basic_flow_mock_stages(mock_mkdtemp, mock_rmtree): + mock_mkdtemp.return_value = "/fake/engine_temp_dir_path" # Path for mkdtemp + + config = create_orchestrator_test_config() + stage1 = MockPassThroughStage("stage1") + stage2 = MockPassThroughStage("stage2") + orchestrator = PipelineOrchestrator(config_obj=config, stages=[stage1, stage2]) + + source_rule = create_orchestrator_test_source_rule("MySourceRule", num_assets=2) + asset1_name = source_rule.assets[0].name + asset2_name = source_rule.assets[1].name + + # Mock asset_metadata to be updated by stages for status check + # The MockPassThroughStage already sets a 'status' = "Processed" if not skipped/failed + # and adds '{stage_name}_executed' = True to asset_metadata. + + results = orchestrator.process_source_rule( + source_rule, Path("/ws"), Path("/out"), False, "inc_val_123", "sha_val_abc" + ) + + assert stage1.execute_call_count == 2 # Called for each asset + assert stage2.execute_call_count == 2 # Called for each asset + + assert asset1_name in results['processed'] + assert asset2_name in results['processed'] + assert not results['skipped'] + assert not results['failed'] + + # Verify context modifications by stages + for i in range(2): # For each asset + # Stage 1 context checks + s1_context_asset = stage1.contexts_called_with[i] + assert s1_context_asset.asset_metadata.get('stage1_executed') is True + assert s1_context_asset.asset_metadata.get('stage2_executed') is None # Stage 2 not yet run for this asset + + # Stage 2 context checks + s2_context_asset = stage2.contexts_called_with[i] + assert s2_context_asset.asset_metadata.get('stage1_executed') is True # From stage 1 + assert s2_context_asset.asset_metadata.get('stage2_executed') is True + assert s2_context_asset.asset_metadata.get('status') == "Processed" + + mock_mkdtemp.assert_called_once() + # The orchestrator creates a subdirectory within the mkdtemp path + expected_temp_path = Path(mock_mkdtemp.return_value) / source_rule.id.hex + mock_rmtree.assert_called_once_with(expected_temp_path, ignore_errors=True) + +@mock.patch('shutil.rmtree') +@mock.patch('tempfile.mkdtemp') +def test_orchestrator_asset_skipping_by_stage(mock_mkdtemp, mock_rmtree): + mock_mkdtemp.return_value = "/fake/engine_temp_dir_path_skip" + + config = create_orchestrator_test_config() + skipper_stage = MockPassThroughStage("skipper_stage") # This stage will set skip_asset = True + stage_after_skip = MockPassThroughStage("stage_after_skip") + + orchestrator = PipelineOrchestrator(config_obj=config, stages=[skipper_stage, stage_after_skip]) + + source_rule = create_orchestrator_test_source_rule("SkipSourceRule", num_assets=1) + asset_to_skip_name = source_rule.assets[0].name + + results = orchestrator.process_source_rule( + source_rule, Path("/ws_skip"), Path("/out_skip"), False, "inc_skip", "sha_skip" + ) + + assert skipper_stage.execute_call_count == 1 # Called for the asset + assert stage_after_skip.execute_call_count == 0 # Not called because asset was skipped + + assert asset_to_skip_name in results['skipped'] + assert not results['processed'] + assert not results['failed'] + + # Verify skip reason in context if needed (MockPassThroughStage stores contexts) + skipped_context = skipper_stage.contexts_called_with[0] + assert skipped_context.status_flags['skip_asset'] is True + assert skipped_context.status_flags['skip_reason'] == "Skipped by skipper_stage" + + mock_mkdtemp.assert_called_once() + expected_temp_path = Path(mock_mkdtemp.return_value) / source_rule.id.hex + mock_rmtree.assert_called_once_with(expected_temp_path, ignore_errors=True) + +@mock.patch('shutil.rmtree') +@mock.patch('tempfile.mkdtemp') +def test_orchestrator_no_assets_in_source_rule(mock_mkdtemp, mock_rmtree): + mock_mkdtemp.return_value = "/fake/engine_temp_dir_no_assets" + + config = create_orchestrator_test_config() + stage1 = MockPassThroughStage("stage1_no_assets") + orchestrator = PipelineOrchestrator(config_obj=config, stages=[stage1]) + + source_rule = create_orchestrator_test_source_rule("NoAssetSourceRule", num_assets=0) + + results = orchestrator.process_source_rule( + source_rule, Path("/ws_no_assets"), Path("/out_no_assets"), False, "inc_no", "sha_no" + ) + + assert stage1.execute_call_count == 0 + assert not results['processed'] + assert not results['skipped'] + assert not results['failed'] + + # mkdtemp should still be called for the source rule processing, even if no assets + mock_mkdtemp.assert_called_once() + expected_temp_path = Path(mock_mkdtemp.return_value) / source_rule.id.hex + mock_rmtree.assert_called_once_with(expected_temp_path, ignore_errors=True) + + +@mock.patch('shutil.rmtree') +@mock.patch('tempfile.mkdtemp') +def test_orchestrator_error_during_stage_execution(mock_mkdtemp, mock_rmtree): + mock_mkdtemp.return_value = "/fake/engine_temp_dir_error" + + config = create_orchestrator_test_config() + error_stage = MockPassThroughStage("error_stage") # This stage will raise an error + stage_after_error = MockPassThroughStage("stage_after_error") + + orchestrator = PipelineOrchestrator(config_obj=config, stages=[error_stage, stage_after_error]) + + # Test with two assets, one fails, one processes (if orchestrator continues) + # The current orchestrator's process_asset is per asset, so an error in one + # should not stop processing of other assets in the same source_rule. + source_rule = create_orchestrator_test_source_rule("ErrorSourceRule", asset_names=["AssetFails", "AssetSucceeds"]) + asset_fails_name = source_rule.assets[0].name + asset_succeeds_name = source_rule.assets[1].name + + # Make only the first asset's processing trigger the error + original_execute = error_stage.execute + def error_execute_side_effect(context: AssetProcessingContext): + if context.asset_rule.name == asset_fails_name: + # The MockPassThroughStage is already configured to raise ValueError for "error_stage" + # but we need to ensure it's only for the first asset. + # We can achieve this by modifying the stage_name temporarily or by checking asset_rule.name + # For simplicity, let's assume the mock stage's error logic is fine, + # and we just need to check the outcome. + # The error_stage will raise ValueError("Simulated error in error_stage") + # The orchestrator's _process_single_asset catches generic Exception. + return original_execute(context) # This will call the erroring logic + else: + # For the second asset, make it pass through without error + context.asset_metadata[f'{error_stage.stage_name}_executed'] = True + context.asset_metadata['status'] = "Processed" + return context + + error_stage.execute = mock.MagicMock(side_effect=error_execute_side_effect) + # stage_after_error should still be called for the successful asset + + results = orchestrator.process_source_rule( + source_rule, Path("/ws_error"), Path("/out_error"), False, "inc_err", "sha_err" + ) + + assert error_stage.execute.call_count == 2 # Called for both assets + # stage_after_error is only called for the asset that didn't fail in error_stage + assert stage_after_error.execute_call_count == 1 + + assert asset_fails_name in results['failed'] + assert asset_succeeds_name in results['processed'] + assert not results['skipped'] + + # Verify the context of the failed asset + failed_context = None + for ctx in error_stage.contexts_called_with: + if ctx.asset_rule.name == asset_fails_name: + failed_context = ctx + break + assert failed_context is not None + assert failed_context.status_flags['asset_failed'] is True + assert "Simulated error in error_stage" in failed_context.status_flags['failure_reason'] + + # Verify the context of the successful asset after stage_after_error + successful_context_after_s2 = None + for ctx in stage_after_error.contexts_called_with: + if ctx.asset_rule.name == asset_succeeds_name: + successful_context_after_s2 = ctx + break + assert successful_context_after_s2 is not None + assert successful_context_after_s2.asset_metadata.get('error_stage_executed') is True # from the non-erroring path + assert successful_context_after_s2.asset_metadata.get('stage_after_error_executed') is True + assert successful_context_after_s2.asset_metadata.get('status') == "Processed" + + + mock_mkdtemp.assert_called_once() + expected_temp_path = Path(mock_mkdtemp.return_value) / source_rule.id.hex + mock_rmtree.assert_called_once_with(expected_temp_path, ignore_errors=True) + + +@mock.patch('shutil.rmtree') +@mock.patch('tempfile.mkdtemp') +def test_orchestrator_asset_processing_context_initialization(mock_mkdtemp, mock_rmtree): + mock_engine_temp_dir = "/fake/engine_temp_dir_context_init" + mock_mkdtemp.return_value = mock_engine_temp_dir + + config = create_orchestrator_test_config() + mock_stage = MockPassThroughStage("context_check_stage") + orchestrator = PipelineOrchestrator(config_obj=config, stages=[mock_stage]) + + source_rule = create_orchestrator_test_source_rule("ContextSourceRule", num_assets=1) + asset_rule = source_rule.assets[0] + + workspace_path = Path("/ws_context") + output_base_path = Path("/out_context") + incrementing_value = "inc_context_123" + sha5_value = "sha_context_abc" + + orchestrator.process_source_rule( + source_rule, workspace_path, output_base_path, False, incrementing_value, sha5_value + ) + + assert mock_stage.execute_call_count == 1 + + # Retrieve the context passed to the mock stage + captured_context = mock_stage.contexts_called_with[0] + + assert captured_context.source_rule == source_rule + assert captured_context.asset_rule == asset_rule + assert captured_context.workspace_path == workspace_path + + # engine_temp_dir for the asset is a sub-directory of the source_rule's temp dir + # which itself is a sub-directory of the main engine_temp_dir from mkdtemp + expected_source_rule_temp_dir = Path(mock_engine_temp_dir) / source_rule.id.hex + expected_asset_temp_dir = expected_source_rule_temp_dir / asset_rule.id.hex + assert captured_context.engine_temp_dir == expected_asset_temp_dir + + assert captured_context.output_base_path == output_base_path + assert captured_context.config_obj == config + assert captured_context.incrementing_value == incrementing_value + assert captured_context.sha5_value == sha5_value + + # Check initial state of other context fields + assert captured_context.asset_metadata == {} # Should be empty initially for an asset + assert captured_context.status_flags == {} # Should be empty initially + assert captured_context.shared_data == {} # Should be empty initially + assert captured_context.current_files == [] # Should be empty initially + + mock_mkdtemp.assert_called_once() + mock_rmtree.assert_called_once_with(expected_source_rule_temp_dir, ignore_errors=True) + +@mock.patch('shutil.rmtree') +@mock.patch('tempfile.mkdtemp') +def test_orchestrator_temp_dir_override_from_config(mock_mkdtemp, mock_rmtree): + # This test verifies that if config.general_settings.temp_dir_override is set, + # mkdtemp is NOT called, and the override path is used and cleaned up. + + config = create_orchestrator_test_config() + override_temp_path_str = "/override/temp/path" + config.general_settings.temp_dir_override = override_temp_path_str + + stage1 = MockPassThroughStage("stage_temp_override") + orchestrator = PipelineOrchestrator(config_obj=config, stages=[stage1]) + + source_rule = create_orchestrator_test_source_rule("TempOverrideRule", num_assets=1) + asset_rule = source_rule.assets[0] + + results = orchestrator.process_source_rule( + source_rule, Path("/ws_override"), Path("/out_override"), False, "inc_override", "sha_override" + ) + + assert stage1.execute_call_count == 1 + assert asset_rule.name in results['processed'] + + mock_mkdtemp.assert_not_called() # mkdtemp should not be called due to override + + # The orchestrator should create its source-rule specific subdir within the override + expected_source_rule_temp_dir_in_override = Path(override_temp_path_str) / source_rule.id.hex + + # Verify the context passed to the stage uses the overridden path structure + captured_context = stage1.contexts_called_with[0] + expected_asset_temp_dir_in_override = expected_source_rule_temp_dir_in_override / asset_rule.id.hex + assert captured_context.engine_temp_dir == expected_asset_temp_dir_in_override + + # rmtree should be called on the source_rule's directory within the override path + mock_rmtree.assert_called_once_with(expected_source_rule_temp_dir_in_override, ignore_errors=True) + +@mock.patch('shutil.rmtree') +@mock.patch('tempfile.mkdtemp') +def test_orchestrator_disabled_asset_rule_is_skipped(mock_mkdtemp, mock_rmtree): + mock_mkdtemp.return_value = "/fake/engine_temp_dir_disabled_asset" + + config = create_orchestrator_test_config() + stage1 = MockPassThroughStage("stage_disabled_check") + orchestrator = PipelineOrchestrator(config_obj=config, stages=[stage1]) + + source_rule = create_orchestrator_test_source_rule("DisabledAssetSourceRule", asset_names=["EnabledAsset", "DisabledAsset"]) + enabled_asset = source_rule.assets[0] + disabled_asset = source_rule.assets[1] + disabled_asset.enabled = False # Disable this asset rule + + results = orchestrator.process_source_rule( + source_rule, Path("/ws_disabled"), Path("/out_disabled"), False, "inc_dis", "sha_dis" + ) + + assert stage1.execute_call_count == 1 # Only called for the enabled asset + + assert enabled_asset.name in results['processed'] + assert disabled_asset.name in results['skipped'] + assert not results['failed'] + + # Verify context for the processed asset + assert stage1.contexts_called_with[0].asset_rule.name == enabled_asset.name + + # Verify skip reason for the disabled asset (this is set by the orchestrator itself) + # The orchestrator's _process_single_asset checks asset_rule.enabled + # We need to inspect the results dictionary for the skip reason if it's stored there, + # or infer it. The current structure of `results` doesn't store detailed skip reasons directly, + # but the test ensures it's in the 'skipped' list. + # For a more detailed check, one might need to adjust how results are reported or mock deeper. + # For now, confirming it's in 'skipped' and stage1 wasn't called for it is sufficient. + + mock_mkdtemp.assert_called_once() + expected_temp_path = Path(mock_mkdtemp.return_value) / source_rule.id.hex + mock_rmtree.assert_called_once_with(expected_temp_path, ignore_errors=True) \ No newline at end of file diff --git a/tests/processing/utils/test_image_processing_utils.py b/tests/processing/utils/test_image_processing_utils.py new file mode 100644 index 0000000..e128b3f --- /dev/null +++ b/tests/processing/utils/test_image_processing_utils.py @@ -0,0 +1,504 @@ +import pytest +from unittest import mock +import numpy as np +from pathlib import Path +import sys + +# Attempt to import the module under test +# This assumes that the 'tests' directory is at the same level as the 'processing' directory, +# and pytest handles the PYTHONPATH correctly. +try: + from processing.utils import image_processing_utils as ipu + import cv2 # Import cv2 here if it's used for constants like cv2.COLOR_BGR2RGB +except ImportError: + # Fallback for environments where PYTHONPATH might not be set up as expected by pytest initially + # This adds the project root to sys.path to find the 'processing' module + # Adjust the number of Path.parent calls if your test structure is deeper or shallower + project_root = Path(__file__).parent.parent.parent.parent + sys.path.insert(0, str(project_root)) + from processing.utils import image_processing_utils as ipu + import cv2 # Import cv2 here as well + +# If cv2 is imported directly in image_processing_utils, you might need to mock it globally for some tests +# For example, at the top of the test file: +# sys.modules['cv2'] = mock.MagicMock() # Basic global mock if needed +# We will use more targeted mocks with @mock.patch where cv2 is used. + +# --- Tests for Mathematical Helpers --- + +def test_is_power_of_two(): + assert ipu.is_power_of_two(1) is True + assert ipu.is_power_of_two(2) is True + assert ipu.is_power_of_two(4) is True + assert ipu.is_power_of_two(16) is True + assert ipu.is_power_of_two(1024) is True + assert ipu.is_power_of_two(0) is False + assert ipu.is_power_of_two(-2) is False + assert ipu.is_power_of_two(3) is False + assert ipu.is_power_of_two(100) is False + +def test_get_nearest_pot(): + assert ipu.get_nearest_pot(1) == 1 + assert ipu.get_nearest_pot(2) == 2 + # Based on current implementation: + # For 3: lower=2, upper=4. (3-2)=1, (4-3)=1. Else branch returns upper_pot. So 4. + assert ipu.get_nearest_pot(3) == 4 + assert ipu.get_nearest_pot(50) == 64 # (50-32)=18, (64-50)=14 -> upper + assert ipu.get_nearest_pot(100) == 128 # (100-64)=36, (128-100)=28 -> upper + assert ipu.get_nearest_pot(256) == 256 + assert ipu.get_nearest_pot(0) == 1 + assert ipu.get_nearest_pot(-10) == 1 + # For 700: value.bit_length() = 10. lower_pot = 1<<(10-1) = 512. upper_pot = 1<<10 = 1024. + # (700-512) = 188. (1024-700) = 324. (188 < 324) is True. Returns lower_pot. So 512. + assert ipu.get_nearest_pot(700) == 512 + assert ipu.get_nearest_pot(6) == 8 # (6-4)=2, (8-6)=2. Returns upper. + assert ipu.get_nearest_pot(5) == 4 # (5-4)=1, (8-5)=3. Returns lower. + + +@pytest.mark.parametrize( + "orig_w, orig_h, target_w, target_h, resize_mode, ensure_pot, allow_upscale, target_max_dim, expected_w, expected_h", + [ + # FIT mode + (1000, 800, 500, None, "fit", False, False, None, 500, 400), # Fit width + (1000, 800, None, 400, "fit", False, False, None, 500, 400), # Fit height + (1000, 800, 500, 500, "fit", False, False, None, 500, 400), # Fit to box (width constrained) + (800, 1000, 500, 500, "fit", False, False, None, 400, 500), # Fit to box (height constrained) + (100, 80, 200, None, "fit", False, False, None, 100, 80), # Fit width, no upscale + (100, 80, 200, None, "fit", False, True, None, 200, 160), # Fit width, allow upscale + (100, 80, 128, None, "fit", True, False, None, 128, 64), # Re-evaluated + (100, 80, 128, None, "fit", True, True, None, 128, 128), # Fit width, ensure_pot, allow upscale (128, 102 -> pot 128, 128) + + # STRETCH mode + (1000, 800, 500, 400, "stretch", False, False, None, 500, 400), + (100, 80, 200, 160, "stretch", False, True, None, 200, 160), # Stretch, allow upscale + (100, 80, 200, 160, "stretch", False, False, None, 100, 80), # Stretch, no upscale + (100, 80, 128, 128, "stretch", True, True, None, 128, 128), # Stretch, ensure_pot, allow upscale + (100, 80, 70, 70, "stretch", True, False, None, 64, 64), # Stretch, ensure_pot, no upscale (70,70 -> pot 64,64) + + # MAX_DIM_POT mode + (1000, 800, None, None, "max_dim_pot", True, False, 512, 512, 512), + (800, 1000, None, None, "max_dim_pot", True, False, 512, 512, 512), + (1920, 1080, None, None, "max_dim_pot", True, False, 1024, 1024, 512), + (100, 100, None, None, "max_dim_pot", True, False, 60, 64, 64), + # Edge cases for calculate_target_dimensions + (0, 0, 512, 512, "fit", False, False, None, 512, 512), + (10, 10, 512, 512, "fit", True, False, None, 8, 8), + (100, 100, 150, 150, "fit", True, False, None, 128, 128), + ] +) +def test_calculate_target_dimensions(orig_w, orig_h, target_w, target_h, resize_mode, ensure_pot, allow_upscale, target_max_dim, expected_w, expected_h): + if resize_mode == "max_dim_pot" and target_max_dim is None: + with pytest.raises(ValueError, match="target_max_dim_for_pot_mode must be provided"): + ipu.calculate_target_dimensions(orig_w, orig_h, target_width=target_w, target_height=target_h, + resize_mode=resize_mode, ensure_pot=ensure_pot, allow_upscale=allow_upscale, + target_max_dim_for_pot_mode=target_max_dim) + elif (resize_mode == "fit" and target_w is None and target_h is None) or \ + (resize_mode == "stretch" and (target_w is None or target_h is None)): + with pytest.raises(ValueError): + ipu.calculate_target_dimensions(orig_w, orig_h, target_width=target_w, target_height=target_h, + resize_mode=resize_mode, ensure_pot=ensure_pot, allow_upscale=allow_upscale, + target_max_dim_for_pot_mode=target_max_dim) + else: + actual_w, actual_h = ipu.calculate_target_dimensions( + orig_w, orig_h, target_width=target_w, target_height=target_h, + resize_mode=resize_mode, ensure_pot=ensure_pot, allow_upscale=allow_upscale, + target_max_dim_for_pot_mode=target_max_dim + ) + assert (actual_w, actual_h) == (expected_w, expected_h), \ + f"Input: ({orig_w},{orig_h}), T=({target_w},{target_h}), M={resize_mode}, POT={ensure_pot}, UPSC={allow_upscale}, TMAX={target_max_dim}" + + +def test_calculate_target_dimensions_invalid_mode(): + with pytest.raises(ValueError, match="Unsupported resize_mode"): + ipu.calculate_target_dimensions(100, 100, 50, 50, resize_mode="invalid_mode") + +@pytest.mark.parametrize( + "ow, oh, rw, rh, expected_str", + [ + (100, 100, 100, 100, "EVEN"), + (100, 100, 200, 200, "EVEN"), + (200, 200, 100, 100, "EVEN"), + (100, 100, 150, 100, "X15Y1"), + (100, 100, 50, 100, "X05Y1"), + (100, 100, 100, 150, "X1Y15"), + (100, 100, 100, 50, "X1Y05"), + (100, 50, 150, 75, "EVEN"), + (100, 50, 150, 50, "X15Y1"), + (100, 50, 100, 75, "X1Y15"), + (100, 50, 120, 60, "EVEN"), + (100, 50, 133, 66, "EVEN"), + (100, 100, 133, 100, "X133Y1"), + (100, 100, 100, 133, "X1Y133"), + (100, 100, 133, 133, "EVEN"), + (100, 100, 67, 100, "X067Y1"), + (100, 100, 100, 67, "X1Y067"), + (100, 100, 67, 67, "EVEN"), + (1920, 1080, 1024, 576, "EVEN"), + (1920, 1080, 1024, 512, "X112Y1"), + (0, 100, 50, 50, "InvalidInput"), + (100, 0, 50, 50, "InvalidInput"), + (100, 100, 0, 50, "InvalidResize"), + (100, 100, 50, 0, "InvalidResize"), + ] +) +def test_normalize_aspect_ratio_change(ow, oh, rw, rh, expected_str): + assert ipu.normalize_aspect_ratio_change(ow, oh, rw, rh) == expected_str + +# --- Tests for Image Manipulation --- + +@mock.patch('cv2.imread') +def test_load_image_success_str_path(mock_cv2_imread): + mock_img_data = np.array([[[1, 2, 3]]], dtype=np.uint8) + mock_cv2_imread.return_value = mock_img_data + + result = ipu.load_image("dummy/path.png") + + mock_cv2_imread.assert_called_once_with("dummy/path.png", cv2.IMREAD_UNCHANGED) + assert np.array_equal(result, mock_img_data) + +@mock.patch('cv2.imread') +def test_load_image_success_path_obj(mock_cv2_imread): + mock_img_data = np.array([[[1, 2, 3]]], dtype=np.uint8) + mock_cv2_imread.return_value = mock_img_data + dummy_path = Path("dummy/path.png") + + result = ipu.load_image(dummy_path) + + mock_cv2_imread.assert_called_once_with(str(dummy_path), cv2.IMREAD_UNCHANGED) + assert np.array_equal(result, mock_img_data) + +@mock.patch('cv2.imread') +def test_load_image_failure(mock_cv2_imread): + mock_cv2_imread.return_value = None + + result = ipu.load_image("dummy/path.png") + + mock_cv2_imread.assert_called_once_with("dummy/path.png", cv2.IMREAD_UNCHANGED) + assert result is None + +@mock.patch('cv2.imread', side_effect=Exception("CV2 Read Error")) +def test_load_image_exception(mock_cv2_imread): + result = ipu.load_image("dummy/path.png") + mock_cv2_imread.assert_called_once_with("dummy/path.png", cv2.IMREAD_UNCHANGED) + assert result is None + + +@mock.patch('cv2.cvtColor') +def test_convert_bgr_to_rgb_3_channel(mock_cv2_cvtcolor): + bgr_image = np.random.randint(0, 255, (10, 10, 3), dtype=np.uint8) + rgb_image_mock = np.random.randint(0, 255, (10, 10, 3), dtype=np.uint8) + mock_cv2_cvtcolor.return_value = rgb_image_mock + + result = ipu.convert_bgr_to_rgb(bgr_image) + + mock_cv2_cvtcolor.assert_called_once_with(bgr_image, cv2.COLOR_BGR2RGB) + assert np.array_equal(result, rgb_image_mock) + +@mock.patch('cv2.cvtColor') +def test_convert_bgr_to_rgb_4_channel_bgra(mock_cv2_cvtcolor): + bgra_image = np.random.randint(0, 255, (10, 10, 4), dtype=np.uint8) + rgb_image_mock = np.random.randint(0, 255, (10, 10, 3), dtype=np.uint8) # cvtColor BGRA2RGB drops alpha + mock_cv2_cvtcolor.return_value = rgb_image_mock # Mocking the output of BGRA2RGB + + result = ipu.convert_bgr_to_rgb(bgra_image) + + mock_cv2_cvtcolor.assert_called_once_with(bgra_image, cv2.COLOR_BGRA2RGB) + assert np.array_equal(result, rgb_image_mock) + + +def test_convert_bgr_to_rgb_none_input(): + assert ipu.convert_bgr_to_rgb(None) is None + +def test_convert_bgr_to_rgb_grayscale_input(): + gray_image = np.random.randint(0, 255, (10, 10), dtype=np.uint8) + result = ipu.convert_bgr_to_rgb(gray_image) + assert np.array_equal(result, gray_image) # Should return as is + +@mock.patch('cv2.cvtColor') +def test_convert_rgb_to_bgr_3_channel(mock_cv2_cvtcolor): + rgb_image = np.random.randint(0, 255, (10, 10, 3), dtype=np.uint8) + bgr_image_mock = np.random.randint(0, 255, (10, 10, 3), dtype=np.uint8) + mock_cv2_cvtcolor.return_value = bgr_image_mock + + result = ipu.convert_rgb_to_bgr(rgb_image) + + mock_cv2_cvtcolor.assert_called_once_with(rgb_image, cv2.COLOR_RGB2BGR) + assert np.array_equal(result, bgr_image_mock) + +def test_convert_rgb_to_bgr_none_input(): + assert ipu.convert_rgb_to_bgr(None) is None + +def test_convert_rgb_to_bgr_grayscale_input(): + gray_image = np.random.randint(0, 255, (10, 10), dtype=np.uint8) + result = ipu.convert_rgb_to_bgr(gray_image) + assert np.array_equal(result, gray_image) # Should return as is + +def test_convert_rgb_to_bgr_4_channel_input(): + rgba_image = np.random.randint(0, 255, (10, 10, 4), dtype=np.uint8) + result = ipu.convert_rgb_to_bgr(rgba_image) + assert np.array_equal(result, rgba_image) # Should return as is + + +@mock.patch('cv2.resize') +def test_resize_image_downscale(mock_cv2_resize): + original_image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8) + resized_image_mock = np.random.randint(0, 255, (50, 50, 3), dtype=np.uint8) + mock_cv2_resize.return_value = resized_image_mock + target_w, target_h = 50, 50 + + result = ipu.resize_image(original_image, target_w, target_h) + + mock_cv2_resize.assert_called_once_with(original_image, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4) + assert np.array_equal(result, resized_image_mock) + +@mock.patch('cv2.resize') +def test_resize_image_upscale(mock_cv2_resize): + original_image = np.random.randint(0, 255, (50, 50, 3), dtype=np.uint8) + resized_image_mock = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8) + mock_cv2_resize.return_value = resized_image_mock + target_w, target_h = 100, 100 + + result = ipu.resize_image(original_image, target_w, target_h) + + mock_cv2_resize.assert_called_once_with(original_image, (target_w, target_h), interpolation=cv2.INTER_CUBIC) + assert np.array_equal(result, resized_image_mock) + +@mock.patch('cv2.resize') +def test_resize_image_custom_interpolation(mock_cv2_resize): + original_image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8) + resized_image_mock = np.random.randint(0, 255, (50, 50, 3), dtype=np.uint8) + mock_cv2_resize.return_value = resized_image_mock + target_w, target_h = 50, 50 + + result = ipu.resize_image(original_image, target_w, target_h, interpolation=cv2.INTER_NEAREST) + + mock_cv2_resize.assert_called_once_with(original_image, (target_w, target_h), interpolation=cv2.INTER_NEAREST) + assert np.array_equal(result, resized_image_mock) + +def test_resize_image_none_input(): + with pytest.raises(ValueError, match="Cannot resize a None image."): + ipu.resize_image(None, 50, 50) + +@pytest.mark.parametrize("w, h", [(0, 50), (50, 0), (-1, 50)]) +def test_resize_image_invalid_dims(w, h): + original_image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8) + with pytest.raises(ValueError, match="Target width and height must be positive."): + ipu.resize_image(original_image, w, h) + + +@mock.patch('cv2.imwrite') +@mock.patch('pathlib.Path.mkdir') # Mock mkdir to avoid actual directory creation +def test_save_image_success(mock_mkdir, mock_cv2_imwrite): + mock_cv2_imwrite.return_value = True + img_data = np.zeros((10,10,3), dtype=np.uint8) # RGB + save_path = "output/test.png" + + # ipu.save_image converts RGB to BGR by default for non-EXR + # So we expect convert_rgb_to_bgr to be called internally, + # and cv2.imwrite to receive BGR data. + # We can mock convert_rgb_to_bgr if we want to be very specific, + # or trust its own unit tests and check the data passed to imwrite. + # For simplicity, let's assume convert_rgb_to_bgr works and imwrite gets BGR. + # The function copies data, so we can check the mock call. + + success = ipu.save_image(save_path, img_data, convert_to_bgr_before_save=True) + + assert success is True + mock_mkdir.assert_called_once_with(parents=True, exist_ok=True) + + # Check that imwrite was called. The first arg to assert_called_once_with is the path. + # The second arg is the image data. We need to compare it carefully. + # Since convert_rgb_to_bgr is called internally, the data passed to imwrite will be BGR. + # Let's create expected BGR data. + expected_bgr_data = cv2.cvtColor(img_data, cv2.COLOR_RGB2BGR) + + args, kwargs = mock_cv2_imwrite.call_args + assert args[0] == str(Path(save_path)) + assert np.array_equal(args[1], expected_bgr_data) + + +@mock.patch('cv2.imwrite') +@mock.patch('pathlib.Path.mkdir') +def test_save_image_success_exr_no_bgr_conversion(mock_mkdir, mock_cv2_imwrite): + mock_cv2_imwrite.return_value = True + img_data_rgb_float = np.random.rand(10,10,3).astype(np.float32) # RGB float for EXR + save_path = "output/test.exr" + + success = ipu.save_image(save_path, img_data_rgb_float, output_format="exr", convert_to_bgr_before_save=False) + + assert success is True + mock_mkdir.assert_called_once_with(parents=True, exist_ok=True) + args, kwargs = mock_cv2_imwrite.call_args + assert args[0] == str(Path(save_path)) + assert np.array_equal(args[1], img_data_rgb_float) # Should be original RGB data + +@mock.patch('cv2.imwrite') +@mock.patch('pathlib.Path.mkdir') +def test_save_image_success_explicit_bgr_false_png(mock_mkdir, mock_cv2_imwrite): + mock_cv2_imwrite.return_value = True + img_data_rgb = np.zeros((10,10,3), dtype=np.uint8) # RGB + save_path = "output/test.png" + + # If convert_to_bgr_before_save is False, it should save RGB as is. + # However, OpenCV's imwrite for PNG might still expect BGR. + # The function's docstring says: "If True and image is 3-channel, converts RGB to BGR." + # So if False, it passes the data as is. + success = ipu.save_image(save_path, img_data_rgb, convert_to_bgr_before_save=False) + + assert success is True + mock_mkdir.assert_called_once_with(parents=True, exist_ok=True) + args, kwargs = mock_cv2_imwrite.call_args + assert args[0] == str(Path(save_path)) + assert np.array_equal(args[1], img_data_rgb) + + +@mock.patch('cv2.imwrite') +@mock.patch('pathlib.Path.mkdir') +def test_save_image_failure(mock_mkdir, mock_cv2_imwrite): + mock_cv2_imwrite.return_value = False + img_data = np.zeros((10,10,3), dtype=np.uint8) + save_path = "output/fail.png" + + success = ipu.save_image(save_path, img_data) + + assert success is False + mock_mkdir.assert_called_once_with(parents=True, exist_ok=True) + mock_cv2_imwrite.assert_called_once() # Check it was called + +def test_save_image_none_data(): + assert ipu.save_image("output/none.png", None) is False + +@mock.patch('cv2.imwrite', side_effect=Exception("CV2 Write Error")) +@mock.patch('pathlib.Path.mkdir') +def test_save_image_exception(mock_mkdir, mock_cv2_imwrite_exception): + img_data = np.zeros((10,10,3), dtype=np.uint8) + save_path = "output/exception.png" + + success = ipu.save_image(save_path, img_data) + + assert success is False + mock_mkdir.assert_called_once_with(parents=True, exist_ok=True) + mock_cv2_imwrite_exception.assert_called_once() + +# Test data type conversions in save_image +@pytest.mark.parametrize( + "input_dtype, input_data_producer, output_dtype_target, expected_conversion_dtype, check_scaling", + [ + (np.uint16, lambda: (np.random.randint(0, 65535, (10,10,3), dtype=np.uint16)), np.uint8, np.uint8, True), + (np.float32, lambda: np.random.rand(10,10,3).astype(np.float32), np.uint8, np.uint8, True), + (np.uint8, lambda: (np.random.randint(0, 255, (10,10,3), dtype=np.uint8)), np.uint16, np.uint16, True), + (np.float32, lambda: np.random.rand(10,10,3).astype(np.float32), np.uint16, np.uint16, True), + (np.uint8, lambda: (np.random.randint(0, 255, (10,10,3), dtype=np.uint8)), np.float16, np.float16, True), + (np.uint16, lambda: (np.random.randint(0, 65535, (10,10,3), dtype=np.uint16)), np.float32, np.float32, True), + ] +) +@mock.patch('cv2.imwrite') +@mock.patch('pathlib.Path.mkdir') +def test_save_image_dtype_conversion(mock_mkdir, mock_cv2_imwrite, input_dtype, input_data_producer, output_dtype_target, expected_conversion_dtype, check_scaling): + mock_cv2_imwrite.return_value = True + img_data = input_data_producer() + original_img_data_copy = img_data.copy() # For checking scaling if needed + + ipu.save_image("output/dtype_test.png", img_data, output_dtype_target=output_dtype_target) + + mock_cv2_imwrite.assert_called_once() + saved_img_data = mock_cv2_imwrite.call_args[0][1] # Get the image data passed to imwrite + + assert saved_img_data.dtype == expected_conversion_dtype + + if check_scaling: + # This is a basic check. More precise checks would require known input/output values. + if output_dtype_target == np.uint8: + if input_dtype == np.uint16: + expected_scaled_data = (original_img_data_copy.astype(np.float32) / 65535.0 * 255.0).astype(np.uint8) + assert np.allclose(saved_img_data, cv2.cvtColor(expected_scaled_data, cv2.COLOR_RGB2BGR), atol=1) # Allow small diff due to float precision + elif input_dtype in [np.float16, np.float32, np.float64]: + expected_scaled_data = (np.clip(original_img_data_copy, 0.0, 1.0) * 255.0).astype(np.uint8) + assert np.allclose(saved_img_data, cv2.cvtColor(expected_scaled_data, cv2.COLOR_RGB2BGR), atol=1) + elif output_dtype_target == np.uint16: + if input_dtype == np.uint8: + expected_scaled_data = (original_img_data_copy.astype(np.float32) / 255.0 * 65535.0).astype(np.uint16) + assert np.allclose(saved_img_data, cv2.cvtColor(expected_scaled_data, cv2.COLOR_RGB2BGR), atol=1) + elif input_dtype in [np.float16, np.float32, np.float64]: + expected_scaled_data = (np.clip(original_img_data_copy, 0.0, 1.0) * 65535.0).astype(np.uint16) + assert np.allclose(saved_img_data, cv2.cvtColor(expected_scaled_data, cv2.COLOR_RGB2BGR), atol=1) + # Add more scaling checks for float16, float32 if necessary + + +# --- Tests for calculate_image_stats --- + +def test_calculate_image_stats_grayscale_uint8(): + img_data = np.array([[0, 128], [255, 10]], dtype=np.uint8) + # Expected normalized: [[0, 0.50196], [1.0, 0.03921]] approx + stats = ipu.calculate_image_stats(img_data) + assert stats is not None + assert np.isclose(stats["min"], 0/255.0) + assert np.isclose(stats["max"], 255/255.0) + assert np.isclose(stats["mean"], np.mean(img_data.astype(np.float64)/255.0)) + +def test_calculate_image_stats_color_uint8(): + img_data = np.array([ + [[0, 50, 100], [10, 60, 110]], + [[255, 128, 200], [20, 70, 120]] + ], dtype=np.uint8) + stats = ipu.calculate_image_stats(img_data) + assert stats is not None + # Min per channel (normalized) + assert np.allclose(stats["min"], [0/255.0, 50/255.0, 100/255.0]) + # Max per channel (normalized) + assert np.allclose(stats["max"], [255/255.0, 128/255.0, 200/255.0]) + # Mean per channel (normalized) + expected_mean = np.mean(img_data.astype(np.float64)/255.0, axis=(0,1)) + assert np.allclose(stats["mean"], expected_mean) + +def test_calculate_image_stats_grayscale_uint16(): + img_data = np.array([[0, 32768], [65535, 1000]], dtype=np.uint16) + stats = ipu.calculate_image_stats(img_data) + assert stats is not None + assert np.isclose(stats["min"], 0/65535.0) + assert np.isclose(stats["max"], 65535/65535.0) + assert np.isclose(stats["mean"], np.mean(img_data.astype(np.float64)/65535.0)) + +def test_calculate_image_stats_color_float32(): + # Floats are assumed to be in 0-1 range already by the function's normalization logic + img_data = np.array([ + [[0.0, 0.2, 0.4], [0.1, 0.3, 0.5]], + [[1.0, 0.5, 0.8], [0.05, 0.25, 0.6]] + ], dtype=np.float32) + stats = ipu.calculate_image_stats(img_data) + assert stats is not None + assert np.allclose(stats["min"], [0.0, 0.2, 0.4]) + assert np.allclose(stats["max"], [1.0, 0.5, 0.8]) + expected_mean = np.mean(img_data.astype(np.float64), axis=(0,1)) + assert np.allclose(stats["mean"], expected_mean) + +def test_calculate_image_stats_none_input(): + assert ipu.calculate_image_stats(None) is None + +def test_calculate_image_stats_unsupported_shape(): + img_data = np.zeros((2,2,2,2), dtype=np.uint8) # 4D array + assert ipu.calculate_image_stats(img_data) is None + +@mock.patch('numpy.mean', side_effect=Exception("Numpy error")) +def test_calculate_image_stats_exception_during_calculation(mock_np_mean): + img_data = np.array([[0, 128], [255, 10]], dtype=np.uint8) + stats = ipu.calculate_image_stats(img_data) + assert stats == {"error": "Error calculating image stats"} + +# Example of mocking ipu.load_image for a function that uses it (if calculate_image_stats used it) +# For the current calculate_image_stats, it takes image_data directly, so this is not needed for it. +# This is just an example as requested in the prompt for a hypothetical scenario. +@mock.patch('processing.utils.image_processing_utils.load_image') +def test_hypothetical_function_using_load_image(mock_load_image): + # This test is for a function that would call ipu.load_image internally + # e.g. def process_image_from_path(path): + # img_data = ipu.load_image(path) + # return ipu.calculate_image_stats(img_data) + + mock_img_data = np.array([[[0.5]]], dtype=np.float32) + mock_load_image.return_value = mock_img_data + + # result = ipu.hypothetical_process_image_from_path("dummy.png") + # mock_load_image.assert_called_once_with("dummy.png") + # assert result["mean"] == 0.5 + pass # This is a conceptual example \ No newline at end of file diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py new file mode 100644 index 0000000..cfc5ffa --- /dev/null +++ b/tests/utils/__init__.py @@ -0,0 +1 @@ +# This file makes the 'tests/utils' directory a Python package. \ No newline at end of file diff --git a/tests/utils/test_path_utils.py b/tests/utils/test_path_utils.py new file mode 100644 index 0000000..56f7a0a --- /dev/null +++ b/tests/utils/test_path_utils.py @@ -0,0 +1,252 @@ +import pytest +from pathlib import Path +from utils.path_utils import sanitize_filename, generate_path_from_pattern + +# Tests for sanitize_filename +def test_sanitize_filename_valid(): + assert sanitize_filename("valid_filename.txt") == "valid_filename.txt" + +def test_sanitize_filename_with_spaces(): + assert sanitize_filename("file name with spaces.txt") == "file_name_with_spaces.txt" + +def test_sanitize_filename_with_special_characters(): + assert sanitize_filename("file!@#$%^&*()[]{};:'\",.<>/?\\|.txt") == "file____________________.txt" + +def test_sanitize_filename_with_leading_trailing_whitespace(): + assert sanitize_filename(" filename_with_spaces .txt") == "filename_with_spaces.txt" + +def test_sanitize_filename_empty_string(): + assert sanitize_filename("") == "" + +def test_sanitize_filename_with_none(): + with pytest.raises(TypeError): + sanitize_filename(None) + +def test_sanitize_filename_mixed_case(): + assert sanitize_filename("MixedCaseFileName.PNG") == "MixedCaseFileName.PNG" + +def test_sanitize_filename_long_filename(): + long_name = "a" * 255 + ".txt" + # Assuming the function doesn't truncate, but sanitizes. + # If it's meant to handle OS limits, this test might need adjustment + # based on the function's specific behavior for long names. + assert sanitize_filename(long_name) == long_name + +def test_sanitize_filename_unicode_characters(): + assert sanitize_filename("文件名前缀_文件名_后缀.jpg") == "文件名前缀_文件名_后缀.jpg" + +def test_sanitize_filename_multiple_extensions(): + assert sanitize_filename("archive.tar.gz") == "archive.tar.gz" + +def test_sanitize_filename_no_extension(): + assert sanitize_filename("filename") == "filename" + +def test_sanitize_filename_only_special_chars(): + assert sanitize_filename("!@#$%^") == "______" + +def test_sanitize_filename_with_hyphens_and_underscores(): + assert sanitize_filename("file-name_with-hyphens_and_underscores.zip") == "file-name_with-hyphens_and_underscores.zip" + +# Tests for generate_path_from_pattern +def test_generate_path_basic(): + result = generate_path_from_pattern( + base_path="output", + pattern="{asset_name}/{map_type}/{filename}", + asset_name="MyAsset", + map_type="Diffuse", + filename="MyAsset_Diffuse.png", + source_rule_name="TestRule", + incrementing_value=None, + sha5_value=None + ) + expected = Path("output/MyAsset/Diffuse/MyAsset_Diffuse.png") + assert Path(result) == expected + +def test_generate_path_all_placeholders(): + result = generate_path_from_pattern( + base_path="project_files", + pattern="{source_rule_name}/{asset_name}/{map_type}_{incrementing_value}_{sha5_value}/{filename}", + asset_name="AnotherAsset", + map_type="Normal", + filename="NormalMap.tif", + source_rule_name="ComplexRule", + incrementing_value="001", + sha5_value="abcde" + ) + expected = Path("project_files/ComplexRule/AnotherAsset/Normal_001_abcde/NormalMap.tif") + assert Path(result) == expected + +def test_generate_path_optional_placeholders_none(): + result = generate_path_from_pattern( + base_path="data", + pattern="{asset_name}/{filename}", + asset_name="SimpleAsset", + map_type="Albedo", # map_type is in pattern but not used if not in string + filename="texture.jpg", + source_rule_name="Basic", + incrementing_value=None, + sha5_value=None + ) + expected = Path("data/SimpleAsset/texture.jpg") + assert Path(result) == expected + +def test_generate_path_optional_incrementing_value_present(): + result = generate_path_from_pattern( + base_path="assets", + pattern="{asset_name}/{map_type}/v{incrementing_value}/{filename}", + asset_name="VersionedAsset", + map_type="Specular", + filename="spec.png", + source_rule_name="VersioningRule", + incrementing_value="3", + sha5_value=None + ) + expected = Path("assets/VersionedAsset/Specular/v3/spec.png") + assert Path(result) == expected + +def test_generate_path_optional_sha5_value_present(): + result = generate_path_from_pattern( + base_path="cache", + pattern="{asset_name}/{sha5_value}/{filename}", + asset_name="HashedAsset", + map_type="Roughness", + filename="rough.exr", + source_rule_name="HashingRule", + incrementing_value=None, + sha5_value="f1234" + ) + expected = Path("cache/HashedAsset/f1234/rough.exr") + assert Path(result) == expected + +def test_generate_path_base_path_is_path_object(): + result = generate_path_from_pattern( + base_path=Path("output_path"), + pattern="{asset_name}/{filename}", + asset_name="ObjectAsset", + map_type="AO", + filename="ao.png", + source_rule_name="PathObjectRule", + incrementing_value=None, + sha5_value=None + ) + expected = Path("output_path/ObjectAsset/ao.png") + assert Path(result) == expected + +def test_generate_path_empty_pattern(): + result = generate_path_from_pattern( + base_path="output", + pattern="", # Empty pattern should just use base_path and filename + asset_name="MyAsset", + map_type="Diffuse", + filename="MyAsset_Diffuse.png", + source_rule_name="TestRule", + incrementing_value=None, + sha5_value=None + ) + expected = Path("output/MyAsset_Diffuse.png") + assert Path(result) == expected + +def test_generate_path_pattern_with_no_placeholders(): + result = generate_path_from_pattern( + base_path="fixed_output", + pattern="some/static/path", # Pattern has no placeholders + asset_name="MyAsset", + map_type="Diffuse", + filename="MyAsset_Diffuse.png", + source_rule_name="TestRule", + incrementing_value=None, + sha5_value=None + ) + expected = Path("fixed_output/some/static/path/MyAsset_Diffuse.png") + assert Path(result) == expected + +def test_generate_path_filename_with_subdirs_in_pattern(): + result = generate_path_from_pattern( + base_path="output", + pattern="{asset_name}", # Filename itself will be appended + asset_name="AssetWithSubdirFile", + map_type="Color", + filename="textures/variant1/color.png", # Filename contains subdirectories + source_rule_name="SubdirRule", + incrementing_value=None, + sha5_value=None + ) + # The function is expected to join pattern result with filename + expected = Path("output/AssetWithSubdirFile/textures/variant1/color.png") + assert Path(result) == expected + +def test_generate_path_no_filename_provided(): + # This test assumes that if filename is None or empty, it might raise an error + # or behave in a specific way, e.g. not append anything or use a default. + # Adjust based on actual function behavior for missing filename. + # For now, let's assume it might raise TypeError if filename is critical. + with pytest.raises(TypeError): # Or ValueError, depending on implementation + generate_path_from_pattern( + base_path="output", + pattern="{asset_name}/{map_type}", + asset_name="MyAsset", + map_type="Diffuse", + filename=None, # No filename + source_rule_name="TestRule", + incrementing_value=None, + sha5_value=None + ) + +def test_generate_path_all_values_are_empty_strings_or_none_where_applicable(): + result = generate_path_from_pattern( + base_path="", # Empty base_path + pattern="{asset_name}/{map_type}/{incrementing_value}/{sha5_value}", + asset_name="", # Empty asset_name + map_type="", # Empty map_type + filename="empty_test.file", + source_rule_name="", # Empty source_rule_name + incrementing_value="", # Empty incrementing_value + sha5_value="" # Empty sha5_value + ) + # Behavior with empty strings might vary. Assuming they are treated as literal empty segments. + # Path("///empty_test.file") might resolve to "/empty_test.file" on POSIX + # or just "empty_test.file" if base_path is current dir. + # Let's assume Path() handles normalization. + # If base_path is "", it means current directory. + # So, "//empty_test.file" relative to current dir. + # Path objects normalize this. e.g. Path('//a') -> Path('/a') on POSIX + # Path('a//b') -> Path('a/b') + # Path('/a//b') -> Path('/a/b') + # Path('//a//b') -> Path('/a/b') + # If base_path is empty, it's like Path('.////empty_test.file') + expected = Path("empty_test.file") # Simplified, actual result might be OS dependent or Path lib norm. + # More robust check: + # result_path = Path(result) + # expected_path = Path.cwd() / "" / "" / "" / "" / "empty_test.file" # This is not quite right + # Let's assume the function joins them: "" + "/" + "" + "/" + "" + "/" + "" + "/" + "empty_test.file" + # which becomes "////empty_test.file" + # Path("////empty_test.file") on Windows becomes "\\empty_test.file" (network path attempt) + # Path("////empty_test.file") on Linux becomes "/empty_test.file" + # Given the function likely uses os.path.join or Path.joinpath, + # and base_path="", asset_name="", map_type="", inc_val="", sha5_val="" + # pattern = "{asset_name}/{map_type}/{incrementing_value}/{sha5_value}" -> "///" + # result = base_path / pattern_result / filename + # result = "" / "///" / "empty_test.file" + # Path("") / "///" / "empty_test.file" -> Path("///empty_test.file") + # This is tricky. Let's assume the function is robust. + # If all path segments are empty, it should ideally resolve to just the filename relative to base_path. + # If base_path is also empty, then filename relative to CWD. + # Let's test the expected output based on typical os.path.join behavior: + # os.path.join("", "", "", "", "", "empty_test.file") -> "empty_test.file" on Windows + # os.path.join("", "", "", "", "", "empty_test.file") -> "empty_test.file" on Linux + assert Path(result) == Path("empty_test.file") + + +def test_generate_path_with_dots_in_placeholders(): + result = generate_path_from_pattern( + base_path="output", + pattern="{asset_name}/{map_type}", + asset_name="My.Asset.V1", + map_type="Diffuse.Main", + filename="texture.png", + source_rule_name="DotsRule", + incrementing_value=None, + sha5_value=None + ) + expected = Path("output/My.Asset.V1/Diffuse.Main/texture.png") + assert Path(result) == expected \ No newline at end of file diff --git a/utils/path_utils.py b/utils/path_utils.py index e20d3c2..b67929f 100644 --- a/utils/path_utils.py +++ b/utils/path_utils.py @@ -154,6 +154,15 @@ def get_next_incrementing_value(output_base_path: Path, output_directory_pattern logger.info(f"Determined next incrementing value: {next_value_str} (Max found: {max_value})") return next_value_str +def sanitize_filename(name: str) -> str: + """Removes or replaces characters invalid for filenames/directory names.""" + if not isinstance(name, str): name = str(name) + name = re.sub(r'[^\w.\-]+', '_', name) # Allow alphanumeric, underscore, hyphen, dot + name = re.sub(r'_+', '_', name) + name = name.strip('_') + if not name: name = "invalid_name" + return name + # --- Basic Unit Tests --- if __name__ == "__main__": print("Running basic tests for path_utils.generate_path_from_pattern...")