MAY199's picture
Update app.py
26ed0d6 verified
import gradio as gr
import numpy as np
import pandas as pd
import torch
from datasets import load_dataset
from transformers import CLIPProcessor, CLIPModel
# Load dataset
ds = load_dataset("amaye15/landscapes")
train_ds = ds["train"]
# Load embeddings
df = pd.read_parquet("image_embeddings_clip.parquet")
image_indices = df["image_index"].values
emb_matrix = df.drop(columns=["image_index"]).values.astype(np.float32)
# Load CLIP
device = "cuda" if torch.cuda.is_available() else "cpu"
model_name = "openai/clip-vit-base-patch32"
processor = CLIPProcessor.from_pretrained(model_name)
model = CLIPModel.from_pretrained(model_name).to(device)
model.eval()
def l2_normalize(x):
return x / np.linalg.norm(x)
@torch.no_grad()
def embed_image(img):
inputs = processor(images=img, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
feats = model.get_image_features(**inputs)
feats = feats / feats.norm(dim=-1, keepdim=True)
return feats.squeeze(0).cpu().numpy()
def recommend(img):
q_emb = embed_image(img)
sims = emb_matrix @ l2_normalize(q_emb)
top = np.argsort(-sims)[1:4]
results = []
for i in top:
results.append(train_ds[int(image_indices[i])]["pixel_values"])
return results
# Gradio interface
demo = gr.Interface(
fn=recommend,
inputs=gr.Image(type="pil", label="Upload a landscape image"),
outputs=[
gr.Image(label="Recommendation 1"),
gr.Image(label="Recommendation 2"),
gr.Image(label="Recommendation 3"),
],
title="Landscape Image Recommendation System",
description="Upload a landscape image and receive visually similar recommendations."
)
# App layout with video below
with gr.Blocks() as app:
demo.render()
gr.Markdown("---")
gr.Markdown("## Project Presentation Video")
gr.HTML("""
<div style="width:100%;max-width:900px;margin:0 auto;">
<div style="position:relative;padding-bottom:56.25%;height:0;overflow:hidden;border-radius:12px;">
<iframe
src="https://www.youtube.com/embed/nwbPAR7UApw"
title="Project presentation"
frameborder="0"
allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share"
allowfullscreen
style="position:absolute;top:0;left:0;width:100%;height:100%;">
</iframe>
</div>
</div>
""")
app.launch()