update'

app.py

@@ -1,17 +1,47 @@
-import gradio as gr
-# from transformers import pipeline
 from pathlib import Path
+import gradio as gr
+from transformers import pipeline
+from transformers import AutoModelForImageClassification
+import shutil
 
-# pipe = pipeline('image-classification','models/davanstrien/autotrain-encyclopaedia-illustrations-blog-post-3327992159')
 
-def upload_file(files):
-    return files
+chk_point = "davanstrien/autotrain-encyclopaedia-illustrations-blog-post-3327992158"
 
-with gr.Blocks() as demo:
-    file_output = gr.File()
-    upload_button = gr.UploadButton("Click to Upload a Directory", file_types=["image", "video"], file_count="directory")
-    upload_button.upload(upload_file, upload_button, file_output)
+model = AutoModelForImageClassification.from_pretrained(chk_point)
 
+pipe = pipeline(
+    "image-classification",
+    "davanstrien/autotrain-encyclopaedia-illustrations-blog-post-3327992158",
+)
+
+
+def make_label_folders():
+    folders = model.config.label2id.keys()
+    for folder in folders:
+        folder = Path(folder)
+        if not folder.exists():
+            folder.mkdir()
+    return folders
 
-demo.launch()
 
+def predictions_into_folders(files):
+    files = [file.name for file in files]
+    folders = make_label_folders()
+    predictions = pipe(files)
+    for file, prediction in zip(files, predictions):
+        label = prediction[0]["label"]
+        file_name = Path(file).name
+        shutil.copy(file, f"{label}/{file_name}")
+    for folder in folders:
+        shutil.make_archive(folder, "zip", ".", folder)
+    return [f"{folder}.zip" for folder in folders]
+
+
+demo = gr.Interface(
+    predictions_into_folders,
+    gr.File(file_count="directory"),
+    gr.Files(),
+    cache_examples=True,
+)
+
+demo.launch()