Spaces:
Running
Running
| from html import escape | |
| import requests | |
| from io import BytesIO | |
| import base64 | |
| from multiprocessing.dummy import Pool | |
| from PIL import Image, ImageDraw | |
| import streamlit as st | |
| import pandas as pd, numpy as np | |
| import torch | |
| from transformers import CLIPProcessor, CLIPModel | |
| from transformers import OwlViTProcessor, OwlViTForObjectDetection | |
| from transformers.image_utils import ImageFeatureExtractionMixin | |
| import tokenizers | |
| DEBUG = False | |
| if DEBUG: | |
| MODEL = "vit-base-patch32" | |
| else: | |
| MODEL = "vit-large-patch14-336" | |
| CLIP_MODEL = f"openai/clip-{MODEL}" | |
| OWL_MODEL = f"google/owlvit-base-patch32" | |
| if not DEBUG and torch.cuda.is_available(): | |
| device = torch.device("cuda") | |
| else: | |
| device = torch.device("cpu") | |
| HEIGHT = 200 | |
| N_RESULTS = 6 | |
| color = st.get_option("theme.primaryColor") | |
| if color is None: | |
| color = (255, 196, 35) | |
| else: | |
| color = tuple(int(color.lstrip("#")[i : i + 2], 16) for i in (0, 2, 4)) | |
| def load(): | |
| df = {0: pd.read_csv("data.csv"), 1: pd.read_csv("data2.csv")} | |
| clip_model = CLIPModel.from_pretrained(CLIP_MODEL) | |
| clip_model.to(device) | |
| clip_model.eval() | |
| clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL) | |
| owl_model = OwlViTForObjectDetection.from_pretrained(OWL_MODEL) | |
| owl_model.to(device) | |
| owl_model.eval() | |
| owl_processor = OwlViTProcessor.from_pretrained(OWL_MODEL) | |
| embeddings = { | |
| 0: np.load(f"embeddings-{MODEL}.npy"), | |
| 1: np.load(f"embeddings2-{MODEL}.npy"), | |
| } | |
| for k in [0, 1]: | |
| embeddings[k] = embeddings[k] / np.linalg.norm( | |
| embeddings[k], axis=1, keepdims=True | |
| ) | |
| return clip_model, clip_processor, owl_model, owl_processor, df, embeddings | |
| clip_model, clip_processor, owl_model, owl_processor, df, embeddings = load() | |
| mixin = ImageFeatureExtractionMixin() | |
| source = {0: "\nSource: Unsplash", 1: "\nSource: The Movie Database (TMDB)"} | |
| def compute_text_embeddings(list_of_strings): | |
| inputs = clip_processor(text=list_of_strings, return_tensors="pt", padding=True).to( | |
| device | |
| ) | |
| with torch.no_grad(): | |
| result = clip_model.get_text_features(**inputs).detach().cpu().numpy() | |
| return result / np.linalg.norm(result, axis=1, keepdims=True) | |
| def image_search(query, corpus, n_results=N_RESULTS): | |
| query_embedding = compute_text_embeddings([query]) | |
| corpus_id = 0 if corpus == "Unsplash" else 1 | |
| dot_product = (embeddings[corpus_id] @ query_embedding.T)[:, 0] | |
| results = np.argsort(dot_product)[-1 : -n_results - 1 : -1] | |
| return [ | |
| ( | |
| df[corpus_id].iloc[i].path, | |
| df[corpus_id].iloc[i].tooltip + source[corpus_id], | |
| df[corpus_id].iloc[i].link, | |
| ) | |
| for i in results | |
| ] | |
| def make_square(img, fill_color=(255, 255, 255)): | |
| x, y = img.size | |
| size = max(x, y) | |
| new_img = Image.new("RGB", (size, size), fill_color) | |
| new_img.paste(img, (int((size - x) / 2), int((size - y) / 2))) | |
| return new_img, x, y | |
| def get_images(paths): | |
| def process_image(path): | |
| return make_square(Image.open(BytesIO(requests.get(path).content))) | |
| processed = Pool(N_RESULTS).map(process_image, paths) | |
| imgs, xs, ys = [], [], [] | |
| for img, x, y in processed: | |
| imgs.append(img) | |
| xs.append(x) | |
| ys.append(y) | |
| return imgs, xs, ys | |
| def apply_owl_model(owl_queries, images): | |
| inputs = owl_processor(text=owl_queries, images=images, return_tensors="pt").to( | |
| device | |
| ) | |
| with torch.no_grad(): | |
| results = owl_model(**inputs) | |
| target_sizes = torch.Tensor([img.size[::-1] for img in images]).to(device) | |
| return owl_processor.post_process(outputs=results, target_sizes=target_sizes) | |
| def keep_best_boxes(boxes, scores, score_threshold=0.1, max_iou=0.8): | |
| candidates = [] | |
| for box, score in zip(boxes, scores): | |
| box = [round(i, 0) for i in box.tolist()] | |
| if score >= score_threshold: | |
| candidates.append((box, float(score))) | |
| to_ignore = set() | |
| for i in range(len(candidates) - 1): | |
| if i in to_ignore: | |
| continue | |
| for j in range(i + 1, len(candidates)): | |
| if j in to_ignore: | |
| continue | |
| xmin1, ymin1, xmax1, ymax1 = candidates[i][0] | |
| xmin2, ymin2, xmax2, ymax2 = candidates[j][0] | |
| if xmax1 < xmin2 or xmax2 < xmin1 or ymax1 < ymin2 or ymax2 < ymin1: | |
| continue | |
| else: | |
| xmin_inter, xmax_inter = sorted([xmin1, xmax1, xmin2, xmax2])[1:3] | |
| ymin_inter, ymax_inter = sorted([ymin1, ymax1, ymin2, ymax2])[1:3] | |
| area_inter = (xmax_inter - xmin_inter) * (ymax_inter - ymin_inter) | |
| area1 = (xmax1 - xmin1) * (ymax1 - ymin1) | |
| area2 = (xmax2 - xmin2) * (ymax2 - ymin2) | |
| iou = area_inter / (area1 + area2 - area_inter) | |
| if iou > max_iou: | |
| if candidates[i][1] > candidates[j][1]: | |
| to_ignore.add(j) | |
| else: | |
| to_ignore.add(i) | |
| break | |
| else: | |
| if area_inter / area1 > 0.9: | |
| if candidates[i][1] < 1.1 * candidates[j][1]: | |
| to_ignore.add(i) | |
| if area_inter / area2 > 0.9: | |
| if 1.1 * candidates[i][1] > candidates[j][1]: | |
| to_ignore.add(j) | |
| return [candidates[i][0] for i in range(len(candidates)) if i not in to_ignore] | |
| def convert_pil_to_base64(image): | |
| img_buffer = BytesIO() | |
| image.save(img_buffer, format="JPEG") | |
| byte_data = img_buffer.getvalue() | |
| base64_str = base64.b64encode(byte_data) | |
| return base64_str | |
| def draw_reshape_encode(img, boxes, x, y): | |
| image = img.copy() | |
| draw = ImageDraw.Draw(image) | |
| new_x, new_y = int(x * HEIGHT / y), HEIGHT | |
| for box in boxes: | |
| draw.rectangle( | |
| (tuple(box[:2]), tuple(box[2:])), outline=color, width=2 * int(y / HEIGHT) | |
| ) | |
| if x > y: | |
| image = image.crop((0, (x - y) / 2, x, x - (x - y) / 2)) | |
| else: | |
| image = image.crop(((y - x) / 2, 0, y - (y - x) / 2, y)) | |
| return convert_pil_to_base64(image.resize((new_x, new_y))) | |
| def get_html(url_list, encoded_images): | |
| html = "<div style='margin-top: 20px; max-width: 1200px; display: flex; flex-wrap: wrap; justify-content: space-evenly'>" | |
| for i in range(len(url_list)): | |
| title, encoded = url_list[i][1], encoded_images[i] | |
| html = ( | |
| html | |
| + f"<img title='{escape(title)}' style='height: {HEIGHT}px; margin: 5px' src='data:image/jpeg;base64,{encoded.decode()}'>" | |
| ) | |
| html += "</div>" | |
| return html | |
| description = """ | |
| # Search and Detect | |
| This demo illustrates how to both retrieve images containing certain objects and locate these objects with a simple text query. | |
| **Enter your query and hit enter** | |
| **Tip 1**: if your query includes "/", the left and right parts will be used to respectively retrieve images and locate objects. For example, if you want to retrieve pictures with several cats but locate individual cats, you can type "cats / cat". | |
| **Tip 2**: change the score threshold to adjust the sensitivity of the object detection step. | |
| *Built with OpenAI's [CLIP](https://openai.com/blog/clip/) model and Google's [OWL-ViT](https://arxiv.org/abs/2205.06230) model, 🤗 Hugging Face's [transformers library](https://huggingface.co/transformers/), [Streamlit](https://streamlit.io/), 25k images from [Unsplash](https://unsplash.com/) and 8k images from [The Movie Database (TMDB)](https://www.themoviedb.org/)* | |
| """ | |
| div_style = { | |
| "display": "flex", | |
| "justify-content": "center", | |
| "flex-wrap": "wrap", | |
| } | |
| def main(): | |
| st.markdown( | |
| """ | |
| <style> | |
| .block-container{ | |
| max-width: 1200px; | |
| } | |
| div.row-widget > div{ | |
| flex-direction:row; | |
| display: flex; | |
| justify-content: center; | |
| } | |
| div.row-widget.stRadio > div > label{ | |
| margin-left: 5px; | |
| margin-right: 5px; | |
| } | |
| .row-widget { | |
| margin-top: -25px; | |
| } | |
| section>div:first-child { | |
| padding-top: 30px; | |
| } | |
| div.appview-container > section:first-child{ | |
| max-width: 320px; | |
| } | |
| #MainMenu { | |
| visibility: hidden; | |
| } | |
| </style>""", | |
| unsafe_allow_html=True, | |
| ) | |
| st.sidebar.markdown(description) | |
| _, c, _ = st.columns((1, 3, 1)) | |
| query = c.text_input("", value="koala") | |
| corpus = c.radio("", ["Unsplash", "Movies"]) | |
| score_threshold = c.slider( | |
| "Score threshold", min_value=0.01, max_value=0.2, value=0.03, step=0.01 | |
| ) | |
| if len(query) > 0: | |
| if "/" in query: | |
| queries = query.split("/") | |
| clip_query, owl_query = ("/").join(queries[:-1]).strip(), queries[ | |
| -1 | |
| ].strip() | |
| else: | |
| clip_query, owl_query = query, query | |
| retrieved = image_search(clip_query, corpus) | |
| imgs, xs, ys = get_images([x[0] for x in retrieved]) | |
| results = apply_owl_model([[owl_query]] * len(imgs), imgs) | |
| encoded_images = [] | |
| for image_idx in range(len(imgs)): | |
| img0, x, y = imgs[image_idx], xs[image_idx], ys[image_idx] | |
| boxes = keep_best_boxes( | |
| results[image_idx]["boxes"], | |
| results[image_idx]["scores"], | |
| score_threshold=score_threshold, | |
| ) | |
| encoded_images.append(draw_reshape_encode(img0, boxes, x, y)) | |
| st.markdown(get_html(retrieved, encoded_images), unsafe_allow_html=True) | |
| if __name__ == "__main__": | |
| main() | |