File size: 962 Bytes
8b2e7f8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
import os
from flash.image import ImageClassifier
# ⚠️ You need this to access the state key
from flash.core.data.data_source import LabelsState
class PreTrainedPipeline():
def __init__(self, path=""):
self.device = 'cpu'
self.model = ImageClassifier.load_from_checkpoint(os.path.join(path, "pytorch_model.bin"))
self.data_pipeline = self.model.build_data_pipeline()
self.labels = self.model._data_pipeline_state._state[LabelsState].labels
self.top_k = 5
def __call__(self, inputs):
x = self.data_pipeline.worker_preprocessor('predict')({'input': inputs})
x = self.model.transfer_batch_to_device(x, self.device, 0)
x = self.data_pipeline.device_preprocessor('predict')(x)
out = self.model.predict_step(x, 0)
proba = out['preds'].softmax(1)[0].tolist()
return [{'score': s, 'label': l} for s, l in sorted(zip(proba, self.labels), reverse=True)[:self.top_k]]
|