app

app.py

@@ -1,12 +1,9 @@
-from fastapi import FastAPI, File, UploadFile
+import gradio as gr
 from PIL import Image
-from torchvision import transforms
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
 from transformers import CLIPProcessor, CLIPModel
 from datasets import load_dataset
 import torch
-import io
-
-app = FastAPI()
 
 # Load the pre-trained CLIP model and its tokenizer
 model_name = "openai/clip-vit-base-patch32"
@@ -17,59 +14,54 @@ model = CLIPModel.from_pretrained(model_name)
 dataset = load_dataset("ashraq/fashion-product-images-small")
 deepfashion_database = dataset["train"]
 
+# Define the preprocessing function for images
 def preprocess_image(image):
-    pil_image = Image.fromarray(image.astype('uint8'), 'RGB')
-    preprocess = transforms.Compose([
-        transforms.Resize(224),
-        transforms.CenterCrop(224),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    preprocess = Compose([
+        Resize(256, interpolation=Image.BICUBIC),
+        CenterCrop(224),
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
     ])
-    return preprocess(pil_image).unsqueeze(0)
-
-def encode_text(text):
-    inputs = processor(text=[text], return_tensors="pt", padding=True, truncation=True)
-    return inputs
-
-def encode_image(image):
-    inputs = processor(images=[image], return_tensors="pt", padding=True, truncation=True)
-    return inputs
+    return preprocess(image).unsqueeze(0)
 
-def calculate_similarities(query_image, query_text):
-    query_image_features = model.get_image_features(query_image)
-    query_text_features = model.get_text_features(query_text)
-
-    similarities = []
-    for product in deepfashion_database:
-        product_image_features = torch.Tensor(product["image_features"])
-        product_text_features = torch.Tensor(product["text_features"])
+# Modify initial_query function to return the same input the user gives 3 times
+def initial_query(image, text):
+    input_product = {"description": text, "image_path": None}
+    top_3_products = [(input_product, 1), (input_product, 1), (input_product, 1)]
+    return top_3_products
 
-        image_similarity = torch.nn.CosineSimilarity(dim=-1)(query_image_features, product_image_features)
-        text_similarity = torch.nn.CosineSimilarity(dim=-1)(query_text_features, product_text_features)
-        
-        similarity_score = image_similarity * text_similarity
-        similarities.append(similarity_score)
+# Keep the rest of the code unchanged
 
-    return similarities
+def generate_output_html(products):
+    html_output = "<ol>"
+    for product in products:
+        html_output += f'<li>{product[0]["description"]}</li>'
+    html_output += "</ol>"
+    return html_output
 
-def initial_query(image, text):
-    query_image = encode_image(image)
-    query_text = encode_text(text)
+def initial_query_wrapper(image, text):
+    top_3_products = initial_query(image, text)
+    return generate_output_html(top_3_products),
 
-    similarities = calculate_similarities(query_image, query_text)
-    sorted_indices = sorted(range(len(similarities)), key=lambda i: similarities[i], reverse=True)
-    top_3_indices = sorted_indices[:3]
+def product_search_wrapper(image=None, text=None, selected_product_index=None, additional_text=None):
+    if image is not None or text is not None:
+        top_3_products = initial_query(image, text)
+        return generate_output_html(top_3_products),
+    else:
+        return "",
 
-    top_3_products = [deepfashion_database[i] for i in top_3_indices]
-    return top_3_products
+iface = gr.Interface(
+    fn=product_search_wrapper,
+    inputs=[
+        gr.inputs.Image(optional=True),
+        gr.inputs.Textbox(lines=3, label="Initial Text Query", optional=True),
+    ],
+    outputs=[
+        gr.outputs.HTML(label="Results")
+    ],
+    title="Product Search",
+    description="Find the best matching products using images and text queries.",
+    layout="vertical"
+)
 
-@app.post("/initial_query/")
-async def api_initial_query(text: str, image: UploadFile = File(None)):
-    if image is not None:
-        image_content = await image.read()
-        image = Image.open(io.BytesIO(image_content))
-        image = preprocess_image(image)
-    else:
-        image = None
-    top_3_products = initial_query(image, text)
-    return {"top_3_products": top_3_products}
+iface.launch()