45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
import os
|
|
|
|
from funcy import rcompose, juxt
|
|
|
|
from image_prediction.classifier.classifier import Classifier
|
|
from image_prediction.classifier.image_classifier import ImageClassifier
|
|
from image_prediction.estimator.adapter.adapter import EstimatorAdapter
|
|
from image_prediction.extractor_classifier.extractor_classifier import ExtractorClassifier
|
|
from image_prediction.formatter.formatters.enum import EnumFormatter
|
|
from image_prediction.image_extractor.extractors.parsable import ParsablePDFImageExtractor
|
|
from image_prediction.label_mapper.mappers.probability import ProbabilityMapper
|
|
|
|
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
|
|
|
|
|
|
def get_image_classifier(model_loader, model_identifier):
|
|
model, classes = juxt(model_loader.load_model, model_loader.load_classes)(model_identifier)
|
|
return ImageClassifier(Classifier(EstimatorAdapter(model), ProbabilityMapper(classes)))
|
|
|
|
|
|
def get_extractor():
|
|
image_extractor = ParsablePDFImageExtractor(verbose=True)
|
|
|
|
return image_extractor
|
|
|
|
|
|
def get_extractor_classifier(model_loader, model_identifier):
|
|
extractor_classifier = ExtractorClassifier(get_extractor(), get_image_classifier(model_loader, model_identifier))
|
|
|
|
return extractor_classifier
|
|
|
|
|
|
def get_formatter():
|
|
formatter = EnumFormatter()
|
|
|
|
return formatter
|
|
|
|
|
|
class Pipeline:
|
|
def __init__(self, model_loader, model_identifier):
|
|
self.pipe = rcompose(get_extractor_classifier(model_loader, model_identifier), get_formatter())
|
|
|
|
def __call__(self, pdf: bytes):
|
|
yield from self.pipe(pdf)
|