Spaces:
Running
Running
| #!/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, ignore2 | |
| def predict(image: PIL.Image.Image): | |
| result_threshold = genTag(image, 0.5) | |
| return result_threshold, ignore2, """<div></div>""" | |
| def predict_api(image: PIL.Image.Image): | |
| result_threshold = genTag(image, 0.5) | |
| result_filter = {key: value for key, value in result_threshold.items() if not is_ignore(key, 2)} | |
| tag = ', '.join(result_filter.keys()) | |
| return str(tag) | |
| def predict_batch(zip_file, 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") or file.endswith(".jpeg") or file.endswith(".webp"): | |
| image_file = zf.open(file) | |
| image = PIL.Image.open(image_file) | |
| image = image.convert("RGBA") | |
| result_threshold = genTag(image, 0.5) | |
| 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(head_paths="head.html") 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', | |
| image_mode="RGBA", | |
| show_fullscreen_button=False, | |
| sources=["upload", "clipboard"]) | |
| result_text = gr.HTML(value="""<div id="m5dd_result"></div>""", padding=False) | |
| result_hide = gr.JSON(visible=False) | |
| result_hide2 = gr.JSON(visible=False) | |
| with gr.Column(scale=2): | |
| result_html = gr.HTML(value="""<div id="m5dd_list"></div>""", padding=False) | |
| result_loading = gr.HTML(value="""<div></div>""", elem_classes='m5dd_html', padding=False) | |
| 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']) | |
| run_button2 = gr.Button('Run') | |
| run_button_api = gr.Button(value='Run', visible=False) | |
| with gr.Column(scale=2): | |
| result_text2 = gr.Textbox(lines=20, | |
| max_lines=20, | |
| label='Result', | |
| show_copy_button=True, | |
| autoscroll=False) | |
| image.upload( | |
| fn=predict, | |
| inputs=[image], | |
| outputs=[result_hide, result_hide2, result_loading], | |
| api_name=False, | |
| js=""" | |
| (image) => { | |
| window.m5Func.clear() | |
| return image; | |
| } | |
| """, | |
| ).success( | |
| fn=None, | |
| inputs=[result_hide, result_hide2], | |
| js=""" | |
| (result, ignore) => { | |
| window.m5Func.refresh(result, ignore) | |
| return [result, ignore]; | |
| } | |
| """, | |
| ) | |
| run_button2.click( | |
| fn=predict_batch, | |
| inputs=[batch_file], | |
| outputs=[result_text2], | |
| api_name=False, | |
| ) | |
| run_button_api.click( | |
| fn=predict_api, | |
| inputs=[image], | |
| outputs=[result_text2], | |
| api_name='predict', | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue(max_size=20).launch() |