Update app.py

app.py

@@ -14,6 +14,7 @@ import huggingface_hub
 import numpy as np
 import PIL.Image
 import tensorflow as tf
+import piexif
 
 TITLE = 'NoCrypt/DeepDanbooru_string'
 DESCRIPTION = 'Cloned from: https://huggingface.co/spaces/hysts/DeepDanbooru'
@@ -68,9 +69,13 @@ def load_labels() -> list[str]:
         labels = [line.strip() for line in f.readlines()]
     return labels
 
+def plaintext_to_html(text):
+    text = "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split('\n')]) + "</p>"
+    return text
 
 def predict(image: PIL.Image.Image, score_threshold: float,
             model: tf.keras.Model, labels: list[str]) -> dict[str, float]:
+    rawimage = image
     _, height, width, _ = model.input_shape
     image = np.asarray(image)
     image = tf.image.resize(image,
@@ -89,7 +94,41 @@ def predict(image: PIL.Image.Image, score_threshold: float,
         res[label] = prob
     b = dict(sorted(res.items(),key=lambda item:item[1], reverse=True))
     a = ', '.join(list(b.keys())).replace('_',' ').replace('(','\(').replace(')','\)')
-    return (a,res)
+    
+    items = rawimage.info
+    geninfo = ''
+    
+    if "exif" in rawimage.info:
+        exif = piexif.load(rawimage.info["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")
+    
+        items['exif comment'] = exif_comment
+        geninfo = exif_comment
+    
+        for field in ['jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',
+                      'loop', 'background', 'timestamp', 'duration']:
+            items.pop(field, None)
+    
+    geninfo = items.get('parameters', geninfo)
+    
+    info = ''
+    for key, text in items.items():
+        info += f"""
+<div>
+<p><b>{plaintext_to_html(str(key))}</b></p>
+<p>{plaintext_to_html(str(text))}</p>
+</div>
+""".strip()+"\n"
+    
+    if len(info) == 0:
+        message = ""
+        info = f"<div><p>{message}<p></div>"
+    
+    return (a,res,geninfo,info)
 
 
 def main():
@@ -112,7 +151,12 @@ def main():
                              default=args.score_threshold,
                              label='Score Threshold'),
         ],
-        [gr.outputs.Textbox(label='Output String'), gr.outputs.Label(label='Output Labels')],
+        [
+            gr.outputs.Textbox(label='Output String'), 
+            gr.outputs.Label(label='Output Labels'),
+            gr.outputs.HTML(visible=len(geninfo)!=0),
+            gr.outputs.HTML(visible=len(info)!=0)
+        ],
         examples=[
         ['miku.jpg',0.5],
         ['miku2.jpg',0.5]