Update app.py

app.py

@@ -1,17 +1,57 @@
+from pathlib import Path
 import gradio as gr
-from transformers import pipeline
-
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
-
-extractor = AutoFeatureExtractor.from_pretrained("kdhht2334/autotrain-diffusion-emotion-facial-expression-recognition-40429105176")
-fer_model = AutoModelForImageClassification.from_pretrained("kdhht2334/autotrain-diffusion-emotion-facial-expression-recognition-40429105176")
-
-pipe = pipeline(task="image-classification", 
-                # model that can do 7-emotion classification
-                model=fer_model)
-gr.Interface.from_pipeline(pipe, 
-                           title="7-emotion classification in DiffusionFER dataset",
-                           description="Facial expression recognition trained with Transformer",
-                           examples = ['happy.png', 'angry.png',],
-                           article = "Author: <a href=\"https://huggingface.co/FER-Universe\">Daeha Kim</a>",
-                           ).launch()
+import torch
+from transformers import AutoModelForImageClassification
+import shutil
+from optimum.pipelines import pipeline
+
+device = 1 if torch.cuda.is_available() else "cpu"
+
+chk_point = "kdhht2334/autotrain-diffusion-emotion-facial-expression-recognition-40429105176"
+
+model = AutoModelForImageClassification.from_pretrained(chk_point)
+
+try:
+    pipe = pipeline(
+        "image-classification",
+        chk_point,
+        accelerator="bettertransformer",
+        device=device,
+    )
+except NotImplementedError:
+    from transformers import pipeline
+
+    pipe = pipeline("image-classification", chk_point, device=device)
+
+
+def make_label_folders():
+    folders = model.config.label2id.keys()
+    for folder in folders:
+        folder = Path(folder)
+        if not folder.exists():
+            folder.mkdir()
+    return folders
+
+
+def predictions_into_folders(files):
+    files = [file.name for file in files]
+    files = [
+        file for file in files if not file.startswith(".") and "DS_Store" not in file
+    ]
+    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.Files(file_count="directory", file_types=["image"]),
+    gr.Files(),
+    cache_examples=True,
+)
+demo.launch(enable_queue=True)