File size: 824 Bytes
6b8d35c
 
bca9f16
6b8d35c
 
bca9f16
6b8d35c
bca9f16
 
 
 
 
6b8d35c
bca9f16
 
 
 
 
 
 
 
 
 
6b8d35c
 
bca9f16
6b8d35c
 
bca9f16
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
from PIL import Image
import requests
import torch
import gradio as gr

from transformers import pipeline

CAPTION_MODELS = {
    'blip-base': 'Salesforce/blip-image-captioning-base',
    'blip-large': 'Salesforce/blip-image-captioning-large',
    'vit-gpt2-coco-en': 'ydshieh/vit-gpt2-coco-en',
}

captioner = pipeline(task="image-to-text",
                        model=CAPTION_MODELS['blip-base'],
                        max_new_tokens=30,
                        device_map="auto", use_fast=True
                        )

# Simple caption creation
def caption_image(captioner, image_path):
    caption = captioner(image_path)[0]['generated_text']
    return str(caption).strip()

def launch(input):
    return caption_image(captioner, input)

iface = gr.Interface(launch, inputs="text", outputs="text")
iface.launch()