Update app.py

app.py

@@ -33,10 +33,52 @@ else:
 # Print some statistics
 print(f"Photos loaded: {len(photo_ids)}")
 
+from PIL import Image
 
-def search_by_text_and_photo(query_text, query_img, query_photo_id=None, photo_weight=0.5):
+
+def encode_search_query(net, search_query):
+    with torch.no_grad():
+        tokenized_query = clip.tokenize(search_query)
+        # print("tokenized_query: ", tokenized_query.shape)
+        # Encode and normalize the search query using CLIP
+        text_encoded = net.encode_text(tokenized_query.to(device))
+        text_encoded /= text_encoded.norm(dim=-1, keepdim=True)
+
+        # Retrieve the feature vector
+        # print("text_encoded: ", text_encoded.shape)
+        return text_encoded
+
+
+def find_best_matches(text_features, photo_features, photo_ids, results_count=5):
+    # Compute the similarity between the search query and each photo using the Cosine similarity
+    # print("text_features: ", text_features.shape)
+    # print("photo_features: ", photo_features.shape)
+    similarities = (photo_features @ text_features.T).squeeze(1)
+
+    # Sort the photos by their similarity score
+    best_photo_idx = (-similarities).argsort()
+    # print("best_photo_idx: ", best_photo_idx.shape)
+    # print("best_photo_idx: ", best_photo_idx[:results_count])
+
+    result_list = [photo_ids[i] for i in best_photo_idx[:results_count]]
+    # print("result_list: ", len(result_list))
+    # Return the photo IDs of the best matches
+    return result_list
+
+
+def search_unslash(net, search_query, photo_features, photo_ids, results_count=10):
+    # Encode the search query
+    text_features = encode_search_query(net, search_query)
+
+    # Find the best matches
+    best_photo_ids = find_best_matches(text_features, photo_features, photo_ids, results_count)
+
+    return best_photo_ids
+
+
+def search_by_text_and_photo(query_text, query_photo=None, query_photo_id=None, photo_weight=0.5):
     # Encode the search query
-    if not query_text and not query_photo_id:
+    if not query_text and query_photo is None and not query_photo_id:
         return []
 
     text_features = encode_search_query(model, query_text)
@@ -53,8 +95,12 @@ def search_by_text_and_photo(query_text, query_img, query_photo_id=None, photo_w
         # Find the best match
         best_photo_ids = find_best_matches(search_features, photo_features, photo_ids, 10)
 
-    elif query_img:
-        query_photo_features = model.encode_image(query_img)
+    elif query_photo is not None:
+        query_photo = preprocess(query_photo)
+        query_photo = torch.tensor(query_photo).permute(2, 0, 1)
+
+        print(query_photo.shape)
+        query_photo_features = model.encode_image(query_photo)
         query_photo_features = query_photo_features / query_photo_features.norm(dim=1, keepdim=True)
 
         # Combine the test and photo queries and normalize again
@@ -66,7 +112,7 @@ def search_by_text_and_photo(query_text, query_img, query_photo_id=None, photo_w
     else:
         # Display the results
         print("Test search result")
-        best_photo_ids = search_unslash(query_text, photo_features, photo_ids, 10)
+        best_photo_ids = search_unslash(model, query_text, photo_features, photo_ids, 10)
 
     return best_photo_ids
 
@@ -76,20 +122,21 @@ with gr.Blocks() as app:
         gr.Markdown(
             """
             # CLIP Image Search Engine!
-            ### Enter search query or/and input image to find the similar images from the database -
+            ### Enter search query or/and select image to find the similar images
             """)
 
     with gr.Row(visible=True):
         with gr.Column():
             with gr.Row():
-                search_text = gr.Textbox(value='', placeholder='Search..', label='Enter Your Query')
+                search_text = gr.Textbox(value='', placeholder='Search..', label='Enter search query')
 
             with gr.Row():
                 submit_btn = gr.Button("Submit", variant='primary')
                 clear_btn = gr.ClearButton()
 
-        with gr.Column():
-            search_image = gr.Image(label='Upload Image or Select from results')
+        with gr.Column(visible=True) as input_image_col:
+            search_image = gr.Image(label='Select from results', interactive=False)
+            search_image_id = gr.State(None)
 
     with gr.Row(visible=True):
         output_images = gr.Gallery(allow_preview=False, label='Results.. ', info='',
@@ -102,44 +149,75 @@ with gr.Blocks() as app:
         return {
             search_image: None,
             output_images: None,
-            search_text: None
+            search_text: None,
+            search_image_id: None,
+            input_image_col: gr.update(visible=True)
         }
 
 
-    clear_btn.click(clear_data, None, [search_image, output_images, search_text])
+    clear_btn.click(clear_data, None, [search_image, output_images, search_text, search_image_id, input_image_col])
 
 
     def on_select(evt: gr.SelectData, output_image_ids):
         return {
-            search_image: f"https://unsplash.com/photos/{output_image_ids[evt.index]}/download?w=100"
+            search_image: f"https://unsplash.com/photos/{output_image_ids[evt.index]}/download?w=320",
+            search_image_id: output_image_ids[evt.index],
+            input_image_col: gr.update(visible=True)
         }
 
 
-    output_images.select(on_select, output_image_ids, search_image)
+    output_images.select(on_select, output_image_ids, [search_image, search_image_id, input_image_col])
 
 
-    def func_search(query, img):
-        best_photo_ids = search_by_text_and_photo(query, img)
-        img_urls = []
-        for p_id in best_photo_ids:
-            url = f"https://unsplash.com/photos/{p_id}/download?w=100"
-            img_urls.append(url)
+    def func_search(query, img, img_id):
+        best_photo_ids = []
+        if img_id:
+            best_photo_ids = search_by_text_and_photo(query, query_photo_id=img_id)
+        elif img is not None:
+            img = Image.open(img)
+            best_photo_ids = search_by_text_and_photo(query, query_photo=img)
+        elif query:
+            best_photo_ids = search_by_text_and_photo(query)
 
-        valid_images = filter_invalid_urls(img_urls, best_photo_ids)
+        if len(best_photo_ids) == 0:
+            print("Invalid Search Request")
+            return {
+                output_image_ids: [],
+                output_images: []
+            }
+        else:
+            img_urls = []
+            for p_id in best_photo_ids:
+                url = f"https://unsplash.com/photos/{p_id}/download?w=20"
+                img_urls.append(url)
 
-        return {
-            output_image_ids: valid_images['image_ids'],
-            output_images: valid_images['image_urls']
-        }
+            valid_images = filter_invalid_urls(img_urls, best_photo_ids)
+
+            return {
+                output_image_ids: valid_images['image_ids'],
+                output_images: valid_images['image_urls']
+            }
 
 
     submit_btn.click(
         func_search,
-        [search_text, search_image],
+        [search_text, search_image, search_image_id],
         [output_images, output_image_ids]
     )
 
+
+    def on_upload(evt: gr.SelectData):
+        return {
+            search_image_id: None
+        }
+
+
+    search_image.upload(on_upload, None, search_image_id)
+
 '''
 Launch the app
 '''
 app.launch()
+
+
+