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