Akbartus's picture
Update app.py
64115b7
raw
history blame
1.73 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(label, 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
tokenizer.src_lang = "en"
encodedText = tokenizer(label, 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]
}
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()