Spaces:
Sleeping
Sleeping
File size: 3,356 Bytes
2198293 5064827 2198293 5064827 2198293 8ef0614 2198293 5064827 2198293 8ef0614 2198293 8ef0614 2198293 8ef0614 2198293 cb0343d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | import gradio as gr
import numpy as np
import pandas as pd
import torch
from transformers import CLIPProcessor, CLIPModel
from sklearn.metrics.pairwise import cosine_similarity
from datasets import load_dataset
# Load model
device = "cuda" if torch.cuda.is_available() else "cpu"
model_name = "openai/clip-vit-base-patch32"
model = CLIPModel.from_pretrained(model_name).to(device)
processor = CLIPProcessor.from_pretrained(model_name)
# Load embeddings
df = pd.read_parquet("nature_embeddings.parquet")
EMBEDDINGS_MATRIX = np.array(df['embedding'].tolist())
CAPTIONS = df['caption'].tolist()
ORIGINAL_INDICES = df['original_index'].tolist()
# Load dataset
dataset = load_dataset("mertcobanov/nature-dataset")
train_data = dataset['train']
def classify_scene(caption):
"""Classify nature scene based on caption keywords"""
caption = caption.lower()
if any(w in caption for w in ['beach', 'ocean', 'sea', 'coast', 'wave']):
return '🌊 Coastal'
elif any(w in caption for w in ['waterfall', 'stream', 'river', 'lake']):
return '💧 Water'
elif any(w in caption for w in ['mountain', 'peak', 'cliff', 'rock']):
return '🏔️ Mountain'
elif any(w in caption for w in ['desert', 'dune', 'sand', 'arid']):
return '🏜️ Desert'
elif any(w in caption for w in ['forest', 'tree', 'jungle', 'wood']):
return '🌲 Forest'
else:
return '🌿 Nature'
def get_image_embedding(image):
inputs = processor(images=image, return_tensors="pt", padding=True).to(device)
with torch.no_grad():
outputs = model.vision_model(**inputs)
features = outputs.pooler_output
features = torch.nn.functional.normalize(features, p=2, dim=-1)
return features.cpu().numpy()
def recommend(image):
query_embedding = get_image_embedding(image)
similarities = cosine_similarity(query_embedding, EMBEDDINGS_MATRIX)[0]
top_indices = np.argsort(similarities)[::-1]
top_indices = [idx for idx in top_indices if similarities[idx] < 0.9999][:3]
results = []
for idx in top_indices:
original_idx = ORIGINAL_INDICES[idx]
img = train_data[int(original_idx)]['image']
cap = CAPTIONS[idx]
sim = float(similarities[idx])
scene = classify_scene(cap)
results.append((img, f"{scene} | {cap} (similarity: {sim:.4f})"))
return results[0][0], results[0][1], results[1][0], results[1][1], results[2][0], results[2][1]
# Gradio interface
with gr.Blocks(title="Nature Scene Recommender") as demo:
gr.Markdown("# 🌿 Nature Scene Recommender")
gr.Markdown("Upload a nature image and get the 3 most similar scenes!")
with gr.Row():
input_image = gr.Image(type="pil", label="Upload Nature Image")
btn = gr.Button("Find Similar Scenes 🔍", variant="primary")
gr.Markdown("### Top 3 Similar Scenes")
with gr.Row():
img1 = gr.Image(label="Match #1")
img2 = gr.Image(label="Match #2")
img3 = gr.Image(label="Match #3")
with gr.Row():
cap1 = gr.Textbox(label="Scene #1")
cap2 = gr.Textbox(label="Scene #2")
cap3 = gr.Textbox(label="Scene #3")
btn.click(fn=recommend, inputs=input_image,
outputs=[img1, cap1, img2, cap2, img3, cap3])
demo.launch(server_name="0.0.0.0", server_port=7860) |