Vishakaraj's picture
Upload folder using huggingface_hub
55eaaa7
raw
history blame contribute delete
821 Bytes
import torch
from PIL import Image
import gradio as gr
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="large_coco", is_eval=True, device=device
)
def predict(image):
pre_processed_image = vis_processors["eval"](image).unsqueeze(0).to(device)
response = model.generate({"image": pre_processed_image}, use_nucleus_sampling=True, num_captions=3)
return image, "\n".join(response)
demo = gr.Interface(
title="Image Captioning - BLIP",
fn=predict,
inputs=gr.Image(type='pil', label="Original Image"),
outputs=[gr.Image(type="pil"), gr.Textbox()],
examples=["example_1.jpg", "example_2.jpg", "example_3.jpg"],
)
demo.launch()