|
from typing import Any, Dict, List |
|
|
|
import numpy as np |
|
from huggingface_hub import from_pretrained_fastai |
|
from PIL import Image |
|
|
|
|
|
class ImageClassificationPipeline(): |
|
def __init__(self, model_id: str): |
|
self.model = from_pretrained_fastai(model_id) |
|
|
|
|
|
self.id2label = self.model.dls.vocab |
|
|
|
|
|
self.top_k = 5 |
|
|
|
def __call__(self, inputs: "Image.Image") -> List[Dict[str, Any]]: |
|
""" |
|
Args: |
|
inputs (:obj:`PIL.Image`): |
|
The raw image representation as PIL. |
|
No transformation made whatsoever from the input. Make all necessary transformations here. |
|
Return: |
|
A :obj:`list`:. The list contains items that are dicts should be liked {"label": "XXX", "score": 0.82} |
|
It is preferred if the returned list is in decreasing `score` order |
|
""" |
|
|
|
_, _, preds = self.model.predict(np.array(inputs)) |
|
preds = preds.tolist() |
|
|
|
labels = [ |
|
{"label": str(self.id2label[i]), "score": float(preds[i])} |
|
for i in range(len(preds)) |
|
] |
|
return labels |
|
|
