import streamlit as st
import torch
import requests
from PIL import Image
from torchvision import transforms
from torchvision.transforms.functional import InterpolationMode

from models.blip import blip_decoder

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

@st.cache(show_spinner=False)
def load_demo_image(image_size, device):
    img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg'
    raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
    w,h = raw_image.size
    transform = transforms.Compose([
        transforms.Resize((image_size, image_size), interpolation=InterpolationMode.BICUBIC),
        transforms.ToTensor(),
        transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
    ])
    image = transform(raw_image).unsqueeze(0).to(device)
    return image, raw_image.resize((w//5,h//5))

def main():
    st.set_page_config(page_title="Image Captioning App")
    st.title("Image Captioning App")
    st.write("This app generates captions for images using a pre-trained model.")
    
    # Load image
    image_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
    if image_file is not None:
        image = Image.open(image_file)
        image_size = 384
        transform = transforms.Compose([
            transforms.Resize((image_size, image_size), interpolation=InterpolationMode.BICUBIC),
            transforms.ToTensor(),
            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
        ])
        image = transform(image).unsqueeze(0).to(device)
        
        # Generate captions
        with torch.no_grad():
            model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_capfilt_large.pth'
            model = blip_decoder(pretrained=model_url, image_size=image_size, vit='base')
            model.eval()
            model = model.to(device)
            num_captions = 3
            captions = []
            for i in range(num_captions):
                caption = model.generate(image, sample=True, top_p=0.9, max_length=20, min_length=5)
                captions.append(caption[0])
            for i, caption in enumerate(captions):
                st.write(f'Caption {i+1}: {caption}')
            
        # Display uploaded image
        st.image(image_file, caption='Uploaded image', use_column_width=True)
        
if __name__ == "__main__":
    main()