Spaces:
Running
Running
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()
|