Create app.py

app.py

@@ -0,0 +1,426 @@
+'''
+!pip install "deepsparse-nightly==1.6.0.20231007"
+!pip install "deepsparse[image_classification]"
+!pip install opencv-python-headless
+!pip uninstall numpy -y
+!pip install numpy
+!pip install gradio
+!pip install pandas
+'''
+
+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
+import json
+import numpy as np
+
+rn50_embedding_pipeline_default = 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=-1, # extracts last layer before projection head and softmax
+)
+
+rn50_embedding_pipeline_last_1 = 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=-1, # extracts last layer before projection head and softmax
+)
+
+rn50_embedding_pipeline_last_2 = 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=-2, # extracts last layer before projection head and softmax
+)
+
+rn50_embedding_pipeline_last_3 = 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
+)
+
+rn50_embedding_pipeline_dict = {
+    "0": rn50_embedding_pipeline_default,
+    "1": rn50_embedding_pipeline_last_1,
+    "2": rn50_embedding_pipeline_last_2,
+    "3": rn50_embedding_pipeline_last_3
+}
+
+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_func(zip_file_name, choose_model,
+    rn50_embedding_pipeline_dict = rn50_embedding_pipeline_dict):
+    print("call file")
+    if zip_file_name is None:
+        return json.dumps({}), None
+    print("zip_file_name :")
+    print(zip_file_name)
+    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()
+    rn50_embedding_pipeline = rn50_embedding_pipeline_dict[choose_model]
+    embeddings = rn50_embedding_pipeline(images=im_name_l)
+    im_l = pd.Series(im_name_l).map(Image.open).values.tolist()
+    if os.path.exists(unzip_path):
+        shutil.rmtree(unzip_path)
+    im_name_l = pd.Series(im_name_l).map(lambda x: x.split("/")[-1]).values.tolist()
+    return json.dumps({
+        "names": im_name_l,
+        "embs": embeddings.embeddings[0]
+    }), im_l
+
+
+def emb_img_func(im, choose_model,
+    rn50_embedding_pipeline_dict = rn50_embedding_pipeline_dict):
+    print("call im :")
+    if im is None:
+        return json.dumps({})
+    im_obj = Image.fromarray(im)
+    im_name = "{}.png".format(uuid1())
+    im_obj.save(im_name)
+    rn50_embedding_pipeline = rn50_embedding_pipeline_dict[choose_model]
+    embeddings = rn50_embedding_pipeline(images=[im_name])
+    os.remove(im_name)
+    return json.dumps({
+        "names": [im_name],
+        "embs": embeddings.embeddings[0]
+    })
+
+def image_grid(imgs, rows, cols):
+    assert len(imgs) <= rows*cols
+    w, h = imgs[0].size
+    grid = Image.new('RGB', size=(cols*w, rows*h))
+    grid_w, grid_h = grid.size
+
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i%cols*w, i//cols*h))
+    return grid
+
+def expand2square(pil_img, background_color):
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+def image_click(images, evt: gr.SelectData,
+    choose_model,
+    rn50_embedding_pipeline_dict = rn50_embedding_pipeline_dict,
+    top_k = 5
+    ):
+
+    images = json.loads(images.model_dump_json())
+    images = list(map(lambda x: {"name": x["image"]["path"]}, images))
+
+    img_selected = images[evt.index]
+    pivot_image_path = images[evt.index]['name']
+
+    im_name_l = list(map(lambda x: x["name"], images))
+    rn50_embedding_pipeline = rn50_embedding_pipeline_dict[choose_model]
+    embeddings = rn50_embedding_pipeline(images=im_name_l)
+    json_text = json.dumps({
+        "names": im_name_l,
+        "embs": embeddings.embeddings[0]
+    })
+
+    assert type(json_text) == type("")
+    assert type(pivot_image_path) in [type(""), type(0)]
+    dd_obj = json.loads(json_text)
+    names = dd_obj["names"]
+    embs = dd_obj["embs"]
+
+    assert pivot_image_path in names
+    corr_df = pd.DataFrame(np.asarray(embs).T).corr()
+    corr_df.columns = names
+    corr_df.index = names
+    arr_l = []
+    for i, r in corr_df.iterrows():
+        arr_ll = sorted(r.to_dict().items(), key = lambda t2: t2[1], reverse = True)
+        arr_l.append(arr_ll)
+    top_k = min(len(corr_df), top_k)
+    cols = pd.Series(arr_l[names.index(pivot_image_path)]).map(lambda x: x[0]).values.tolist()[:top_k]
+    corr_array_df = pd.DataFrame(arr_l).applymap(lambda x: x[0])
+    corr_array_df.index = names
+    #### corr_array
+    corr_array = corr_array_df.loc[cols].iloc[:, :top_k].values
+    l_list = pd.Series(corr_array.reshape([-1])).values.tolist()
+    l_list = pd.Series(l_list).map(Image.open).map(lambda x: expand2square(x, (0, 0, 0))).values.tolist()
+    l_dist_list = []
+    for ele in l_list:
+        if ele not in l_dist_list:
+            l_dist_list.append(ele)
+    return l_dist_list, l_list
+
+import gradio as gr
+from Lex import *
+'''
+lex = Lexica(query="man woman fire snow").images()
+'''
+from PIL import Image
+import imagehash
+import requests
+
+from zipfile import ZipFile
+
+from time import sleep
+sleep_time = 0.5
+
+hash_func_name = list(filter(lambda x: x.endswith("hash") and
+"hex" not in x ,dir(imagehash)))
+hash_func_name = ['average_hash', 'colorhash', 'dhash', 'phash', 'whash', 'crop_resistant_hash',]
+
+def min_dim_to_size(img, size = 512):
+    h, w = img.size
+    ratio = size / max(h, w)
+    h, w = map(lambda x: int(x * ratio), [h, w])
+    return ( ratio ,img.resize((h, w)) )
+
+#ratio_size = 512
+#ratio, img_rs = min_dim_to_size(img, ratio_size)
+
+'''
+def image_click(images, evt: gr.SelectData):
+    img_selected = images[evt.index]
+    return images[evt.index]['name']
+
+def swap_gallery(im, images, func_name):
+    #### name data is_file
+    #print(images[0].keys())
+    if im is None:
+        return list(map(lambda x: x["name"], images))
+    hash_func = getattr(imagehash, func_name)
+
+    im_hash = hash_func(Image.fromarray(im))
+    t2_list = sorted(images, key = lambda imm:
+    hash_func(Image.open(imm["name"])) - im_hash, reverse = False)
+    return list(map(lambda x: x["name"], t2_list))
+'''
+
+
+
+def lexica(prompt, limit_size = 128, ratio_size = 256 + 128):
+    lex = Lexica(query=prompt).images()
+    lex = lex[:limit_size]
+    lex = list(map(lambda x: x.replace("full_jpg", "sm2"), lex))
+    lex_ = []
+    for ele in lex:
+        try:
+            im = Image.open(
+                            requests.get(ele, stream = True).raw
+                            )
+            lex_.append(im)
+        except:
+            print("err")
+        sleep(sleep_time)
+    assert lex_
+    lex = list(map(lambda x: min_dim_to_size(x, ratio_size)[1], lex_))
+    return lex
+
+def enterpix(prompt, limit_size = 100, ratio_size = 256 + 128, use_key = "bigThumbnailUrl"):
+    resp = requests.post(
+    url = "https://www.enterpix.app/enterpix/v1/image/prompt-search",
+    data= {
+        "length": limit_size,
+        "platform": "stable-diffusion,midjourney",
+        "prompt": prompt,
+        "start": 0
+        }
+    )
+    resp = resp.json()
+    resp = list(map(lambda x: x[use_key], resp["images"]))
+    lex_ = []
+    for ele in resp:
+        try:
+            im = Image.open(
+                            requests.get(ele, stream = True).raw
+                            )
+            lex_.append(im)
+        except:
+            print("err")
+        sleep(sleep_time)
+    assert lex_
+    resp = list(map(lambda x: min_dim_to_size(x, ratio_size)[1], lex_))
+    return resp
+
+#def search(prompt, search_name, im, func_name):
+def search(prompt, search_name,):
+    if search_name == "lexica":
+        im_l = lexica(prompt)
+    else:
+        im_l = enterpix(prompt)
+    return im_l
+    '''
+    if im is None:
+        return im_l
+    hash_func = getattr(imagehash, func_name)
+
+    im_hash = hash_func(Image.fromarray(im))
+    t2_list = sorted(im_l, key = lambda imm:
+    hash_func(imm) - im_hash, reverse = False)
+    return t2_list
+    #return list(map(lambda x: x["name"], t2_list))
+    '''
+
+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
+
+with gr.Blocks(css="custom.css") as demo:
+    title = gr.HTML(
+            """<h1><img src="https://i.imgur.com/52VJ8vS.png" alt="SD"> StableDiffusion Search by Prompt order by Image</h1>""",
+            elem_id="title",
+    )
+
+    with gr.Row():
+        with gr.Column():
+            with gr.Row():
+                search_func_name = gr.Radio(choices=["lexica", "enterpix"],
+                    value="lexica", label="Search by", elem_id="search_radio")
+            with gr.Row():
+                #inputs = gr.Textbox(label = 'Enter prompt to search Lexica.art')
+                inputs = gr.Textbox(label="Prompt", show_label=False, lines=1, max_lines=20, min_width = 256,
+                placeholder="Enter prompt to search", elem_id="prompt")
+                #gr.Slider(label='Number of images ', minimum = 4, maximum = 20, step = 1, value = 4)]
+                text_button = gr.Button("Retrieve Images", elem_id="run_button")
+            #i = gr.Image(elem_id="result-image", label = "Image upload or selected", height = 768 - 256 - 32)
+            with gr.Row():
+                title = gr.Markdown(
+                        value="### Click on a Image in the gallery to select it",
+                        visible=True,
+                        elem_id="selected_model",
+                    )
+                choose_model = gr.Radio(choices=["0", "1", "2", "3"],
+                    value="0", label="Choose embedding layer", elem_id="layer_radio")
+                g_outputs = gr.Gallery(lable='Output gallery', elem_id="gallery",).style(grid=5,height=768 + 64 + 32,
+                    allow_preview=False, label = "retrieve Images")
+            with gr.Row():
+                with gr.Tab(label = "Download"):
+                    zip_button = gr.Button("Zip Images to Download", elem_id="zip_button")
+                    downloads = gr.File(label = "Image zipped", elem_id = "zip_file")
+
+        with gr.Column():
+            sdg_outputs = gr.Gallery(label='Sort Distinct gallery', elem_id="gallery",
+            columns=[5],object_fit="contain", height="auto")
+            sg_outputs = gr.Gallery(label='Sort gallery', elem_id="gallery",
+            columns=[5],object_fit="contain", height="auto")
+            #order_func_name = gr.Radio(choices=hash_func_name,
+            #value=hash_func_name[0], label="Order by", elem_id="order_radio")
+            #gr.Dataframe(label='prompts for corresponding images')]
+
+
+    with gr.Row():
+        '''
+        gr.Examples(
+            [
+            ["chinese zodiac signs", "lexica", "images/chinese_zodiac_signs.png", "average_hash"],
+            ["trending digital art", "lexica", "images/trending_digital_art.png", "colorhash"],
+            ["masterpiece, best quality, 1girl, solo, crop top, denim shorts, choker, (graffiti:1.5), paint splatter, arms behind back, against wall, looking at viewer, armband, thigh strap, paint on body, head tilt, bored, multicolored hair, aqua eyes, headset,", "lexica", "images/yuzu_girl0.png", "average_hash"],
+            ["beautiful home", "enterpix", "images/beautiful_home.png", "whash"],
+            ["interior design of living room", "enterpix", "images/interior_design_of_living_room.png", "whash"],
+            ["1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, hoop earrings, jewelry, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt",
+             "enterpix", "images/waifu_girl0.png", "phash"],
+            ],
+            inputs = [inputs, search_func_name, i, order_func_name],
+            label = "Examples"
+        )
+        '''
+        gr.Examples(
+            [
+            ["chinese zodiac signs", "lexica", ],
+            ["trending digital art", "lexica", ],
+            ["masterpiece, best quality, 1girl, solo, crop top, denim shorts, choker, (graffiti:1.5), paint splatter, arms behind back, against wall, looking at viewer, armband, thigh strap, paint on body, head tilt, bored, multicolored hair, aqua eyes, headset,", "lexica",],
+            ["beautiful home", "enterpix", ],
+            ["interior design of living room", "enterpix", ],
+            ["1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, hoop earrings, jewelry, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt",
+             "enterpix", ],
+            ],
+            inputs = [inputs, search_func_name,],
+            label = "Examples"
+        )
+
+
+    #outputs.select(image_click, outputs, i, _js="(x) => x.splice(0,x.length)")
+    #outputs.select(image_click, outputs, i,)
+    '''
+    i.change(
+                fn=swap_gallery,
+                inputs=[i, outputs, order_func_name],
+                outputs=outputs,
+                queue=False
+                )
+    order_func_name.change(
+                fn=swap_gallery,
+                inputs=[i, outputs, order_func_name],
+                outputs=outputs,
+                queue=False
+                )
+    '''
+
+
+    g_outputs.select(image_click,
+        inputs = [g_outputs, choose_model],
+        outputs = [sdg_outputs, sg_outputs],)
+
+    #### gr.Textbox().submit().success()
+
+    ### lexica
+    #text_button.click(lexica, inputs=inputs, outputs=outputs)
+    ### enterpix
+    #text_button.click(enterpix, inputs=inputs, outputs=outputs)
+    text_button.click(search, inputs=[inputs, search_func_name,], outputs=g_outputs)
+
+    zip_button.click(
+        zip_ims, inputs = sdg_outputs, outputs=downloads
+    )
+
+demo.launch("0.0.0.0")