Spaces:
Runtime error
Runtime error
| import streamlit as st | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| from sentence_transformers import SentenceTransformer | |
| from transformers import CLIPTokenizer, CLIPModel | |
| from scipy.spatial import distance | |
| import numpy as np | |
| import torch | |
| import json | |
| import random | |
| model_id = "openai/clip-vit-base-patch32" | |
| clip_model = CLIPModel.from_pretrained(model_id) | |
| clip_tokenizer = CLIPTokenizer.from_pretrained(model_id) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # move the model to the device | |
| clip_model.to(device) | |
| def get_single_text_embedding(text): | |
| inputs = clip_tokenizer(text, return_tensors = "pt") | |
| text_embeddings = clip_model.to(device).get_text_features(**inputs.to(device)) | |
| embedding_as_np = text_embeddings.cpu().detach().numpy() | |
| return embedding_as_np | |
| def get_top_N_images(query, image_vectors, top_K=4): | |
| query_vect = get_single_text_embedding(query) | |
| data = cosine_similarity(query_vect, image_vectors)[0] | |
| most_similar_images = sorted(range(len(data)), key=lambda i: data[i], reverse=True)[:top_K] | |
| return most_similar_images | |
| image_data = [] | |
| with open('Property-images.jsonl', 'r', encoding="utf-8") as f: | |
| for line in f: | |
| image_data.append(json.loads(line)) | |
| with open('property-image-vectors.npy', 'rb') as f: | |
| image_vectors = np.load(f) | |
| def combine(data): | |
| new_data = [] | |
| window = 4 # number of sentences to combine | |
| stride = 1 # number of sentences to 'stride' over, used to create overlap | |
| for i in range(0, len(data), stride): | |
| i_end = min(len(data)-1, i+window) | |
| if data[i]['title'] != data[i_end]['title']: | |
| # in this case we skip this entry as we have start/end of two videos | |
| continue | |
| text = ' ' | |
| for x in data[i:i_end]: | |
| text += x['text'] | |
| new_data.append({ | |
| 'start': data[i]['start'], | |
| 'end': data[i_end]['end'], | |
| 'title': data[i]['title'], | |
| 'text': text, | |
| 'id': data[i]['id'], | |
| 'url': data[i]['url'], | |
| 'published': data[i]['published'] | |
| }) | |
| return new_data | |
| model_id = "multi-qa-mpnet-base-dot-v1" | |
| model = SentenceTransformer(model_id) | |
| meta_data = [] | |
| with open('Property-transcription.jsonl', 'r', encoding="utf-8") as f: | |
| for line in f: | |
| meta_data.append(json.loads(line)) | |
| meta_data = combine(meta_data) | |
| with open('property-vectors.npy', 'rb') as f: | |
| text_vector = np.load(f) | |
| def card(thumbnail: str, title: str, urls: list, contexts: list, starts: list, ends: list): | |
| meta = [(e, s, u, c) for e, s, u, c in zip(ends, starts, urls, contexts)] | |
| meta.sort(reverse=False) | |
| text_content = [] | |
| current_start = 0 | |
| current_end = 0 | |
| for end, start, url, context in meta: | |
| # reformat seconds to timestamp | |
| time = start / 60 | |
| mins = f"0{int(time)}"[-2:] | |
| secs = f"0{int(round((time - int(mins))*60, 0))}"[-2:] | |
| timestamp = f"{mins}:{secs}" | |
| if start < current_end and start > current_start: | |
| # this means it is a continuation of the previous sentence | |
| text_content[-1][0] = text_content[-1][0].split(context[:10])[0] | |
| text_content.append([f"[{timestamp}] {context.capitalize()}", url]) | |
| else: | |
| text_content.append(["xxLINEBREAKxx", ""]) | |
| text_content.append([f"[{timestamp}] {context}", url]) | |
| current_start = start | |
| current_end = end | |
| html_text = "" | |
| for text, url in text_content: | |
| if text == "xxLINEBREAKxx": | |
| html_text += "<br>" | |
| else: | |
| html_text += f"<small><a href={url}>{text.strip()}... </a></small><br>" | |
| print(html_text) | |
| html = f""" | |
| <div class="container-fluid"> | |
| <div class="row align-items-start"> | |
| <div class="col-md-4 col-sm-4"> | |
| <div class="position-relative"> | |
| <a href={urls[0]}><img src={thumbnail} class="img-fluid" style="width: 192px; height: 106px"></a> | |
| </div> | |
| </div> | |
| <div class="col-md-8 col-sm-8"> | |
| <h2>{title}</h2> | |
| </div> | |
| <div> | |
| {html_text} | |
| <br><br> | |
| """ | |
| return st.markdown(html, unsafe_allow_html=True) | |
| st.write(""" | |
| # Property Video Search 🏘️ | |
| Utilize AI to quickly locate amenities and features within your residence.🤗 | |
| """) | |
| st.markdown(""" | |
| <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/bootstrap@4.0.0/dist/css/bootstrap.min.css" integrity="sha384-Gn5384xqQ1aoWXA+058RXPxPg6fy4IWvTNh0E263XmFcJlSAwiGgFAW/dAiS6JXm" crossorigin="anonymous"> | |
| """, unsafe_allow_html=True) | |
| query = st.text_input("Search!", "") | |
| if query != "": | |
| vector1 = model.encode(query) | |
| text_cosines = dict() | |
| for i,vec in enumerate(text_vector): | |
| text_cosines[str(i)+'-text'] = 1 - distance.cosine(vector1, vec) | |
| if random.randint(0, 1): | |
| vector1 = get_single_text_embedding(query) | |
| image_cosines = dict() | |
| for i,vec in enumerate(image_vectors): | |
| image_cosines[str(i)+'-image'] = 1 - distance.cosine(vector1.reshape((512,)), vec) | |
| text_cosines.update(image_cosines) | |
| sorted_cosines = sorted(text_cosines.items(), key=lambda x:x[1], reverse=True) | |
| converted_dict = dict(sorted_cosines) | |
| results = {} | |
| order = [] | |
| for vec_index in list(converted_dict.keys())[:7]: | |
| idx = int(vec_index.split('-')[0]) | |
| video_id = image_data[idx]['url'].split('/')[-1] | |
| if vec_index.split('-')[-1] == 'image': | |
| if video_id not in results: | |
| results[video_id] = { | |
| 'title': image_data[idx]['title'], | |
| 'urls': [f"{image_data[idx]['url']}?t={int(image_data[idx]['sec'])}"], | |
| 'contexts': ['Image-query'], | |
| 'starts': [int(image_data[idx]['sec'])], | |
| 'ends': [int(image_data[idx]['sec']+6)] | |
| } | |
| order.append(video_id) | |
| else: | |
| results[video_id]['urls'].append( | |
| f"{image_data[idx]['url']}?t={int(image_data[idx]['sec'])}" | |
| ) | |
| results[video_id]['contexts'].append('Image-query') | |
| results[video_id]['starts'].append(int(image_data[idx]['sec'])) | |
| results[video_id]['ends'].append(int(image_data[idx]['sec']+6)) | |
| elif vec_index.split('-')[-1] == 'text': | |
| if video_id not in results: | |
| results[video_id] = { | |
| 'title': meta_data[idx]['title'], | |
| 'urls': [f"{meta_data[idx]['url']}?t={int(meta_data[idx]['start'])}"], | |
| 'contexts': [meta_data[idx]['text']], | |
| 'starts': [int(meta_data[idx]['start'])], | |
| 'ends': [int(meta_data[idx]['end'])] | |
| } | |
| order.append(video_id) | |
| else: | |
| results[video_id]['urls'].append( | |
| f"{meta_data[idx]['url']}?t={int(meta_data[idx]['start'])}" | |
| ) | |
| results[video_id]['contexts'].append( | |
| meta_data[idx]['text'] | |
| ) | |
| results[video_id]['starts'].append(int(meta_data[idx]['start'])) | |
| results[video_id]['ends'].append(int(meta_data[idx]['end'])) | |
| # now display cards | |
| for video_id in order: | |
| card( | |
| thumbnail=f"https://img.youtube.com/vi/{video_id}/maxresdefault.jpg", | |
| title=results[video_id]['title'], | |
| urls=results[video_id]['urls'], | |
| contexts=results[video_id]['contexts'], | |
| starts=results[video_id]['starts'], | |
| ends=results[video_id]['ends'] | |
| ) | |
| else: | |
| st.warning('💡Try searching: huge balcony, swimming pool, spiral staircase, panoramic view ... etc') | |