add batch handle

README.md

@@ -4,7 +4,7 @@ emoji: 🏃
 colorFrom: gray
 colorTo: purple
 sdk: gradio
-sdk_version: 3.37.0
+sdk_version: 4.19.2
 app_file: app.py
 pinned: false
 ---

app.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import gradio as gr
 import PIL.Image
+import zipfile
 from genTag import genTag
 
 def predict(image: PIL.Image.Image, score_threshold: float):
@@ -16,47 +17,51 @@ def predict(image: PIL.Image.Image, score_threshold: float):
     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
-        }
-    })
-}
-"""
+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)
+                tag = ', '.join(result_threshold.keys())
+                result = result + str(file) + '\n' + str(tag) + '\n'
+    return result
 
-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>")
+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',
+                                 sources=["upload", "clipboard"],
+                                 height='20em')
+                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=2):
+                result_html = gr.HTML(value="<div></div>")
+    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'],
+                                     height='20em')
+                score_threshold2 = gr.Slider(label='Score threshold',
+                                             minimum=0,
+                                             maximum=1,
+                                             step=0.05,
+                                             value=0.5)
+                run_button2 = gr.Button('Run')
+            with gr.Column(scale=2):
+                result_text2 = gr.Textbox(lines=5, show_copy_button=True)
 
     run_button.click(
         fn=predict,
@@ -64,6 +69,11 @@ with gr.Blocks(css="style.css") as demo:
         outputs=[result_html, result_text],
         api_name='predict',
     )
-    demo.load(None,None,None,_js=js)
+    run_button2.click(
+        fn=predict_batch,
+        inputs=[batch_file, score_threshold2],
+        outputs=[result_text2],
+        api_name='predict_batch',
+    )
 
 demo.queue().launch()

script.js

@@ -0,0 +1,24 @@
+
+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
+  }
+})