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