#!/Volumes/SAM2/CODE/pynamer/.venv/bin/python import argparse import base64 import os import sys from pathlib import Path import yaml from typing import Dict, List, Optional, Union import litellm from litellm import completion import logging # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger('pynamer') class PyNamer: """A tool to generate descriptive filenames for images using LLMs.""" def __init__(self, config_path: str = 'config.yaml'): """Initialize the PyNamer with configuration. Args: config_path: Path to the YAML configuration file """ self.config = self._load_config(config_path) self._setup_llm() def _load_config(self, config_path: str) -> Dict: """Load configuration from YAML file. Args: config_path: Path to the configuration file Returns: Dict containing configuration """ try: with open(config_path, 'r') as f: config = yaml.safe_load(f) logger.info(f"Loaded configuration from {config_path}") return config except Exception as e: logger.error(f"Failed to load configuration: {e}") sys.exit(1) def _setup_llm(self) -> None: """Set up the LLM client based on configuration.""" llm_config = self.config.get('llm', {}) # Set API key if provided in config api_key = llm_config.get('api_key') if api_key: os.environ["OPENAI_API_KEY"] = api_key # Set custom endpoint if provided endpoint = llm_config.get('endpoint') if endpoint: os.environ["OPENAI_API_BASE"] = endpoint self.model = llm_config.get('model', 'gpt-4-vision-preview') self.max_tokens = llm_config.get('max_tokens', 100) self.temperature = llm_config.get('temperature', 0.7) logger.info(f"LLM setup complete. Using model: {self.model}") def _encode_image(self, image_path: str) -> str: """Encode image to base64 for API submission. Args: image_path: Path to the image file Returns: Base64 encoded image string """ with open(image_path, "rb") as image_file: return base64.b64encode(image_file.read()).decode('utf-8') def _is_supported_format(self, file_path: str) -> bool: """Check if the file format is supported. Args: file_path: Path to the file Returns: True if supported, False otherwise """ supported_formats = self.config.get('image', {}).get('supported_formats', []) file_ext = os.path.splitext(file_path)[1].lower() return file_ext in supported_formats def generate_filename(self, image_path: str) -> str: """Generate a descriptive filename for the image using LLM. Args: image_path: Path to the image file Returns: Generated filename (without extension) """ if not os.path.exists(image_path): logger.error(f"Image not found: {image_path}") return None if not self._is_supported_format(image_path): logger.error(f"Unsupported file format: {image_path}") return None try: # Encode image base64_image = self._encode_image(image_path) # Prepare messages for LLM system_message = self.config.get('prompt', {}).get('system_message', '') user_message = self.config.get('prompt', {}).get('user_message', '') messages = [ {"role": "system", "content": system_message}, { "role": "user", "content": [ {"type": "text", "text": user_message}, { "type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"} } ] } ] # Call LLM response = completion( model=self.model, messages=messages, max_tokens=self.max_tokens, temperature=self.temperature ) # Extract filename from response filename = response.choices[0].message.content.strip() logger.info(f"Generated filename: {filename}") return filename except Exception as e: logger.error(f"Error generating filename: {e}") return None def rename_image(self, image_path: str, dry_run: bool = False) -> Optional[str]: """Rename the image with a generated descriptive filename. Args: image_path: Path to the image file dry_run: If True, don't actually rename the file Returns: New path if successful, None otherwise """ # Generate filename new_filename = self.generate_filename(image_path) if not new_filename: return None # Clean up the filename (ensure snake_case, no special chars) new_filename = new_filename.lower() new_filename = ''.join(c if c.isalnum() else '_' for c in new_filename) new_filename = new_filename.replace('__', '_').strip('_') # Get original path components path = Path(image_path) directory = path.parent extension = path.suffix # Create new path new_path = directory / f"{new_filename}{extension}" # Rename file if not dry_run: try: # Handle case where the new filename already exists counter = 1 while new_path.exists(): new_path = directory / f"{new_filename}_{counter}{extension}" counter += 1 path.rename(new_path) logger.info(f"Renamed: {image_path} -> {new_path}") except Exception as e: logger.error(f"Error renaming file: {e}") return None else: logger.info(f"[DRY RUN] Would rename: {image_path} -> {new_path}") return str(new_path) def main(): """Main entry point for the script.""" parser = argparse.ArgumentParser(description='Generate descriptive filenames for images using LLMs') parser.add_argument('images', nargs='+', help='Paths to image files') parser.add_argument('-c', '--config', default='config.yaml', help='Path to configuration file') parser.add_argument('-d', '--dry-run', action='store_true', help='Preview changes without renaming files') parser.add_argument('-v', '--verbose', action='store_true', help='Enable verbose output') args = parser.parse_args() # Set logging level if args.verbose: logger.setLevel(logging.DEBUG) # Initialize PyNamer namer = PyNamer(config_path=args.config) # Process each image - handle image paths that might contain spaces image_paths = [] # First, try to handle the case where the entire argument list is a single path with spaces if len(args.images) > 1: combined_path = ' '.join(args.images) if os.path.exists(combined_path): image_paths = [combined_path] logger.info(f"Found image by combining all arguments: {combined_path}") # If that didn't work, try to handle each argument individually if not image_paths: # Use a set to avoid duplicate processing processed_paths = set() for path in args.images: # Check if the path exists as is if os.path.exists(path) and path not in processed_paths: image_paths.append(path) processed_paths.add(path) else: # Try to find files that match the pattern import glob matching_files = glob.glob(f"*{path}*") for file in matching_files: if file not in processed_paths: image_paths.append(file) processed_paths.add(file) if not matching_files: logger.debug(f"Could not find any file matching '{path}'") # Process each valid image path if not image_paths: print("Error: No valid image files found to process.") return # Remove duplicates while preserving order unique_image_paths = [] seen = set() for path in image_paths: if path not in seen: unique_image_paths.append(path) seen.add(path) # Process each image for image_path in unique_image_paths: if not os.path.exists(image_path): logger.warning(f"Skipping non-existent file: {image_path}") continue new_path = namer.rename_image(image_path, dry_run=args.dry_run) if new_path: print(f"{'[DRY RUN] ' if args.dry_run else ''}Renamed: {image_path} -> {new_path}") else: print(f"Failed to process: {image_path}") if __name__ == "__main__": main()