Update app.py

app.py

@@ -7,6 +7,13 @@ from PIL import Image
 import imagehash
 import requests
 
+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)
@@ -20,24 +27,34 @@ def image_click(images, evt: gr.SelectData):
     img_selected = images[evt.index]
     return images[evt.index]['name']
 
-def swap_gallery(im, images):
+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 = imagehash.average_hash(Image.fromarray(im))
+    im_hash = hash_func(Image.fromarray(im))
     t2_list = sorted(images, key = lambda imm:
-    imagehash.average_hash(Image.open(imm["name"])) - im_hash, reverse = False)
+    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 ele: Image.open(
-                    requests.get(ele, stream = True).raw
-                    ), lex))
-    lex = list(map(lambda x: min_dim_to_size(x, ratio_size)[1], lex))
+    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"):
@@ -52,15 +69,37 @@ def enterpix(prompt, limit_size = 100, ratio_size = 256 + 128, use_key = "bigThu
     )
     resp = resp.json()
     resp = list(map(lambda x: x[use_key], resp["images"]))
-    resp = list(map(lambda ele: Image.open(
-                    requests.get(ele, stream = True).raw
-                    ), resp))
-    resp = list(map(lambda x: min_dim_to_size(x, ratio_size)[1], resp))
+    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
 
-with gr.Blocks() as demo:
+def search(prompt, search_name):
+    if search_name == "lexica":
+        return lexica(prompt)
+    else:
+        return enterpix(prompt)
+
+with gr.Blocks(css="custom.css") as demo:
+    title = gr.HTML(
+            """<h1><img src="https://i.imgur.com/52VJ8vS.png" alt="SD"> StableDiffusion Image Search </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,
@@ -75,7 +114,10 @@ with gr.Blocks() as demo:
                     visible=True,
                     elem_id="selected_model",
                 )
-            outputs = gr.Gallery(lable='Output gallery', elem_id="gallery",).style(grid=3,height=768,
+            order_func_name = gr.Radio(choices=hash_func_name,
+            value=hash_func_name[0], label="Order by", elem_id="order_radio")
+
+            outputs = gr.Gallery(lable='Output gallery', elem_id="gallery",).style(grid=5,height=768,
                     allow_preview=False, label = "retrieve Images")
             #gr.Dataframe(label='prompts for corresponding images')]
 
@@ -84,7 +126,13 @@ with gr.Blocks() as demo:
 
         i.change(
                 fn=swap_gallery,
-                inputs=[i, outputs],
+                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
                 )
@@ -93,6 +141,8 @@ with gr.Blocks() as demo:
         ### lexica
         #text_button.click(lexica, inputs=inputs, outputs=outputs)
         ### enterpix
-        text_button.click(enterpix, inputs=inputs, outputs=outputs)
+        #text_button.click(enterpix, inputs=inputs, outputs=outputs)
+        text_button.click(search, inputs=[inputs, search_func_name], outputs=outputs)
+
 
 demo.launch("0.0.0.0")