import logging import cv2 import numpy as np from pathlib import Path from typing import List, Dict, Any, Tuple, Optional # Potentially import ipu from ...utils import image_processing_utils as ipu # Assuming ipu is available in the same utils directory or parent try: from . import image_processing_utils as ipu except ImportError: # Fallback for different import structures if needed, adjust based on actual project structure # For this project structure, the relative import should work. logging.warning("Could not import image_processing_utils using relative path. Attempting absolute import.") try: from processing.utils import image_processing_utils as ipu except ImportError: logging.error("Could not import image_processing_utils.") ipu = None # Handle case where ipu is not available logger = logging.getLogger(__name__) def save_image_variants( source_image_data: np.ndarray, base_map_type: str, # Filename-friendly map type source_bit_depth_info: List[Optional[int]], image_resolutions: Dict[str, int], file_type_defs: Dict[str, Dict[str, Any]], output_format_8bit: str, output_format_16bit_primary: str, output_format_16bit_fallback: str, png_compression_level: int, jpg_quality: int, output_filename_pattern_tokens: Dict[str, Any], # Must include 'output_base_directory': Path and 'asset_name': str output_filename_pattern: str, # Consider adding ipu or relevant parts of it if not importing globally ) -> List[Dict[str, Any]]: """ Centralizes image saving logic, generating and saving various resolution variants according to configuration. Args: source_image_data (np.ndarray): High-res image data (in memory, potentially transformed). base_map_type (str): Final map type (e.g., "COL", "ROUGH", "NORMAL", "MAP_NRMRGH"). This is the filename-friendly map type. source_bit_depth_info (List[Optional[int]]): List of original source bit depth(s) (e.g., [8], [16], [8, 16]). Can contain None. image_resolutions (Dict[str, int]): Dictionary mapping resolution keys (e.g., "4K") to max dimensions (e.g., 4096). file_type_defs (Dict[str, Dict[str, Any]]): Dictionary defining properties for map types, including 'bit_depth_rule'. output_format_8bit (str): File extension for 8-bit output (e.g., "jpg", "png"). output_format_16bit_primary (str): Primary file extension for 16-bit output (e.g., "png", "tif"). output_format_16bit_fallback (str): Fallback file extension for 16-bit output. png_compression_level (int): Compression level for PNG output (0-9). jpg_quality (int): Quality level for JPG output (0-100). output_filename_pattern_tokens (Dict[str, Any]): Dictionary of tokens for filename pattern replacement. Must include 'output_base_directory' (Path) and 'asset_name' (str). output_filename_pattern (str): Pattern string for generating output filenames (e.g., "[assetname]_[maptype]_[resolution].[ext]"). Returns: List[Dict[str, Any]]: A list of dictionaries, each containing details about a saved file. Example: [{'path': str, 'resolution_key': str, 'format': str, 'bit_depth': int, 'dimensions': (w,h)}, ...] """ if ipu is None: logger.error("image_processing_utils is not available. Cannot save images.") return [] saved_file_details = [] source_h, source_w = source_image_data.shape[:2] source_max_dim = max(source_h, source_w) # 1. Use provided configuration inputs (already available as function arguments) logger.info(f"Saving variants for map type: {base_map_type}") # 2. Determine Target Bit Depth target_bit_depth = 8 # Default bit_depth_rule = file_type_defs.get(base_map_type, {}).get('bit_depth_rule', 'force_8bit') if bit_depth_rule not in ['force_8bit', 'respect_inputs']: logger.warning(f"Unknown bit_depth_rule '{bit_depth_rule}' for map type '{base_map_type}'. Defaulting to 'force_8bit'.") bit_depth_rule = 'force_8bit' if bit_depth_rule == 'respect_inputs': # Check if any source bit depth is > 8, ignoring None if any(depth is not None and depth > 8 for depth in source_bit_depth_info): target_bit_depth = 16 else: target_bit_depth = 8 logger.info(f"Bit depth rule 'respect_inputs' applied. Source bit depths: {source_bit_depth_info}. Target bit depth: {target_bit_depth}") else: # force_8bit target_bit_depth = 8 logger.info(f"Bit depth rule 'force_8bit' applied. Target bit depth: {target_bit_depth}") # 3. Determine Output File Format(s) if target_bit_depth == 8: output_ext = output_format_8bit.lstrip('.').lower() elif target_bit_depth == 16: # Prioritize primary, fallback to fallback if primary is not supported/desired # For now, just use primary. More complex logic might be needed later. output_ext = output_format_16bit_primary.lstrip('.').lower() # Basic fallback logic example (can be expanded) if output_ext not in ['png', 'tif']: # Assuming common 16-bit formats output_ext = output_format_16bit_fallback.lstrip('.').lower() logger.warning(f"Primary 16-bit format '{output_format_16bit_primary}' might not be suitable. Using fallback '{output_format_16bit_fallback}'.") else: logger.error(f"Unsupported target bit depth: {target_bit_depth}. Defaulting to 8-bit format.") output_ext = output_format_8bit.lstrip('.').lower() logger.info(f"Target bit depth: {target_bit_depth}, Output format: {output_ext}") # 4. Generate and Save Resolution Variants # Sort resolutions by max dimension descending sorted_resolutions = sorted(image_resolutions.items(), key=lambda item: item[1], reverse=True) for res_key, res_max_dim in sorted_resolutions: logger.info(f"Processing resolution variant: {res_key} ({res_max_dim} max dim)") # Calculate target dimensions, ensuring no upscaling if source_max_dim <= res_max_dim: # If source is smaller or equal, use source dimensions target_w_res, target_h_res = source_w, source_h if source_max_dim < res_max_dim: logger.info(f"Source image ({source_w}x{source_h}) is smaller than target resolution {res_key} ({res_max_dim}). Saving at source resolution.") else: # Downscale, maintaining aspect ratio aspect_ratio = source_w / source_h if source_w > source_h: target_w_res = res_max_dim target_h_res = int(res_max_dim / aspect_ratio) else: target_h_res = res_max_dim target_w_res = int(res_max_dim * aspect_ratio) logger.info(f"Resizing source image ({source_w}x{source_h}) to {target_w_res}x{target_h_res} for {res_key} variant.") # Resize source_image_data # Use INTER_AREA for downscaling, INTER_LINEAR or INTER_CUBIC for upscaling (though we avoid upscaling here) interpolation_method = cv2.INTER_AREA # Good for downscaling # If we were allowing upscaling, we might add logic like: # if target_w_res > source_w or target_h_res > source_h: # interpolation_method = cv2.INTER_LINEAR # Or INTER_CUBIC try: variant_data = ipu.resize_image(source_image_data, (target_w_res, target_h_res), interpolation=interpolation_method) logger.debug(f"Resized variant data shape: {variant_data.shape}") except Exception as e: logger.error(f"Error resizing image for {res_key} variant: {e}") continue # Skip this variant if resizing fails # Filename Construction current_tokens = output_filename_pattern_tokens.copy() current_tokens['maptype'] = base_map_type current_tokens['resolution'] = res_key current_tokens['ext'] = output_ext try: # Replace placeholders in the pattern filename = output_filename_pattern for token, value in current_tokens.items(): # Ensure value is string for replacement, handle Path objects later filename = filename.replace(f"[{token}]", str(value)) # Construct full output path output_base_directory = current_tokens.get('output_base_directory') if not isinstance(output_base_directory, Path): logger.error(f"'output_base_directory' token is missing or not a Path object: {output_base_directory}. Cannot save file.") continue # Skip this variant output_path = output_base_directory / filename logger.info(f"Constructed output path: {output_path}") # Ensure parent directory exists output_path.parent.mkdir(parents=True, exist_ok=True) logger.debug(f"Ensured directory exists: {output_path.parent}") except Exception as e: logger.error(f"Error constructing filepath for {res_key} variant: {e}") continue # Skip this variant if path construction fails # Prepare Save Parameters save_params_cv2 = [] if output_ext == 'jpg': save_params_cv2.append(cv2.IMWRITE_JPEG_QUALITY) save_params_cv2.append(jpg_quality) logger.debug(f"Using JPG quality: {jpg_quality}") elif output_ext == 'png': save_params_cv2.append(cv2.IMWRITE_PNG_COMPRESSION) save_params_cv2.append(png_compression_level) logger.debug(f"Using PNG compression level: {png_compression_level}") # Add other format specific parameters if needed (e.g., TIFF compression) # Bit Depth Conversion (just before saving) image_data_for_save = variant_data try: if target_bit_depth == 8: image_data_for_save = ipu.convert_to_uint8(variant_data) logger.debug("Converted variant data to uint8.") elif target_bit_depth == 16: # ipu.convert_to_uint16 might handle different input types (float, uint8) # Assuming variant_data might be float after resizing, convert to uint16 image_data_for_save = ipu.convert_to_uint16(variant_data) logger.debug("Converted variant data to uint16.") # Add other bit depth conversions if needed except Exception as e: logger.error(f"Error converting image data to target bit depth {target_bit_depth} for {res_key} variant: {e}") continue # Skip this variant if conversion fails # Saving try: # ipu.save_image is expected to handle the actual cv2.imwrite call success = ipu.save_image(str(output_path), image_data_for_save, params=save_params_cv2) if success: logger.info(f"Successfully saved {res_key} variant to {output_path}") # Collect details for the returned list saved_file_details.append({ 'path': str(output_path), 'resolution_key': res_key, 'format': output_ext, 'bit_depth': target_bit_depth, 'dimensions': (target_w_res, target_h_res) }) else: logger.error(f"Failed to save {res_key} variant to {output_path}") except Exception as e: logger.error(f"Error saving image for {res_key} variant to {output_path}: {e}") # Continue to next variant even if one fails # Discard in-memory variant after saving (Python's garbage collection handles this) del variant_data del image_data_for_save # 5. Return List of Saved File Details logger.info(f"Finished saving variants for map type: {base_map_type}. Saved {len(saved_file_details)} variants.") return saved_file_details # Optional Helper Functions (can be added here if needed) # def _determine_target_bit_depth(...): ... # def _determine_output_format(...): ... # def _construct_variant_filepath(...): ...