File size: 654 Bytes
aa1f5e1
 
 
 
 
4ea7291
 
aa1f5e1
 
 
 
4ea7291
 
aa1f5e1
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from transformers import AutoImageProcessor, DetrForSegmentation
from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation


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


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