Spaces:
Runtime error
Runtime error
Update app.py
Browse files
app.py
CHANGED
|
@@ -20,7 +20,6 @@ import piexif.helper
|
|
| 20 |
|
| 21 |
TITLE = 'DeepDanbooru String'
|
| 22 |
|
| 23 |
-
TOKEN = os.environ['TOKEN']
|
| 24 |
MODEL_REPO = 'yo2266911/DeepDanbooru_string'
|
| 25 |
MODEL_FILENAME = 'model-resnet_custom_v3.h5'
|
| 26 |
LABEL_FILENAME = 'tags.txt'
|
|
@@ -47,33 +46,30 @@ def load_sample_image_paths() -> list[pathlib.Path]:
|
|
| 47 |
dataset_repo = 'hysts/sample-images-TADNE'
|
| 48 |
path = huggingface_hub.hf_hub_download(dataset_repo,
|
| 49 |
'images.tar.gz',
|
| 50 |
-
repo_type='dataset'
|
| 51 |
-
use_auth_token=TOKEN)
|
| 52 |
with tarfile.open(path) as f:
|
| 53 |
f.extractall()
|
| 54 |
return sorted(image_dir.glob('*'))
|
| 55 |
|
| 56 |
|
| 57 |
def load_model() -> tf.keras.Model:
|
| 58 |
-
path = huggingface_hub.hf_hub_download(MODEL_REPO,
|
| 59 |
-
MODEL_FILENAME,
|
| 60 |
-
use_auth_token=TOKEN)
|
| 61 |
model = tf.keras.models.load_model(path)
|
| 62 |
return model
|
| 63 |
|
| 64 |
|
| 65 |
def load_labels() -> list[str]:
|
| 66 |
-
path = huggingface_hub.hf_hub_download(MODEL_REPO,
|
| 67 |
-
LABEL_FILENAME,
|
| 68 |
-
use_auth_token=TOKEN)
|
| 69 |
with open(path) as f:
|
| 70 |
labels = [line.strip() for line in f.readlines()]
|
| 71 |
return labels
|
| 72 |
|
|
|
|
| 73 |
def plaintext_to_html(text):
|
| 74 |
text = "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split('\n')]) + "</p>"
|
| 75 |
return text
|
| 76 |
|
|
|
|
| 77 |
def predict(image: PIL.Image.Image, score_threshold: float,
|
| 78 |
model: tf.keras.Model, labels: list[str]) -> dict[str, float]:
|
| 79 |
rawimage = image
|
|
@@ -93,13 +89,13 @@ def predict(image: PIL.Image.Image, score_threshold: float,
|
|
| 93 |
if prob < score_threshold:
|
| 94 |
continue
|
| 95 |
res[label] = prob
|
| 96 |
-
b = dict(sorted(res.items(),key=lambda item:item[1], reverse=True))
|
| 97 |
-
a = ', '.join(list(b.keys())).replace('_',' ').replace('(','\(').replace(')','\)')
|
| 98 |
c = ', '.join(list(b.keys()))
|
| 99 |
-
|
| 100 |
items = rawimage.info
|
| 101 |
geninfo = ''
|
| 102 |
-
|
| 103 |
if "exif" in rawimage.info:
|
| 104 |
exif = piexif.load(rawimage.info["exif"])
|
| 105 |
exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b'')
|
|
@@ -107,16 +103,16 @@ def predict(image: PIL.Image.Image, score_threshold: float,
|
|
| 107 |
exif_comment = piexif.helper.UserComment.load(exif_comment)
|
| 108 |
except ValueError:
|
| 109 |
exif_comment = exif_comment.decode('utf8', errors="ignore")
|
| 110 |
-
|
| 111 |
items['exif comment'] = exif_comment
|
| 112 |
geninfo = exif_comment
|
| 113 |
-
|
| 114 |
for field in ['jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',
|
| 115 |
'loop', 'background', 'timestamp', 'duration']:
|
| 116 |
items.pop(field, None)
|
| 117 |
-
|
| 118 |
geninfo = items.get('parameters', geninfo)
|
| 119 |
-
|
| 120 |
info = f"""
|
| 121 |
<p><h4>PNG Info</h4></p>
|
| 122 |
"""
|
|
@@ -126,13 +122,13 @@ def predict(image: PIL.Image.Image, score_threshold: float,
|
|
| 126 |
<p><b>{plaintext_to_html(str(key))}</b></p>
|
| 127 |
<p>{plaintext_to_html(str(text))}</p>
|
| 128 |
</div>
|
| 129 |
-
""".strip()+"\n"
|
| 130 |
-
|
| 131 |
if len(info) == 0:
|
| 132 |
message = "Nothing found in the image."
|
| 133 |
info = f"<div><p>{message}<p></div>"
|
| 134 |
-
|
| 135 |
-
return (a,c,res,info)
|
| 136 |
|
| 137 |
|
| 138 |
def main():
|
|
@@ -154,14 +150,14 @@ def main():
|
|
| 154 |
label='Score Threshold'),
|
| 155 |
],
|
| 156 |
[
|
| 157 |
-
gr.outputs.Textbox(label='Output (string)'),
|
| 158 |
-
gr.outputs.Textbox(label='Output (raw string)'),
|
| 159 |
gr.outputs.Label(label='Output (label)'),
|
| 160 |
gr.outputs.HTML()
|
| 161 |
],
|
| 162 |
examples=[
|
| 163 |
-
|
| 164 |
-
|
| 165 |
],
|
| 166 |
title=TITLE,
|
| 167 |
description='''
|
|
@@ -180,4 +176,4 @@ PNG Info code forked from [AUTOMATIC1111/stable-diffusion-webui](https://github.
|
|
| 180 |
|
| 181 |
|
| 182 |
if __name__ == '__main__':
|
| 183 |
-
main()
|
|
|
|
| 20 |
|
| 21 |
TITLE = 'DeepDanbooru String'
|
| 22 |
|
|
|
|
| 23 |
MODEL_REPO = 'yo2266911/DeepDanbooru_string'
|
| 24 |
MODEL_FILENAME = 'model-resnet_custom_v3.h5'
|
| 25 |
LABEL_FILENAME = 'tags.txt'
|
|
|
|
| 46 |
dataset_repo = 'hysts/sample-images-TADNE'
|
| 47 |
path = huggingface_hub.hf_hub_download(dataset_repo,
|
| 48 |
'images.tar.gz',
|
| 49 |
+
repo_type='dataset')
|
|
|
|
| 50 |
with tarfile.open(path) as f:
|
| 51 |
f.extractall()
|
| 52 |
return sorted(image_dir.glob('*'))
|
| 53 |
|
| 54 |
|
| 55 |
def load_model() -> tf.keras.Model:
|
| 56 |
+
path = huggingface_hub.hf_hub_download(MODEL_REPO, MODEL_FILENAME)
|
|
|
|
|
|
|
| 57 |
model = tf.keras.models.load_model(path)
|
| 58 |
return model
|
| 59 |
|
| 60 |
|
| 61 |
def load_labels() -> list[str]:
|
| 62 |
+
path = huggingface_hub.hf_hub_download(MODEL_REPO, LABEL_FILENAME)
|
|
|
|
|
|
|
| 63 |
with open(path) as f:
|
| 64 |
labels = [line.strip() for line in f.readlines()]
|
| 65 |
return labels
|
| 66 |
|
| 67 |
+
|
| 68 |
def plaintext_to_html(text):
|
| 69 |
text = "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split('\n')]) + "</p>"
|
| 70 |
return text
|
| 71 |
|
| 72 |
+
|
| 73 |
def predict(image: PIL.Image.Image, score_threshold: float,
|
| 74 |
model: tf.keras.Model, labels: list[str]) -> dict[str, float]:
|
| 75 |
rawimage = image
|
|
|
|
| 89 |
if prob < score_threshold:
|
| 90 |
continue
|
| 91 |
res[label] = prob
|
| 92 |
+
b = dict(sorted(res.items(), key=lambda item: item[1], reverse=True))
|
| 93 |
+
a = ', '.join(list(b.keys())).replace('_', ' ').replace('(', '\(').replace(')', '\)')
|
| 94 |
c = ', '.join(list(b.keys()))
|
| 95 |
+
|
| 96 |
items = rawimage.info
|
| 97 |
geninfo = ''
|
| 98 |
+
|
| 99 |
if "exif" in rawimage.info:
|
| 100 |
exif = piexif.load(rawimage.info["exif"])
|
| 101 |
exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b'')
|
|
|
|
| 103 |
exif_comment = piexif.helper.UserComment.load(exif_comment)
|
| 104 |
except ValueError:
|
| 105 |
exif_comment = exif_comment.decode('utf8', errors="ignore")
|
| 106 |
+
|
| 107 |
items['exif comment'] = exif_comment
|
| 108 |
geninfo = exif_comment
|
| 109 |
+
|
| 110 |
for field in ['jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',
|
| 111 |
'loop', 'background', 'timestamp', 'duration']:
|
| 112 |
items.pop(field, None)
|
| 113 |
+
|
| 114 |
geninfo = items.get('parameters', geninfo)
|
| 115 |
+
|
| 116 |
info = f"""
|
| 117 |
<p><h4>PNG Info</h4></p>
|
| 118 |
"""
|
|
|
|
| 122 |
<p><b>{plaintext_to_html(str(key))}</b></p>
|
| 123 |
<p>{plaintext_to_html(str(text))}</p>
|
| 124 |
</div>
|
| 125 |
+
""".strip() + "\n"
|
| 126 |
+
|
| 127 |
if len(info) == 0:
|
| 128 |
message = "Nothing found in the image."
|
| 129 |
info = f"<div><p>{message}<p></div>"
|
| 130 |
+
|
| 131 |
+
return (a, c, res, info)
|
| 132 |
|
| 133 |
|
| 134 |
def main():
|
|
|
|
| 150 |
label='Score Threshold'),
|
| 151 |
],
|
| 152 |
[
|
| 153 |
+
gr.outputs.Textbox(label='Output (string)'),
|
| 154 |
+
gr.outputs.Textbox(label='Output (raw string)'),
|
| 155 |
gr.outputs.Label(label='Output (label)'),
|
| 156 |
gr.outputs.HTML()
|
| 157 |
],
|
| 158 |
examples=[
|
| 159 |
+
['miku.jpg', 0.5],
|
| 160 |
+
['miku2.jpg', 0.5]
|
| 161 |
],
|
| 162 |
title=TITLE,
|
| 163 |
description='''
|
|
|
|
| 176 |
|
| 177 |
|
| 178 |
if __name__ == '__main__':
|
| 179 |
+
main()
|