File size: 1,725 Bytes
08fa61a bb26976 08fa61a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
import warnings
import traceback
import sys
import numpy as np
import os
from PIL import Image
from exceptions.NotFaceError import NotFaceError
from transformers import pipeline, SegformerForSemanticSegmentation, SegformerImageProcessor, SegformerFeatureExtractor
def warning_with_traceback(message, category, filename, lineno, file=None, line=None):
log = file if hasattr(file,'write') else sys.stderr
traceback.print_stack(file=log)
log.write(warnings.formatwarning(message, category, filename, lineno, line))
# warnings.showwarning = warning_with_traceback
class FaceSegmentationModel:
def __init__(self):
model_checkpoint = os.path.join("models","segformer-b0-finetuned-segments-skin-outputs", "checkpoint-1640")
self.model = SegformerForSemanticSegmentation.from_pretrained(model_checkpoint, local_files_only=True)
self.image_processor = SegformerImageProcessor.from_pretrained(model_checkpoint, local_files_only=True)
self.pipeline = pipeline("image-segmentation", model=self.model, image_processor=self.image_processor)
def infer(self, image:Image.Image):
'''
Infer the input image. it will return list of {'score', 'label', and 'mask'}
Example:
[{'score': None,
'label': 'background',
'mask': <PIL.Image.Image image mode=L size=500x500>},
{'score': None,
'label': 'acne',
'mask': <PIL.Image.Image image mode=L size=500x500>},
{'score': None,
'label': 'dry',
'mask': <PIL.Image.Image image mode=L size=500x500>}]
'''
results = self.pipeline(image)
return results
|