patricebechard's picture
add file for streamlit app
9a3a339 unverified
import requests
import streamlit as st
from PIL import Image
from transformers import CLIPModel, CLIPProcessor
st.title("Most Similar Kanye West Song")
# read lyrics from all kanye west songs
with open("kanye_lyrics.txt", "r") as f:
all_lyrics = f.read().split("\n\n\n\n")
# load model and processor
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
# get user input
url = st.text_input("Enter an image url: ", "https://i.redd.it/0vsrxk3gcfr91.jpg")
# get image from url
image = Image.open(requests.get(url, stream=True).raw)
st.image(url, width=700)
# get the image-text similarity score
inputs = processor(text=all_lyrics, images=image, return_tensors="pt", padding=True, truncation=True)
outputs = model(**inputs)
logits_per_image = outputs.logits_per_image # this is the image-text similarity score
probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities
# get the most similar kanye song
most_similar_song = all_lyrics[probs.argmax().item()]
st.write(most_similar_song)