Image-Colour-Extract / predict.py
ADITYA KUMAR
Update predict.py
fb1e70e
raw
history blame contribute delete
654 Bytes
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
from PIL import Image
from io import BytesIO
import requests
import numpy as np
model = AutoModelForSequenceClassification.from_pretrained("./extract-colors.py")
tokenizer = AutoTokenizer.from_pretrained("./tokenizer", tokenizer_func=tokenizer_function)
model_class = pipeline('image-classification', model=model, tokenizer=tokenizer)
def get_colors_and_closest_to_white(image_url):
response = requests.get(image_url)
img = Image.open(BytesIO(response.content))
img_array = np.array(img)
result = model_class(image_array=img_array)
return result