from pathlib import Path
import gradio as gr
import torch
from transformers import AutoModelForImageClassification
import shutil
from optimum.pipelines import pipeline

device = 1 if torch.cuda.is_available() else "cpu"

chk_point = "davanstrien/autotrain-ia-useful-covers-3665397856"

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)