app.py

import streamlit as st
from PIL import Image
import numpy as np
import torch

import open_clip

#from transformers import CLIPProcessor, CLIPModel

knnpath = '20241204-ams-no-env-open_clip_ViT-H-14-378-quickgelu.npz'
clip_model_name = 'ViT-H-14-378-quickgelu'
pretrained_name = 'dfn5b'

categories = ['walkability', 'bikeability', 'pleasantness', 'greenness', 'safety']

# Set page config
st.set_page_config(
    page_title="Percept",
    layout="wide"
)

#model, preprocess = open_clip.create_model_from_pretrained('hf-hub:laion/CLIP-ViT-g-14-laion2B-s12B-b42K')
#tokenizer = open_clip.get_tokenizer('hf-hub:laion/CLIP-ViT-g-14-laion2B-s12B-b42K')

#model, preprocess = open_clip.create_model_from_pretrained(clip_model_name)
#tokenizer = open_clip.get_tokenizer(clip_model_name)

#st.write("Available models:", open_clip.list_models())

@st.cache_resource
def load_model():
    """Load the OpenCLIP model and return model and processor"""
    model, _, preprocess = open_clip.create_model_and_transforms(
        clip_model_name, pretrained=pretrained_name
    )
    tokenizer = open_clip.get_tokenizer(clip_model_name)
    return model, preprocess, tokenizer

def process_image(image, preprocess):
    """Process image and return tensor"""
    if isinstance(image, str):
        # If image is a URL
        response = requests.get(image)
        image = Image.open(BytesIO(response.content))
    # Ensure image is in RGB mode
    if image.mode != 'RGB':
        image = image.convert('RGB')
    processed_image = preprocess(image).unsqueeze(0)
    return processed_image

def knn_get_score(knn, k, cat, vec):
    allvecs = knn[f'{cat}_vecs']
    st.write('allvecs.shape', allvecs.shape)
    scores = knn[f'{cat}_scores']
    st.write('scores.shape', scores.shape)
    # Compute cosine similiarity of vec against allvecs
    # (both are already normalized)
    cos_sim_table = vec @ allvecs.T
    st.write('cos_sim_table.shape', cos_sim_table.shape)
    # Get sorted array indices by similiarity in descending order
    sortinds = np.flip(np.argsort(cos_sim_table))
    st.write('sortinds.shape', sortinds.shape)
    # Get corresponding scores for the sorted vectors
    kscores = scores[sortinds][:k]
    st.write('kscores.shape', kscores.shape)
    # Get actual sorted similiarity scores
    ksims = np.expand_dims(cos_sim_table[sortinds][:k], axis=0)
    st.write('ksims.shape', ksims.shape)
    # Apply normalization after exponential formula
    ksims = softmax(10**ksims)
    # Weighted sum
    kweightedscore = np.sum(kscores * ksims)
    return kweightedscore


@st.cache_resource
def load_knn():
    return np.load(knnpath)

def main():
    st.title("Percept: Human Perception of Street View Image Analyzer")
 
    try:
        with st.spinner('Loading CLIP model... This may take a moment.'):
            model, preprocess, tokenizer = load_model()
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model = model.to(device)
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        st.info("Please make sure you have enough memory and the correct dependencies installed.")

    with st.spinner('Loading KNN model... This may take a moment.'):
        knn = load_knn()
    st.write(knn['walkability_vecs'].shape)

    file = st.file_uploader('Upload An Image')

    if file:
        try:
            image = Image.open(file)

            st.image(image, caption="Uploaded Image", width=640)
                
            # Process image
            with st.spinner('Processing image...'):
                processed_image = process_image(image, preprocess)
                processed_image = processed_image.to(device)

                # Encode into CLIP vector
                with torch.no_grad():
                    vec = model.encode_image(processed_image)
                    
                    # Normalize vector
                    vec /= vec.norm(dim=-1, keepdim=True)
                st.write(vec.shape)
                k = 40
                for cat in ['walkability']:
                    st.write(cat, 'rating =', knn_get_score(knn, k, cat, vec))

        except Exception as e:
            st.error(f"Error processing image: {str(e)}")

if __name__ == "__main__":
    main()