Add application file

app.py

@@ -0,0 +1,121 @@
+import os
+import sys
+import gradio as gr
+import numpy as np
+from PIL import Image
+from embedding import get_device, get_model_and_processor, to_embedding
+from datasets import Dataset, load_dataset
+
+def initialize_model():
+    device = get_device()
+    model, processor = get_model_and_processor("patrickjohncyh/fashion-clip", device)
+    ref_dataset = load_dataset("HadrienCr/embeddeDior", split="train")
+    ref_dataset = ref_dataset.add_faiss_index("embeddings")
+    return model, processor, ref_dataset, device
+
+def search(
+    image: np.ndarray,
+    reference_dataset: Dataset,
+    model,
+    processor,
+    device: str,
+    num_neighbors: int = 4,
+):
+    """a function that embeds a new image and returns the most probable results"""
+    scores, retrieved_examples = reference_dataset.get_nearest_examples(
+        "embeddings",
+        to_embedding(np.expand_dims(image, 0), processor, model, device),
+        k=num_neighbors,
+    )
+    return retrieved_examples
+
+def process_image(image, num_results, remove_bg, model, processor, ref_dataset, device):
+    if image is None:
+        return []  # Return an empty list if no image is provided
+    
+    # Ensure the input image is a numpy array
+    if isinstance(image, Image.Image):
+        image = np.array(image)
+    
+    # Handle background removal
+    if remove_bg:
+        from rembg import remove
+        image = remove(image)[:,:,0:3]
+    # Perform the search
+    results = search(
+        image,
+        ref_dataset,
+        model,
+        processor,
+        device,
+        num_neighbors=num_results
+    )
+    images = results['image']
+    paths = results['path']
+    # Prepare the output
+    output_images = []
+
+    for img, path in zip(images, paths):
+        output_images.append((np.array(img), os.path.basename(path)))
+    
+    return output_images  # Return the list of tuples
+
+
+def main():
+    print("Initializing model and loading reference dataset...")
+    model, processor, ref_dataset, device = initialize_model()
+    print("Initialization complete!")
+
+    # Path to the examples folder
+    examples_folder = "examples/"
+    example_files = [
+        os.path.join(examples_folder, fname) 
+        for fname in os.listdir(examples_folder) 
+        if fname.lower().endswith(('png', 'jpg', 'jpeg'))
+    ]
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# Image Retrieval System")
+        gr.Markdown("Upload an image to find similar images in the reference dataset.")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image = gr.Image(
+                    label="Upload Image",
+                    type="pil"  # Changed to PIL format
+                )
+                num_results = gr.Slider(
+                    minimum=1, 
+                    maximum=10, 
+                    value=5, 
+                    step=1, 
+                    label="Number of results"
+                )
+                remove_bg = gr.Checkbox(label="Remove Background")
+                submit_btn = gr.Button("Search")
+            
+            with gr.Column(scale=2):
+                gallery = gr.Gallery(
+                    label="Retrieved Images",
+                    show_label=True,
+                    columns=3,
+                    object_fit="contain"
+                )
+
+        # Add the Examples component
+        gr.Examples(
+            examples=example_files,
+            inputs=input_image,
+            label="Example Images"
+        )
+
+        submit_btn.click(
+            fn=lambda img, num, bg: process_image(img, num, bg, model, processor, ref_dataset, device),
+            inputs=[input_image, num_results, remove_bg],
+            outputs=gallery
+        )
+
+    demo.launch(share=True)
+
+if __name__ == "__main__":
+    main()