Akbartus's picture
Update app.py
ed80744
raw
history blame
1.46 kB
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
from turtle import title
from transformers import pipeline
import numpy as np
from PIL import Image
import sys, os
import gradio as gr
model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
pipe = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32")
images="dog.jpg"
def shot(image, labels_text):
PIL_image = Image.fromarray(np.uint8(image)).convert('RGB')
labels = labels_text.split(",")
# Translate
tokenizer.src_lang = "en"
encodedText = tokenizer(labels_text, return_tensors="pt")
generatedTokens = model.generate(**encodedText, forced_bos_token_id=tokenizer.get_lang_id("ru"))
return tokenizer.batch_decode(generatedTokens, skip_special_tokens=True)[0]
res = pipe(images=PIL_image,
candidate_labels=labels,
hypothesis_template= "This is a photo of a {}")
return {dic["label"]: dic["score"] for dic in res}
iface = gr.Interface(shot,
["image", "text"],
"label",
examples=[["dog.jpg", "dog,cat,bird"]],
description="Add a picture and a list of labels separated by commas",
title="Zero-shot Image Classification")
iface.launch()