46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
import os
|
|
|
|
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
|
|
|
|
from image_prediction.classifier.classifier import Classifier
|
|
from image_prediction.classifier.image_classifier import ImageClassifier
|
|
from image_prediction.config import CONFIG
|
|
from image_prediction.estimator.adapter.adapter import EstimatorAdapter
|
|
from image_prediction.extractor_classifier.extractor_classifier import ExtractorClassifier
|
|
from image_prediction.image_extractor.extractors.parsable import ParsablePDFImageExtractor
|
|
from image_prediction.locations import MLRUNS_DIR
|
|
from image_prediction.model_loader.loader import ModelLoader
|
|
from image_prediction.model_loader.loaders.mlflow import MlflowConnector
|
|
from image_prediction.redai_adapter.mlflow import MlflowModelReader
|
|
|
|
|
|
def get_image_classifier():
|
|
model_loader = ModelLoader(MlflowConnector(MlflowModelReader(MLRUNS_DIR)))
|
|
model = model_loader.load_model(CONFIG.service.run_id)
|
|
classes = model_loader.load_classes(CONFIG.service.run_id)
|
|
classifier = Classifier(EstimatorAdapter(model), classes)
|
|
image_classifier = ImageClassifier(classifier)
|
|
|
|
return image_classifier
|
|
|
|
|
|
def get_extractor():
|
|
image_extractor = ParsablePDFImageExtractor(verbose=True)
|
|
|
|
return image_extractor
|
|
|
|
|
|
def get_extractor_classifier():
|
|
extractor_classifier = ExtractorClassifier(get_extractor(), get_image_classifier())
|
|
|
|
return extractor_classifier
|
|
|
|
|
|
class Pipeline:
|
|
|
|
def __init__(self):
|
|
self.pipeline = get_extractor_classifier()
|
|
|
|
def __call__(self, pdf: bytes):
|
|
return self.pipeline(pdf)
|