Spaces:
Runtime error
Runtime error
File size: 2,107 Bytes
45e49d7 3e6bee3 68d068c 45e49d7 0a35212 4a73a40 0a35212 57a597e 45e49d7 f49cac7 45e49d7 |
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 |
import streamlit as st
import random
import os
import clip
from funcs.get_similarity import get_similarity_score, create_filelist, load_embeddings, find_matches
from funcs.fiass_similaruty import load_embeddings, encode_text, find_matches_fiass
import torch.nn.functional as F
import pandas as pd
import zipfile
from unzip import unzip
zip_filename = 'img.zip'
target_dir = 'photos'
# Create the target directory if it doesn't exist
if not os.path.exists(target_dir):
os.makedirs(target_dir)
# Open the zip file and extract its contents to the target directory
with zipfile(zip_filename, 'r') as zip_ref:
zip_ref.extractall(target_dir)
device = 'cpu'
model_path = "weights/ViT-B-32.pt"
model, preprocess = clip.load('ViT-B/32', device)
file_name = create_filelist('img')
features = load_embeddings('embeddings/emb_images_5000.npy')
df = pd.read_csv('data/results.csv')
random_queries = ['friends playing cards', 'rock band playing on guitars', 'policeman cross the road',
'sleeping kids', 'football team playing on the grass' , 'learning programming'
]
st.header('Find my pic!')
request = st.text_input('Write a description of the picture', ' Two people at the photo')
if request:
unzip('img.zip')
img_count = st.slider('How much pic you need?', 4, 8, 6, 2)
matches = find_matches_fiass(features, request, file_name, img_count)
row1, row2 = st.columns(2)
if st.button('Find!'):
selected_filenames = matches
for i in range(int(img_count/2)):
filename = selected_filenames[i]
img_path = filename
img_discription = df[df['image_name'] == filename.split('/')[1]]['comment'].iloc[0]
with row1:
st.image(img_path, width=300, caption=img_discription)
# display next 3 images in the second row
for i in range(int(img_count/2), img_count):
filename = selected_filenames[i]
img_path = filename
img_discription = df[df['image_name'] == filename.split('/')[1]]['comment'].iloc[0]
with row2:
st.image(img_path, width=300, caption=img_discription) |