danielalonk's picture
Update app.py
9525457 verified
Raw
History Blame Contribute Delete
5.73 kB
import torch
import torch.nn.functional as F
import numpy as np
import pandas as pd
import gradio as gr
from transformers import CLIPModel, CLIPTokenizer
from sklearn.metrics.pairwise import cosine_similarity
from datasets import load_dataset
from PIL import Image
import traceback
MODEL_ID = "openai/clip-vit-base-patch32"
tokenizer = CLIPTokenizer.from_pretrained(MODEL_ID)
model = CLIPModel.from_pretrained(MODEL_ID)
model.eval()
embeddings = np.load("car_embeddings.npz")["embeddings"]
metadata = pd.read_csv("car_metadata.csv")
ds = load_dataset("DamianBoborzi/car_images", split="train").select(range(10000))
top3_default_indices = (
metadata
.nlargest(3, "aesthetic_score")["dataset_index"]
.astype(int)
.tolist()
)
def get_default_images():
imgs = []
for ds_idx in top3_default_indices:
img = ds[ds_idx]["image"]
if not isinstance(img, Image.Image):
img = Image.fromarray(img)
imgs.append(img.convert("RGB"))
return imgs
def get_text_vector(user_query: str) -> np.ndarray:
inputs = tokenizer(user_query, return_tensors="pt",
padding=True, truncation=True, max_length=77)
with torch.no_grad():
text_out = model.text_model(**inputs)
feat = model.text_projection(text_out.pooler_output)
return F.normalize(feat, p=2, dim=-1).cpu().numpy()
def recommend(user_query: str):
if not user_query or not user_query.strip():
defaults = get_default_images()
rows = [metadata[metadata["dataset_index"] == i].iloc[0]
for i in top3_default_indices]
results = []
for rank, (img, row) in enumerate(zip(defaults, rows)):
score = float(row["aesthetic_score"])
car_name = str(row["text"]).split(",")[0].split("(")[0].strip()[:55]
label = f"#{rank+1} Aesthetic Score: {score:.2f}\n{car_name}"
results += [img, label]
return results[0], results[1], results[2], results[3], results[4], results[5]
try:
query_vec = get_text_vector(user_query)
scores = cosine_similarity(query_vec, embeddings).flatten()
top3 = scores.argsort()[-3:][::-1]
results = []
for rank, df_idx in enumerate(top3):
row = metadata.iloc[int(df_idx)]
score = float(scores[int(df_idx)])
ds_idx = int(row["dataset_index"])
img = ds[ds_idx]["image"]
if not isinstance(img, Image.Image):
img = Image.fromarray(img)
img = img.convert("RGB")
car_name = str(row["text"]).split(",")[0].split("(")[0].strip()[:55]
label = f"#{rank+1} Score: {score:.4f} | Cluster {int(row['cluster'])}\n{car_name}"
results += [img, label]
return results[0], results[1], results[2], results[3], results[4], results[5]
except Exception as e:
traceback.print_exc()
msg = f"Error: {e}"
return None, msg, None, msg, None, msg
default_imgs = get_default_images()
default_rows = [metadata[metadata["dataset_index"] == i].iloc[0]
for i in top3_default_indices]
default_labels = [
f"#{r+1} Aesthetic Score: {float(row['aesthetic_score']):.2f}\n"
f"{str(row['text']).split(',')[0].split('(')[0].strip()[:55]}"
for r, row in enumerate(default_rows)
]
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown("# 🏎️ Car Recommendation Engine")
gr.Markdown(
"Describe the car you're looking for and CLIP will find the 3 most similar cars "
"from a dataset of **10,000 real car images**.\n\n"
"*Default view shows the top 3 highest-rated cars by aesthetic score.*"
)
with gr.Row():
query_box = gr.Textbox(
label="Describe your ideal car",
placeholder="e.g. A red aggressive supercar with wide body kit",
scale=5
)
search_btn = gr.Button("🔍 Find Matches", variant="primary", scale=1)
gr.Markdown("### 🏆 Top 3 Matches")
with gr.Row():
with gr.Column():
img1 = gr.Image(label="Match 1", type="pil", value=default_imgs[0])
lbl1 = gr.Textbox(show_label=False, interactive=False, value=default_labels[0])
with gr.Column():
img2 = gr.Image(label="Match 2", type="pil", value=default_imgs[1])
lbl2 = gr.Textbox(show_label=False, interactive=False, value=default_labels[1])
with gr.Column():
img3 = gr.Image(label="Match 3", type="pil", value=default_imgs[2])
lbl3 = gr.Textbox(show_label=False, interactive=False, value=default_labels[2])
search_btn.click(fn=recommend, inputs=query_box,
outputs=[img1, lbl1, img2, lbl2, img3, lbl3])
gr.Markdown("---")
gr.Markdown("## 📽️ Project Presentation Video")
with gr.Row():
gr.HTML("""
<div style="display: flex; justify-content: center; align-items: center; width: 100%;">
<iframe width="760" height="415"
src="https://www.youtube.com/embed/q1pvsyz1A9g"
title="YouTube video player"
frameborder="0"
allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture"
allowfullscreen>
</iframe>
</div>
""")
gr.Examples(
examples=[
["A red aggressive supercar with wide body kit"],
["Vintage blue convertible classic car"],
["White family SUV with clean modern design"],
["Black 3D rendered sports coupe low angle"],
],
inputs=query_box,
)
demo.launch()