import numpy as np
from PIL import Image


class NormalBaeDetector:
    def __init__(self, model):
        self.model = model

    @classmethod
    def from_pretrained(cls, pretrained_model_or_path="fal/teed", cache_dir=None, local_files_only=False): # pylint: disable=unused-argument
        from installer import install
        install('controlnet-aux', quiet=True)
        from controlnet_aux.normalbae import NormalBaeDetector as _NormalBaeDetector
        model = _NormalBaeDetector.from_pretrained(pretrained_model_or_path, filename="5_model.pth")
        return cls(model)

    def __call__(self, image, output_type="pil", **kwargs):
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        result = self.model(image, output_type=output_type)
        return result