Spaces:
Sleeping
Sleeping
| """ | |
| Module for extracting metadata from image files, particularly focusing on | |
| Stable Diffusion metadata from PNG files. | |
| """ | |
| import io | |
| from PIL import Image, PngImagePlugin | |
| import re | |
| class MetadataExtractor: | |
| """Class for extracting and parsing metadata from images.""" | |
| def extract_metadata(image_path): | |
| """ | |
| Extract metadata from an image file. | |
| Args: | |
| image_path: path to the image file | |
| Returns: | |
| dict: dictionary with extracted metadata | |
| """ | |
| try: | |
| # Open image with PIL | |
| image = Image.open(image_path) | |
| # Extract metadata from PNG info | |
| metadata_text = image.info.get("parameters", "") | |
| # Parse the metadata | |
| parsed_metadata = MetadataExtractor.parse_metadata(metadata_text) | |
| # Add basic image info | |
| parsed_metadata.update({ | |
| 'width': image.width, | |
| 'height': image.height, | |
| 'format': image.format, | |
| 'mode': image.mode, | |
| }) | |
| return parsed_metadata | |
| except Exception as e: | |
| print(f"Error extracting metadata from {image_path}: {e}") | |
| return {'error': str(e)} | |
| def parse_metadata(metadata_text): | |
| """ | |
| Parse Stable Diffusion metadata text into structured data. | |
| Args: | |
| metadata_text: raw metadata text from image | |
| Returns: | |
| dict: structured metadata | |
| """ | |
| if not metadata_text: | |
| return {'raw_text': ''} | |
| result = {'raw_text': metadata_text} | |
| # Extract prompt | |
| prompt_end = metadata_text.find("Negative prompt:") | |
| if prompt_end > 0: | |
| result['prompt'] = metadata_text[:prompt_end].strip() | |
| negative_prompt_end = metadata_text.find("\n", prompt_end) | |
| if negative_prompt_end > 0: | |
| result['negative_prompt'] = metadata_text[prompt_end + len("Negative prompt:"):negative_prompt_end].strip() | |
| else: | |
| result['prompt'] = metadata_text.strip() | |
| # Extract model name | |
| model_match = re.search(r'Model: ([^,\n]+)', metadata_text) | |
| if model_match: | |
| result['model'] = model_match.group(1).strip() | |
| # Extract other parameters | |
| params = { | |
| 'steps': r'Steps: (\d+)', | |
| 'sampler': r'Sampler: ([^,\n]+)', | |
| 'cfg_scale': r'CFG scale: ([^,\n]+)', | |
| 'seed': r'Seed: ([^,\n]+)', | |
| 'size': r'Size: ([^,\n]+)', | |
| 'model_hash': r'Model hash: ([^,\n]+)', | |
| } | |
| for key, pattern in params.items(): | |
| match = re.search(pattern, metadata_text) | |
| if match: | |
| result[key] = match.group(1).strip() | |
| return result | |
| def group_images_by_model(metadata_list): | |
| """ | |
| Group images by model name. | |
| Args: | |
| metadata_list: list of (image_path, metadata) tuples | |
| Returns: | |
| dict: dictionary with model names as keys and lists of image paths as values | |
| """ | |
| result = {} | |
| for image_path, metadata in metadata_list: | |
| model = metadata.get('model', 'unknown') | |
| if model not in result: | |
| result[model] = [] | |
| result[model].append(image_path) | |
| return result | |
| def group_images_by_prompt(metadata_list): | |
| """ | |
| Group images by prompt. | |
| Args: | |
| metadata_list: list of (image_path, metadata) tuples | |
| Returns: | |
| dict: dictionary with prompts as keys and lists of image paths as values | |
| """ | |
| result = {} | |
| for image_path, metadata in metadata_list: | |
| prompt = metadata.get('prompt', 'unknown') | |
| # Use first 50 chars as key to avoid extremely long keys | |
| prompt_key = prompt[:50] + ('...' if len(prompt) > 50 else '') | |
| if prompt_key not in result: | |
| result[prompt_key] = [] | |
| result[prompt_key].append((image_path, metadata.get('model', 'unknown'))) | |
| return result | |
| def update_metadata(image_path, new_metadata, output_path=None): | |
| """ | |
| Update metadata in an image file. | |
| Args: | |
| image_path: path to the input image file | |
| new_metadata: new metadata text to write | |
| output_path: path to save the updated image (if None, overwrites input) | |
| Returns: | |
| bool: True if successful, False otherwise | |
| """ | |
| try: | |
| # Open image with PIL | |
| image = Image.open(image_path) | |
| # Create a PngInfo object to store metadata | |
| pnginfo = PngImagePlugin.PngInfo() | |
| pnginfo.add_text("parameters", new_metadata) | |
| # Save the image with the updated metadata | |
| save_path = output_path if output_path else image_path | |
| image.save(save_path, format="PNG", pnginfo=pnginfo) | |
| return True | |
| except Exception as e: | |
| print(f"Error updating metadata: {e}") | |
| return False | |