import os

os.system("pip uninstall numpy -y")
os.system("pip install numpy")
os.system("pip install pandas")

import gradio as gr
import sys
from uuid import uuid1
from PIL import Image
from zipfile import ZipFile
import pathlib
import shutil
import pandas as pd
import deepsparse

rn50_embedding_pipeline = deepsparse.Pipeline.create(
    task="embedding-extraction",
    base_task="image-classification", # tells the pipeline to expect images and normalize input with ImageNet means/stds
    model_path="zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/channel20_pruned75_quant-none-vnni",
    emb_extraction_layer=-3, # extracts last layer before projection head and softmax
)

def zip_ims(g):
    from uuid import uuid1
    if g is None:
        return None
    l = list(map(lambda x: x["name"], g))
    if not l:
        return None
    zip_file_name ="tmp.zip"
    with ZipFile(zip_file_name ,"w") as zipObj:
        for ele in l:
            zipObj.write(ele, "{}.png".format(uuid1()))
        #zipObj.write(file2.name, "file2")
    return zip_file_name

def unzip_ims(zip_file_name):
    print("call file")
    if zip_file_name is None:
        return {}
    unzip_path = "img_dir"
    if os.path.exists(unzip_path):
        shutil.rmtree(unzip_path)
    with ZipFile(zip_file_name) as archive:
        archive.extractall(unzip_path)
    im_name_l = pd.Series(
    list(pathlib.Path(unzip_path).rglob("*.png")) + \
    list(pathlib.Path(unzip_path).rglob("*.jpg")) + \
    list(pathlib.Path(unzip_path).rglob("*.jpeg"))
    ).map(str).values.tolist()
    embeddings = rn50_embedding_pipeline(images=im_name_l)
    if os.path.exists(unzip_path):
        shutil.rmtree(unzip_path)
    return {
        "names": im_name_l,
        "embs": embeddings
    }


def emb_img_func(im):
    print("call im :")
    if im is None:
        return {}
    im_obj = Image.fromarray(im)
    im_name = "{}.png".format(uuid1())
    im_obj.save(im_name)
    embeddings = rn50_embedding_pipeline(images=[im_name])
    os.remove(im_name)
    return {
        "names": [im_name],
        "embs": embeddings
    }

'''
def emb_gallery_func(gallery):
    print("call ga :")
    if gallery is None:
        return []
    im_name_l = list(map(lambda x: x["name"], images))
    embeddings = rn50_embedding_pipeline(images=im_name_l)
    return embeddings
'''

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            inputs_0 = gr.Image(label = "Input Image for embed")
            button_0 = gr.Button("Image button")
        with gr.Column():
            inputs_1 = gr.File(label = "Input Images zip file for embed")
            button_1 = gr.Button("Image File button")
    with gr.Row():
        outputs = gr.JSON(label = "Output Embeddings")

    button_0.click(fn = emb_img_func, inputs = inputs_0, outputs = outputs)
    button_1.click(fn = unzip_ims, inputs = inputs_1, outputs = outputs)

demo.launch("0.0.0.0")