import streamlit as st
import torch
from PIL import Image
from lavis.models import load_model_and_preprocess

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model, vis_processors, _ = load_model_and_preprocess(name="blip_caption", model_type="base_coco", is_eval=True, device=device)


def generate_caption(raw_image):
    raw_image = raw_image.convert('RGB')
    image = vis_processors["eval"](raw_image).unsqueeze(0).to(device)
    return model.generate({"image": image})[0]


st.title("EcomGenius")

file_name = st.file_uploader("Upload a Image")

if file_name is not None:
    col1, col2 = st.columns(2)

    image = Image.open(file_name)
    col1.image(image, use_column_width=True)
    caption = generate_caption(image)

    col2.header("Probabilities")
    col2.subheader(caption)