app.py

import streamlit as st
import folium
from streamlit_folium import st_folium
import requests
from PIL import Image
from io import BytesIO
import numpy as np
import torch
from sklearn.utils.extmath import softmax
import open_clip
import os

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"
)

# Securely get the token from environment variables
MAPILLARY_ACCESS_TOKEN = os.environ.get('MAPILLARY_ACCESS_TOKEN')

# Verify token exists
if not MAPILLARY_ACCESS_TOKEN:
    st.error("Mapillary access token not found. Please configure it in the Space secrets.")
    st.stop()

def get_bounding_box(lat, lon):
    """
    Create a bounding box around a point that extends roughly 25 meters in each direction
    at Amsterdam's latitude (52.37°N):
    - 0.000224 degrees latitude = 25 meters N/S
    - 0.000368 degrees longitude = 25 meters E/W
    """
    lat_offset = 0.000224  # 25 meters in latitude
    lon_offset = 0.000368  # 25 meters in longitude
    return [
        lon - lon_offset,  # min longitude
        lat - lat_offset,  # min latitude
        lon + lon_offset,  # max longitude
        lat + lat_offset   # max latitude
    ]

def get_nearest_image(lat, lon):
    """
    Get the nearest Mapillary image to given coordinates
    """
    bbox = get_bounding_box(lat, lon)
    params = {
        'fields': 'id,thumb_1024_url',
        'limit': 1,
        'is_pano': False,
        'bbox': f'{bbox[0]},{bbox[1]},{bbox[2]},{bbox[3]}'
    }

    header = {'Authorization' : 'OAuth {}'.format(MAPILLARY_ACCESS_TOKEN)}
    try:
        response = requests.get(
            "https://graph.mapillary.com/images",
            params=params,
            headers=header
        )
        response.raise_for_status()
        data = response.json()
        
        if 'data' in data and len(data['data']) > 0:
            return data['data'][0]
        return None
        
    except requests.exceptions.RequestException as e:
        st.error(f"Error fetching Mapillary data: {str(e)}")
        return None

@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']
    if debug: st.write('allvecs.shape', allvecs.shape)
    scores = knn[f'{cat}_scores']
    if debug: st.write('scores.shape', scores.shape)
    # Compute cosine similiarity of vec against allvecs
    # (both are already normalized)
    cos_sim_table = vec @ allvecs.T
    if debug: 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, axis=1), axis=1)
    if debug: st.write('sortinds.shape', sortinds.shape)
    # Get corresponding scores for the sorted vectors
    kscores = scores[sortinds][:,:k]
    if debug: st.write('kscores.shape', kscores.shape)
    # Get actual sorted similiarity scores
    # (line copied from clip_retrieval_knn.py even though sortinds.shape[0] == 1 here)
    ksims = cos_sim_table[np.expand_dims(np.arange(sortinds.shape[0]), axis=1), sortinds]
    ksims = ksims[:,:k]
    if debug: 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: Map Explorer")
    
    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()

    # Initialize the map centered on Amsterdam
    amsterdam_coords = [52.3676, 4.9041]
    m = folium.Map(location=amsterdam_coords, zoom_start=13)
    
    # Add a marker for Amsterdam city center
    folium.Marker(
        amsterdam_coords,
        popup="Amsterdam City Center",
        icon=folium.Icon(color="red", icon="info-sign")
    ).add_to(m)
    
    # Display the map and get clicked coordinates
    map_data = st_folium(m, height=400, width=700)
    
    # Check if a location was clicked
    if map_data['last_clicked']:
        lat = map_data['last_clicked']['lat']
        lng = map_data['last_clicked']['lng']
        
        st.write(f"Selected coordinates: {lat:.4f}, {lng:.4f}")
        
        # Get nearest Mapillary image
        with st.spinner('Fetching street view image...'):
            image_data = get_nearest_image(lat, lng)
            
            if image_data:
                # Display the image
                try:
                    response = requests.get(image_data['thumb_1024_url'])
                    image = Image.open(BytesIO(response.content))
                    st.image(image, caption="Street View", width=400)
                    
                    # Add download button
                    st.download_button(
                        label="Download Image",
                        data=response.content,
                        file_name=f"streetview_{lat}_{lng}.jpg",
                        mime="image/jpeg"
                    )
                    
                except Exception as e:
                    st.error(f"Error displaying image: {str(e)}")
            else:
                st.warning("No street view images found at this location. Try a different spot.")

if __name__ == "__main__":
    main()