commit 2a8aaff874a7222feb766288ec319f36b356b3d0 Author: jianjiang Date: Fri Apr 21 17:59:47 2023 +0800 ok diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..5b47bb2 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,11 @@ +#FROM python:3.8.13 +FROM artifacts.iflytek.com/docker-private/atp/base_image_for_ailab:0.0.1 + +WORKDIR /app + +COPY . /app + +RUN pip config set global.index-url https://pypi.mirrors.ustc.edu.cn/simple +RUN pip install -r requirements.txt + +CMD ["python", "app.py"] diff --git a/app.py b/app.py new file mode 100644 index 0000000..e3d6822 --- /dev/null +++ b/app.py @@ -0,0 +1,51 @@ +import torch +import requests +from PIL import Image +from transformers import ViTFeatureExtractor, AutoTokenizer, VisionEncoderDecoderModel +import gradio as gr +from gradio.themes.utils import sizes + + +theme = gr.themes.Default(radius_size=sizes.radius_none).set( + block_label_text_color = '#4D63FF', + block_title_text_color = '#4D63FF', + button_primary_text_color = '#4D63FF', + button_primary_background_fill='#FFFFFF', + button_primary_border_color='#4D63FF', + button_primary_background_fill_hover='#EDEFFF', +) + + +loc = "ydshieh/vit-gpt2-coco-en" + +feature_extractor = ViTFeatureExtractor.from_pretrained(loc) +tokenizer = AutoTokenizer.from_pretrained(loc) +model = VisionEncoderDecoderModel.from_pretrained(loc) +model.eval() + + +def predict(image): + pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values + + with torch.no_grad(): + output_ids = model.generate(pixel_values, max_length=16, num_beams=4, return_dict_in_generate=True).sequences + + preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True) + + total_caption = "" + for pred in preds: + total_caption = total_caption + pred.strip() + total_caption = total_caption + "\r\n" + + return total_caption + +demo = gr.Interface(fn=predict, + inputs='image', + outputs='text', + title = "image2text", + theme = theme, + examples = ['soccer.jpg']) + + +if __name__ == "__main__": + demo.queue(concurrency_count=1).launch(server_name = "0.0.0.0") diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..7e2fba5 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +Pillow diff --git a/soccer.jpg b/soccer.jpg new file mode 100644 index 0000000..3f2ec82 Binary files /dev/null and b/soccer.jpg differ