Spaces:
Sleeping
Sleeping
File size: 5,466 Bytes
f89e218 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 |
"""
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."""
@staticmethod
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)}
@staticmethod
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
@staticmethod
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
@staticmethod
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
@staticmethod
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
|