read stealth image info

app.py

@@ -205,7 +205,7 @@ def util_ui():
     with gr.Blocks():
         with gr.Row(equal_height=False):
             with gr.Column(variant='panel'):
-                image = gr.Image(label="上传图片", sources=["upload"], interactive=True, type="pil")
+                image = gr.Image(label="上传图片", image_mode="RGBA", sources=["upload"], interactive=True, type="pil")
             with gr.Column(variant='panel'):
                 info = gr.HTML('')
                 items = gr.JSON(value={}, visible=False)

pnginfo.py

@@ -1,7 +1,9 @@
 from PIL import Image
+import piexif
 import json
 import html
 import re
+import gzip
 
 re_param_code = r'\s*([\w ]+):\s*("(?:\\.|[^\\"])+"|[^,]*)(?:,|$)'
 re_param = re.compile(re_param_code)
@@ -9,7 +11,7 @@ re_imagesize = re.compile(r"^(\d+)x(\d+)$")
 IGNORED_INFO_KEYS = {
     'jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',
     'loop', 'background', 'timestamp', 'duration', 'progressive', 'progression',
-    'icc_profile', 'chromaticity', 'photoshop',
+    'icc_profile', 'chromaticity', 'photoshop', 'srgb', 'gamma', 'dpi'
 }
 
 def plaintext_to_html(text, classname=None):
@@ -22,7 +24,114 @@ def to_digit(v):
         return float(v)
     except:
         return v
-    
+
+def read_stealth_info(image):
+    width, height = image.size
+    pixels = image.load()
+
+    has_alpha = True if image.mode == 'RGBA' else False
+    mode = None
+    compressed = False
+    binary_data = ''
+    buffer_a = ''
+    buffer_rgb = ''
+    index_a = 0
+    index_rgb = 0
+    sig_confirmed = False
+    confirming_signature = True
+    reading_param_len = False
+    reading_param = False
+    read_end = False
+    for x in range(width):
+        for y in range(height):
+            if has_alpha:
+                r, g, b, a = pixels[x, y]
+                buffer_a += str(a & 1)
+                index_a += 1
+            else:
+                r, g, b = pixels[x, y]
+            buffer_rgb += str(r & 1)
+            buffer_rgb += str(g & 1)
+            buffer_rgb += str(b & 1)
+            index_rgb += 3
+            if confirming_signature:
+                if index_a == len('stealth_pnginfo') * 8:
+                    decoded_sig = bytearray(int(buffer_a[i:i + 8], 2) for i in
+                                            range(0, len(buffer_a), 8)).decode('utf-8', errors='ignore')
+                    if decoded_sig in {'stealth_pnginfo', 'stealth_pngcomp'}:
+                        confirming_signature = False
+                        sig_confirmed = True
+                        reading_param_len = True
+                        mode = 'alpha'
+                        if decoded_sig == 'stealth_pngcomp':
+                            compressed = True
+                        buffer_a = ''
+                        index_a = 0
+                    else:
+                        read_end = True
+                        break
+                elif index_rgb == len('stealth_pnginfo') * 8:
+                    decoded_sig = bytearray(int(buffer_rgb[i:i + 8], 2) for i in
+                                            range(0, len(buffer_rgb), 8)).decode('utf-8', errors='ignore')
+                    if decoded_sig in {'stealth_rgbinfo', 'stealth_rgbcomp'}:
+                        confirming_signature = False
+                        sig_confirmed = True
+                        reading_param_len = True
+                        mode = 'rgb'
+                        if decoded_sig == 'stealth_rgbcomp':
+                            compressed = True
+                        buffer_rgb = ''
+                        index_rgb = 0
+            elif reading_param_len:
+                if mode == 'alpha':
+                    if index_a == 32:
+                        param_len = int(buffer_a, 2)
+                        reading_param_len = False
+                        reading_param = True
+                        buffer_a = ''
+                        index_a = 0
+                else:
+                    if index_rgb == 33:
+                        pop = buffer_rgb[-1]
+                        buffer_rgb = buffer_rgb[:-1]
+                        param_len = int(buffer_rgb, 2)
+                        reading_param_len = False
+                        reading_param = True
+                        buffer_rgb = pop
+                        index_rgb = 1
+            elif reading_param:
+                if mode == 'alpha':
+                    if index_a == param_len:
+                        binary_data = buffer_a
+                        read_end = True
+                        break
+                else:
+                    if index_rgb >= param_len:
+                        diff = param_len - index_rgb
+                        if diff < 0:
+                            buffer_rgb = buffer_rgb[:diff]
+                        binary_data = buffer_rgb
+                        read_end = True
+                        break
+            else:
+                # impossible
+                read_end = True
+                break
+        if read_end:
+            break
+    if sig_confirmed and binary_data != '':
+        # Convert binary string to UTF-8 encoded text
+        byte_data = bytearray(int(binary_data[i:i + 8], 2) for i in range(0, len(binary_data), 8))
+        try:
+            if compressed:
+                decoded_data = gzip.decompress(bytes(byte_data)).decode('utf-8')
+            else:
+                decoded_data = byte_data.decode('utf-8', errors='ignore')
+            geninfo = decoded_data
+        except:
+            pass
+    return json.loads(geninfo)
+
 def read_info_from_image(image):
     if image is None:
         return '', {}
@@ -34,6 +143,21 @@ def read_info_from_image(image):
     for field in IGNORED_INFO_KEYS:
         items.pop(field, None)
 
+    if len(items) == 0:
+        items = read_stealth_info(image)
+
+    if "exif" in items:
+        exif = piexif.load(items["exif"])
+        exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b'')
+        try:
+            exif_comment = piexif.helper.UserComment.load(exif_comment)
+        except ValueError:
+            exif_comment = exif_comment.decode('utf8', errors="ignore")
+
+        if exif_comment:
+            items['exif comment'] = exif_comment
+            geninfo = exif_comment
+
     if items.get("Software", None) == "NovelAI":
         json_info = json.loads(items["Comment"])
 
@@ -101,7 +225,7 @@ Steps: {json_info["steps"]}, Sampler: {json_info['sampler']}, CFG scale: {json_i
 
 
 def send_paras(*args):
-    if len(args[0]) == 0:
+    if args[0] is None or len(args[0]) == 0:
         return args[1:]
     items, prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale, variety = args
     paras = [prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale, variety]