Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python | |
| from __future__ import annotations | |
| import gradio as gr | |
| import PIL.Image | |
| import zipfile | |
| from genTag import genTag | |
| from checkIgnore import is_ignore | |
| from createTagDom import create_tag_dom | |
| 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 += create_tag_dom(label, is_ignore(label, 1), prob) | |
| result_html = '<div>' + result_html + '</div>' | |
| result_filter = {key: value for key, value in result_threshold.items() if not is_ignore(key, 1)} | |
| result_text = '<div id="m5dd_result">' + ', '.join(result_filter.keys()) + '</div>' | |
| return result_html, result_text | |
| def predict_batch(zip_file, score_threshold: float, progress=gr.Progress()): | |
| result = '' | |
| with zipfile.ZipFile(zip_file) as zf: | |
| for file in progress.tqdm(zf.namelist()): | |
| print(file) | |
| if file.endswith(".png") or file.endswith(".jpg"): | |
| image_file = zf.open(file) | |
| image = PIL.Image.open(image_file) | |
| image = image.convert("RGB") | |
| result_threshold = genTag(image, score_threshold) | |
| result_filter = {key: value for key, value in result_threshold.items() if not is_ignore(key, 2)} | |
| tag = ', '.join(result_filter.keys()) | |
| result = result + str(file) + '\n' + str(tag) + '\n\n' | |
| return result | |
| with gr.Blocks(css="style.css", js="script.js") as demo: | |
| with gr.Tab(label='Single'): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| image = gr.Image(label='Upload a image', | |
| type='pil', | |
| elem_classes='m5dd_image', | |
| sources=["upload", "clipboard"]) | |
| score_threshold = gr.Slider(label='Score threshold', | |
| minimum=0, | |
| maximum=1, | |
| step=0.1, | |
| value=0.3) | |
| run_button = gr.Button('Run') | |
| result_text = gr.HTML(value="") | |
| with gr.Column(scale=2): | |
| result_html = gr.HTML(value="") | |
| with gr.Tab(label='Batch'): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| batch_file = gr.File(label="Upload a ZIP file containing images", | |
| file_types=['.zip']) | |
| score_threshold2 = gr.Slider(label='Score threshold', | |
| minimum=0, | |
| maximum=1, | |
| step=0.1, | |
| value=0.3) | |
| run_button2 = gr.Button('Run') | |
| with gr.Column(scale=2): | |
| result_text2 = gr.Textbox(lines=20, | |
| max_lines=20, | |
| label='Result', | |
| show_copy_button=True, | |
| autoscroll=False) | |
| run_button.click( | |
| fn=predict, | |
| inputs=[image, score_threshold], | |
| outputs=[result_html, result_text], | |
| api_name='predict', | |
| ) | |
| run_button2.click( | |
| fn=predict_batch, | |
| inputs=[batch_file, score_threshold2], | |
| outputs=[result_text2], | |
| api_name='predict_batch', | |
| ) | |
| demo.queue().launch() | |