image-classification-service/fb_detr/utils/estimator.py

import os

from fb_detr.config import CONFIG
from fb_detr.locations import DATA_DIR, TORCH_HOME
from fb_detr.predictor import Predictor


def suppress_userwarnings():
    import warnings

    warnings.filterwarnings("ignore")


def load_classes():
    classes = CONFIG.estimator.classes
    id2class = dict(zip(range(1, len(classes) + 1), classes))
    return id2class


def get_checkpoint():
    return DATA_DIR / CONFIG.estimator.checkpoint


def set_torch_env():
    os.environ["TORCH_HOME"] = str(TORCH_HOME)


def initialize_predictor(resume):
    set_torch_env()
    checkpoint = get_checkpoint() if not resume else resume
    predictor = Predictor(checkpoint, classes=load_classes(), rejection_class=CONFIG.estimator.rejection_class)
    return predictor