koclip / text2patch.py
jaketae's picture
feature: add intro page, cleanup descriptions
import os
import jax
import jax.numpy as jnp
import numpy as np
import requests
import streamlit as st
from PIL import Image
from utils import load_model
def split_image(im, num_rows=3, num_cols=3):
im = np.array(im)
row_size = im.shape[0] // num_rows
col_size = im.shape[1] // num_cols
tiles = [
im[row : row + row_size, col : col + col_size]
for row in range(0, num_rows * row_size, row_size)
for col in range(0, num_cols * col_size, col_size)
return tiles
def app(model_name):
model, processor = load_model(f"koclip/{model_name}")
st.title("Patch-based Relevance Ranking")
Given a piece of text, the CLIP model finds the part of an image that best explains the text.
To try it out, you can
1. Upload an image
2. Explain a part of the image in text
which will yield the most relevant image tile from a grid of the image. You can specify how
granular you want to be with your search by specifying the number of rows and columns that
make up the image grid.
query1 = st.text_input(
"Enter a URL to an image...",
query2 = st.file_uploader("or upload an image...", type=["jpg", "jpeg", "png"])
captions = st.text_input(
"Enter a prompt to query the image.",
value="이건 서울의 경복궁 사진이다.",
col1, col2 = st.beta_columns(2)
with col1:
num_rows = st.slider(
"Number of rows", min_value=1, max_value=5, value=3, step=1
with col2:
num_cols = st.slider(
"Number of columns", min_value=1, max_value=5, value=3, step=1
if st.button("질문 (Query)"):
if not any([query1, query2]):
st.error("Please upload an image or paste an image URL.")
with st.spinner("Computing..."):
image_data = (
if query2 is not None
else requests.get(query1, stream=True).raw
image = Image.open(image_data)
images = split_image(image, num_rows, num_cols)
inputs = processor(
text=captions, images=images, return_tensors="jax", padding=True
inputs["pixel_values"] = jnp.transpose(
inputs["pixel_values"], axes=[0, 2, 3, 1]
outputs = model(**inputs)
probs = jax.nn.softmax(outputs.logits_per_image, axis=0)
for idx, prob in sorted(
enumerate(probs), key=lambda x: x[1], reverse=True
st.text(f"Score: {prob[0]:.3f}")