models.py

from transformers import AutoImageProcessor, DetrForSegmentation
from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation


def get_detr():
    image_processor = AutoImageProcessor.from_pretrained("facebook/detr-resnet-50-panoptic")
    model = DetrForSegmentation.from_pretrained("facebook/detr-resnet-50-panoptic")
    return model, image_processor


def get_maskformer():
    feature_extractor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-small-coco")
    model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-small-coco")
    return model, feature_extractor