''' !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 ''' print(g[0]) print(g[0][0]) print(g[0][1]) ''' g = json.loads(g.model_dump_json()) g = list(map(lambda x: {"name": x["image"]["path"]}, g)) 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( """

SD StableDiffusion Search by Prompt order by Image Embedding

""", 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(): with gr.Column(): 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") with gr.Row(): g_outputs = gr.Gallery(label='Output gallery', elem_id="gallery", columns=[5],object_fit="contain", height="auto") 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(): 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.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 ink painting", "lexica", ], ["silk road", "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")