27 lines
821 B
Python
27 lines
821 B
Python
import os
|
|
|
|
from funcy import rcompose
|
|
|
|
from image_prediction.config import CONFIG
|
|
from image_prediction.default_objects import get_extractor_classifier, get_formatter, get_mlflow_model_loader
|
|
from image_prediction.locations import MLRUNS_DIR
|
|
|
|
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
|
|
|
|
|
|
def load_pipeline(**kwargs):
|
|
model_loader = get_mlflow_model_loader(MLRUNS_DIR)
|
|
model_identifier = CONFIG.service.run_id
|
|
|
|
pipeline = Pipeline(model_loader, model_identifier, **kwargs)
|
|
|
|
return pipeline
|
|
|
|
|
|
class Pipeline:
|
|
def __init__(self, model_loader, model_identifier, **kwargs):
|
|
self.pipe = rcompose(get_extractor_classifier(model_loader, model_identifier, **kwargs), get_formatter())
|
|
|
|
def __call__(self, pdf: bytes, page_range: range = None):
|
|
yield from self.pipe(pdf, page_range=page_range)
|