DeepDanbooru / app.py
MonkeyJuice's picture
split code
7c078a3
raw
history blame
2.43 kB
#!/usr/bin/env python
from __future__ import annotations
import gradio as gr
import PIL.Image
from genTag import genTag
def predict(image: PIL.Image.Image, score_threshold: float):
result_threshold = genTag(image, score_threshold)
result_html = ''
for label, prob in result_threshold.items():
result_html = result_html + '<p class="m5dd_list use"><span>' + str(label) + '</span><span>' + str(round(prob, 3)) + '</span></p>'
result_html = '<div>' + str(result_html) + '</div>'
result_text = ', '.join(result_threshold.keys())
result_text = '<div id="m5dd_result">' + str(result_text) + '</div>'
return result_html, result_text
js = """
async () => {
document.addEventListener('click', function(event) {
let tagItem = event.target.closest('.m5dd_list')
let resultArea = event.target.closest('#m5dd_result')
if (tagItem){
if (tagItem.classList.contains('use')){
tagItem.classList.remove('use')
}else{
tagItem.classList.add('use')
}
document.getElementById('m5dd_result').innerText =
Array.from(document.querySelectorAll('.m5dd_list.use>span:nth-child(1)'))
.map(v=>v.innerText)
.join(', ')
}else if (resultArea){
const selection = window.getSelection()
selection.removeAllRanges()
const range = document.createRange()
range.selectNodeContents(resultArea)
selection.addRange(range)
}else{
return
}
})
}
"""
with gr.Blocks(css="style.css") as demo:
with gr.Row():
with gr.Column(scale=1):
image = gr.Image(label='Input', type='pil')
score_threshold = gr.Slider(label='Score threshold',
minimum=0,
maximum=1,
step=0.05,
value=0.5)
run_button = gr.Button('Run')
result_text = gr.HTML(value="<div></div>")
with gr.Column(scale=3):
result_html = gr.HTML(value="<div></div>")
run_button.click(
fn=predict,
inputs=[image, score_threshold],
outputs=[result_html, result_text],
api_name='predict',
)
demo.load(None,None,None,_js=js)
demo.queue().launch()